Skip to content

[lora] refactor - native LoRA - #2017

Merged
yushengsu-thu merged 19 commits into
zhichen/lora-nativefrom
yusheng/lora-plugin-refactor
Aug 3, 2026
Merged

[lora] refactor - native LoRA#2017
yushengsu-thu merged 19 commits into
zhichen/lora-nativefrom
yusheng/lora-plugin-refactor

Conversation

@yushengsu-thu

@yushengsu-thu yushengsu-thu commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Refactor PR into #1792's branch (zhichen/lora-native), so the native-LoRA extraction ships as a plugin from the start. Adapter math is unchanged — a side-by-side harness against the base monolith shows bit-identical adapter parameters, forward outputs, backward gradients, and HF exports for full-target GQA (TP=1, TP=2 ± SP) and MLA configs. Bridge-mode LoRA is untouched.

Includes #2072 (function-level redesign + fixes from validating four models end-to-end), the lora-native CI suite (#2073), and — after merging main's Inkling LoRA (#2122) — a rewrite of Inkling onto the plugin architecture, replacing the model-specific monolith.

Layout

miles_plugins/lora/
    __init__.py
    lora.py              # Single public entry point and lifecycle orchestration
    config.py            # LoRAConfig: rank, alpha, dropout, and target modules
    registry.py          # model_type -> ModelEntry(spec, status); launcher helpers
    hf_adapter.py        # HF/PEFT adapter naming, import, and export
    serving.py           # SGLang serving export (fused-buffer zero-fill, target expansion)
    checkpointing.py     # Adapter checkpoint state dicts and per-rank shard naming
    distributed.py       # TP/EP/SP communication and gradient handling (ParallelGather)

    spec/
      base.py            # Contracts: ProjectionSpec, AttachContext, protocols, enums
      layout.py          # Declarative layouts: dim resolvers, bindings, attach walk, spec bases
      attention.py       # GQA / MLA / GDN / hybrid attention spec classes
      mlp.py             # Fused gate/up + down spec class
      moe.py             # Routed-expert per-layer policy; Inkling grouped experts

    modules/
      linear.py          # LoRALinear + LoRASplitAdapter (SplitQKV / SplitFC1)
      moe.py             # Grouped-GEMM expert, shared-experts, and lm_head adapters

    kernel/              # Future fused and grouped-expert kernels

Design

Classes carry the taxonomy; every fact is declared once and derived everywhere else.

  • Declarative layouts (spec/layout.py): an architecture spec is a class holding a ModuleLayout table — which projections exist, on which physical linear, with which shard geometry. One shared attach_layout() implements the existence-check → target-filter → guard → build → hook walk; the dim-resolver section is the only spec-side reader of MCore attribute names. A new layout is ~15 lines of facts.
  • Spec hierarchy (LayoutSpecAttentionSpecBaseGQA/MLA/GDN/Inkling, HybridGQAGDN(GQAAttentionSpec), SharedOuterExpertMoESpec(GeneralExpertMoESpec)): supported_targets, the canonical --target-modules order, and SGLang's fused-family expansion all derive from the layout declaration; a new architecture's fused groups reach serving-side target expansion automatically. ShardLayout/AttentionFamily str-enums replace bare strings; no parallel name tables in the exporters.
  • ModelEntry(spec, status, reason) with VALIDATED / STRUCTURAL / UNSTABLE: support status is one first-class field; UNSTABLE entries carry their reason and warn at attach.
  • Public export descriptors: adapters expose exports() -> ProjectionExport (with SGLangFusedGroup metadata); serving.py consumes only that API.
  • Adapter IO hooks: modules whose export shape isn't a per-projection column/row gather override export_plan(gather) / load_plan_custom(take); hf_adapter.py stays generic. ParallelGather batches async all-gathers over the TP or EP group behind one request(tensor, dim, group=...) call.
  • Escape hatches for non-table structure: ModuleLayout.hf_block_prefix and ProjectionBinding.adapter_class override HF naming and the module class per binding; LoRAArchSpec gained moe.attach(...) and an optional lm_head spec so a fully custom family still routes through the one orchestrator in lora.py.
  • Launcher helpers: default_target_modules(hf_checkpoint) and preflight_native_lora(...) (registry ∧ mbridge bridge ∧ model-args audit, no GPU) — launch scripts stop re-declaring per-family target strings and fail in seconds when a conversion bridge is missing.

Inkling on the plugin (replaces the #2122 monolith)

miles_plugins/models/inkling/lora.py (759 lines, its own attach/export/IO stack) is deleted; Inkling's specs live in the role files next to the base classes they inherit from (InklingAttentionSpec in spec/attention.py, InklingDenseMLPSpec in spec/mlp.py, InklingMoESpec in spec/moe.py, InklingLMHeadSpec beside its registry entry), registered as inkling_mm_model:

  • Attention: 4-way fused [q;k;v;r] declared as one FusedAttach — Inkling's fusion is a plain concat (no GQA group-permute), so the split-adapter row table {q,k,v,r} is the entire model-specific code. TML export names (attn.wq_du/wk_dv/wv_dv/wr_du/wo_ud) come straight from the ProjectionSpecs.
  • Dense MLP: fused gate/up exports as the single TML gate_up_proj projection via an adapter_class override whose export_plan gathers the gate/up halves separately (a plain dim-0 TP gather would interleave ranks).
  • Routed experts: LoRAGroupedFC1/FC2 — shared-A, per-expert-B grouped-GEMM deltas (F.grouped_mm over the router's token offsets), EP-aware export/load. Shared experts: one adapter with per-sub-expert B slices, hooked into each sub-expert's forward.
  • lm_head: LoRAOutputHead with muP logit scaling and unpadded-vocab trim.
  • Serving: TML names are auto-detected by SGLang, so serving passes lora_target_modules=["all"] (sentinel keyed off the checkpoint, not the model list).
  • normalize_targets maps any request (all-linear or HF-style names) to the full native TML set, matching the provider it replaces.

Fixes found by running the models natively

  • miles_plugins/mbridge/kimi_k25.py (new): raw-mode K2.5 conversion needs an mbridge bridge (bridge mode uses megatron.bridge's own KimiK25VL bridge and never converts). Thin DeepseekV3Bridge subclass: text_config descent + language_model. prefix remap.
  • megatron_to_hf kimi dispatch: _convert_to_hf_core matched only kimi_k25, but the direct weight iterator passes kimik25config (dead branch); the converter also lacked the raw text-only provider's unprefixed top-level weight names.
  • convert_kimi_int4_to_bf16.py strips quantization_config: a surviving one sends SGLang through its CompressedTensors path, which serves the BF16 checkpoint with a degenerate context-free forward → garbage rollouts, train/rollout logprob_abs_diff ~2.2. Verified by a megatron-vs-HF-vs-SGLang three-way logprob comparison.
  • torch_dsa_topk clamps k to the key length: short sequences (< index_topk keys) crashed torch.topk in raw-mode log-prob compute.
  • run_glm5_2_744b_a40b_lora_native.py (new): raw-mode GLM-5.2 launcher; its toy conversion must be single-rank (any PP split of the 5-layer toy starts a stage on a DSA skip layer).

CI

tests/e2e/lora_native/ (label lora-native, trigger run-ci-lora-native): qwen3, qwen3.5, glm4.7, glm5.2-5layer, kimi-2.5-2layer, each a 20-rollout native run checking logprob parity and weight-update equality (vision towers and engine-transformed params skip-listed). The pre-existing inkling 4-layer CI runs against the plugin path.

Validation

20-rollout native GRPO runs (wandb ch271828n-team/miles_lora_native), acceptance ≈ 0.01 train/rollout logprob_abs_diff:

model logprob_abs_diff
Qwen3-8B 0.0099
Qwen3.5-35B-A3B 0.0105 (grad_norm ~1e-2 throughout; the recorded GDN raw-backward divergence did not reproduce → qwen3_5_moe VALIDATED)
GLM-5.2 5-layer 0.0096
Kimi-K2.5 2-layer 0.0124 (after the quantization_config fix)
Inkling-Small (4-node LoRA) running

tests/fast/miles_plugins/lora: 64 passed (includes Inkling spec locks). Behavior-preservation smoke: a 2-rollout Qwen3-8B run reproduces the baseline's first-rollout logprob_abs_diff (0.00967 vs 0.00978). Adapter attribute names are pinned, so checkpoint keys are unchanged.

Moves miles/backends/megatron_utils/lora_native.py (906 lines, one file)
into a miles_plugins/lora package, so model-architecture support lives in
the plugin layer like mbridge and megatron_bridge already do. Bridge-mode
LoRA is untouched: the shared helpers both paths use stay in
lora_utils.py, and nothing under miles/ imports the plugin at module
level -- core reaches it only through the resolve_lora_provider dotted
path, whose default is now miles_plugins.lora.lora.

Layout, split along the seams the code already had:

* adapter.py   -- NativeLoRAAdapter and the _ADAPTER_LAYOUT projection
                  table. Adapters stay self-describing, so export/load
                  never know which spec attached them.
* io.py        -- HF/PEFT export (TP-gathered) and per-rank load.
* naming.py    -- the checkpoint-index naming heuristic, plus an optional
                  mbridge cross-check (lazy import, warns on disagreement,
                  no-op when mbridge is absent).
* spec/        -- attach paths: attention.py (GQA fused-qkv + MLA, with
                  their guards), mlp.py (gated MLP / shared experts),
                  base.py (per-run _Spec and shared primitives).
* registry.py  -- new: HF model_type -> spec family, keyed the same way
                  mbridge's register_model is. Unregistered architectures
                  fail at startup listing the supported set and the
                  escape hatches; checkpoints with no config.json (the
                  numerical harnesses) fall back to structural dispatch
                  with a warning. The config x parallelism guards
                  (num_query_groups >= TP, MLA q_lora_rank, replicated
                  down-projections) are unchanged and still apply on top.
* lora.py      -- the provider entrypoint: apply_native_lora, the run
                  guards, and the three-function provider protocol.

Function bodies are unchanged from lora_native.py except apply_native_lora
gaining the registry lookup and arch=... in its startup log line.
resolve_lora_provider moves to lora_utils.py, next to the other helpers
its call sites (model.py, hf_weight_iterator_direct.py,
save_lora_checkpoint) already import.

Tests move to tests/fast/miles_plugins/lora/ -- the 47 existing tests
unchanged apart from imports, plus 17 new registry/naming tests. Verified
against the parent branch in the same environment: identical pass/fail
sets on the moved suite and on the four lora_utils-dependent suites
(test_lora_utils, test_lora_model_branches, test_lora_update_weight,
test_lora_hf_weight_iterator); the only local failure either way is
megatron-core's triton import on macOS. black / isort / ruff clean.

Co-authored-by: Cursor <cursoragent@cursor.com>
Copilot AI review requested due to automatic review settings July 31, 2026 02:57
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Comment thread miles/backends/megatron_utils/lora_utils.py Outdated
…fety, API cleanup

- MoE without shared experts (qwen3_moe): parser-expanded all-linear MLP
  targets now skip with a log line instead of asserting at startup,
  mirroring the MLA generic-qkv normalization; explicit gate/up/down
  requests still fail closed.
- Registry: drop deepseek_v4 (DeepSeek-V4-Flash attention is wq_a/wq_b/wkv,
  not mcore MLA — docs already declare it out of scope), and report
  "config.json has no model_type" distinctly from "config.json missing".
- Resume safety: load_native_adapter_state_dict now reports missing tensors
  (shard saved with narrower targets) in addition to unexpected ones, and
  the optimizer/scheduler-restore decision is agreed across ranks via
  all_gather_object so per-rank shard differences cannot desync the resume
  iteration or LR schedule.
- Provider path: the package itself now exports the provider protocol, the
  default becomes miles_plugins.lora (old miles_plugins.lora.lora pins keep
  working, including for the native save fast path).
- Delete all 17 transitional underscore aliases; tests and the numerical
  harness import the public names.
- Extract parse_lora_target_modules out of miles_validate_args so
  tests/fast/utils/test_lora_arguments.py exercises the production function
  instead of a hand-copied replica (now also covers MLA all-linear
  expansion).
- New tests/fast/miles_plugins/lora/test_orchestration.py: apply_native_lora
  attach counts / freeze / grad flow on fake decoders, registry-gated
  startup failures, _assert_supported_run rejections, and MLA attach for
  both q_lora_rank branches (previously 0% fast-test coverage on lora.py
  and the MLA attach path).
- docs: VL wrappers resolve via text_config.model_type (vision never
  adapted; unregistered VL types fail closed), shared-expert-less MoE
  all-linear behavior, provider default.
@yushengsu-thu
yushengsu-thu force-pushed the yusheng/lora-plugin-refactor branch from f3c0c73 to dd9d77c Compare July 31, 2026 20:11
- SharedExpertOnlyMoESpec splits into SharedOuterExpertMoESpec (layers with a
  shared/outer expert: LoRA adapts the shared expert's MLP) and
  GeneralExpertMoESpec (routed-only layers: all-linear skips, explicit MLP
  targets fail closed). The shared spec delegates routed-only layers to the
  general spec, so one registry entry still covers mixed-layer models.
  Module docstring spells out the distinction from the bridge-only
  --experts-shared-outer-loras adapter layout (SGLang PR #21466).
- SplitQKV -> LoRASplitQKV, SplitFC1 -> LoRASplitFC1, aligning with
  Megatron-Bridge's LoRALinearSplitQKV / LoRALinearSplitFC1UpGate naming
  (LoRALinear already matches).
@yushengsu-thu yushengsu-thu changed the title [lora] extract native LoRA into the miles_plugins/lora plugin [lora] refactor - native LoRA Jul 31, 2026
…pported combos

Findings from a parallelism audit of the native path. TP-only handling is
correct for what the plugin attaches (ETP is moot because MCore builds the MoE
shared expert on the attention TP group; CP needs nothing because adapters are
ordinary replicated DDP-bucket params reduced over dp_cp), but the surrounding
plumbing had two real bugs:

- EP>1 resume found no shard. The adapter shard name carried ep_rank while only
  dense-DP-rank-0 ranks wrote, so under scripts/run_lora_native.py's shipped
  TP2/EP4 defaults exactly {tp0_pp0.pt, tp1_pp0_ep1.pt} were written and the
  EP 2/3 ranks found nothing (verified against the pinned fork's RankGenerator).
  Native adapter state is EP-invariant, so its name is now keyed on (tp, pp)
  only; bridge shards keep the EP component because bridge PEFT can attach
  genuinely expert-parallel adapters. Writers are now elected per distinct
  shard filename, which covers every name some rank will request under any
  layout and collapses to the old behavior when the name has no EP component.
- --overlap-grad-reduce silently corrupted adapter gradients: the TP grad sum
  runs before finalize_model_grads and writes the same main_grad buffer MCore's
  per-bucket reduce-scatter is already streaming into. It cannot move after
  finish_grad_sync (--use-distributed-optimizer leaves only the local shard
  valid), so it is refused at startup.

Guards for combinations that silently did the wrong thing:
- --experts-shared-outer-loras now requires bridge mode or a custom provider;
  native adapts no routed experts, so SGLang was being configured for buffers
  the export never fills.
- --train-backend fsdp + --lora-rank now fails instead of silently
  full-finetuning while the rollout engine expects an adapter.
- qwen3_5/qwen3_6/qwen3_next warn that raw mode's frozen-base backward is known
  to diverge there (scripts/run_lora_native.py documents grad_norm 1e7-1e10 /
  NaN), so a run says so instead of blowing up at step 1.

Also records, as roadmap notes next to the code that would change, the
dimensions Megatron-Bridge PEFT covers and native does not: routed-expert LoRA
in three layouts and router LoRA (spec/moe.py), the expert-TP group and the
structural reason overlap flags are refused (distributed.py), expert-axis
export (codec/hf.py), GDN in_proj and uncompressed-MLA q_proj (spec/attention.py),
each with the bridge file:line to port from.
create_lora_instance passed both target_modules and exclude_modules to the
bridge LoRA, but Megatron-Bridge treats them as mutually exclusive selection
mechanisms and asserts the latter is empty whenever the former is set
(peft/module_matcher.py). Since --target-modules is required whenever LoRA is
on, any --exclude-modules value made every bridge-mode LoRA run die during the
module walk with "exclude_modules should be empty when using target_modules".
Reproduced against the pinned radixark/Megatron-Bridge matcher.

The forwarded value was redundant anyway: parse_lora_target_modules already
subtracts the excluded names from --target-modules, which is the only mechanism
both LoRA paths share, and miles can never reach the bridge branch that honors
exclude_modules. So the kwarg is dropped (verified against the real matcher:
linear_qkv still matches, the excluded linear_proj returns None, no assert), and
parse_exclude_modules goes with it — its only caller was that kwarg.

Wildcard excludes now fail at startup instead of being silently ignored: exact
names are all the subtraction can remove, the native provider matches leaf
names, and bridge refuses patterns alongside a target list.

@yushengsu-thu yushengsu-thu left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fix them

Comment thread miles_plugins/lora/distributed.py Outdated
Comment thread miles_plugins/lora/spec/moe.py Outdated
@yushengsu-thu

yushengsu-thu commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

To-do

  • dsk-v4 support: because we have not support dsk-v4 lora on sglang side yet. Thus, I reomve dsk-v4 part and we can add it after sglang side support.

  • Refer to megatron-bridge to compelete the features at below layouts:

  • [I added] moe lora: in native branch, it only supports shared moe lora (bth, that's tml special design moe lora). Another general moe lora is non-shared and this is implemented in sglang and megatron-bridge already as well. We can refer to that part when we refactor megatron-bridge lora in miles. (Refer to different type of lora in megatron-bridge).

yushengsu-thu and others added 4 commits July 31, 2026 18:15
Applies the style already used in distributed.py and spec/moe.py to the places
that still carried long roadmap prose: codec/hf.py, codec/checkpoint.py,
modules/linear.py, modules/moe.py, spec/base.py, the GDN and MLA spec
docstrings, the plugin __init__, five docstrings and one comment in
lora_utils.py, parse_lora_target_modules, two test docstrings, and the docs
parallelism section.

Assertion and warning text is untouched: those are user-facing diagnostics where
naming the reason and the escape hatch is the point.

No behavior change; 177 LoRA fast tests pass.
…g/lora-plugin-refactor

# Conflicts:
#	tests/fast/fixtures/generation_fixtures.py
#	tests/fast/fixtures/rollout_fixtures.py
@Zhichenzzz Zhichenzzz added the run-ci-lora-native Run native (raw-mode) LoRA plugin e2e tests label Aug 2, 2026
Comment thread miles_plugins/lora/checkpointing.py Outdated

from miles_plugins.lora.modules.linear import iter_adapters

_MODEL_CHUNK_PREFIX = "_miles_model_chunks."

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here we don't need to make the miles-sepcific lora weights, how about aligning with the megatron-bridge naming ways?

@yushengsu-thu yushengsu-thu Aug 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use mcore checkpoint and sharding

Zhichenzzz and others added 4 commits August 2, 2026 21:48
The ci-test weight checker runs in equality mode; Qwen3.5's visual.* and
K2.5's vision_tower./mm_projector. weights are never re-shipped by the
text-only frozen-base run and are transformed by the engine at load, so raw
equality against the checkpoint cannot hold (all text weights pass). Same
pattern as the GLM-5.2 bridge CI's skip list.
Replace the per-rank .pt shard format (and per-rank optimizer pickles) for
the native provider with one MCore torch_dist checkpoint holding adapter
weights, fully_reshardable optimizer state, and scheduler/iteration:

- NativeLoRAAdapter.sharded_state_dict is metadata-gated: invisible to base
  checkpoint loads as before, and emits real ShardedTensors (TP axes from
  partition_dim) when the adapter checkpoint walk opts in.
- Standard MCore walk supplies global layer keys and PP offsets; the merged
  dict is keyed by (ShardedTensor.key, global_offset) since VPP chunks share
  chunk-local state-dict keys. A completeness assert fails loudly if a
  provider's module tree hides adapters from the walk.
- Checkpoints reshard on load across TP/PP/DP changes; optimizer fp32 mains
  and Adam moments travel with the weights, so resume no longer depends on
  rank-count-pinned files or post-load main-param refresh ordering.
- The miles-specific '_miles_model_chunks.' key namespace is gone; both
  legacy shard schemes (flat and namespaced) remain readable. Path selection
  between dist and legacy formats is agreed across ranks before entering the
  collective load.

Bridge/custom providers keep their existing shard + training-state format;
HF PEFT export for rollout sync is unchanged.
…imizer roundtrip

tests/fast (gloo, CPU): metadata gating, save/load roundtrip with optimizer
plumbing, TP2-save -> TP1-load resharding, VPP chunks with identical
chunk-local dict keys, both legacy shard schemes, and the loud-failure path
for module trees that hide adapters from the sharded walk. A scoped fixture
shims two upstream GPU/Linux assumptions in MCore's async writer for
CUDA-less lanes and macOS.

tests/fast-gpu: real MCore DDP + DistributedOptimizer (bf16 params, fp32
mains, fully_reshardable state) saved under TP2 and reloaded under TP1;
one more identical-gradient Adam step on both sides must match bitwise,
covering weights, fp32 mains, and Adam moments.
@yushengsu-thu
yushengsu-thu force-pushed the yusheng/lora-plugin-refactor branch from 60c2a0d to 9f173d8 Compare August 3, 2026 08:51

@yushengsu-thu yushengsu-thu left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

clean here

Comment thread miles/backends/megatron_utils/megatron_to_hf/__init__.py Outdated
Comment thread miles_plugins/lora/checkpointing.py Outdated
Comment thread miles_plugins/lora/checkpointing.py
Comment thread miles_plugins/lora/checkpointing.py
Comment thread miles_plugins/lora/checkpointing.py
Comment thread miles_plugins/lora/checkpointing.py Outdated
…st, fix isort

- checkpointing.py docstrings cut to abstracts per review.
- Drop the kimik25 dispatch comment per review.
- tests/fast-gpu/test_native_lora_dist_ckpt_optimizer.py: add the required
  register_cuda_ci (stage-b-2-gpu-h200, lora-native) — its absence failed the
  registry sanity check in every CPU suite — and apply isort's import order.
@yushengsu-thu
yushengsu-thu merged commit 11a3d8b into zhichen/lora-native Aug 3, 2026
37 checks passed
@yushengsu-thu
yushengsu-thu deleted the yusheng/lora-plugin-refactor branch August 3, 2026 16:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-lora run-ci-lora-native Run native (raw-mode) LoRA plugin e2e tests run-ci-megatron run-ci-miles-plugin Run CI tests labeled miles-plugin

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants