[lora] Native LoRA for standard Megatron-core raw mode models - #1792
Open
Zhichenzzz wants to merge 53 commits into
Open
[lora] Native LoRA for standard Megatron-core raw mode models#1792Zhichenzzz wants to merge 53 commits into
Zhichenzzz wants to merge 53 commits into
Conversation
Zhichenzzz
requested review from
Shi-Dong,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
July 25, 2026 00:25
Zhichenzzz
force-pushed
the
zhichen/lora-native
branch
from
July 25, 2026 00:37
ea5ae5b to
1465ef1
Compare
Contributor
Author
|
Warning Parts of this comment are corrected in a later comment: #1792 (comment) Ran this on an 8×H200 devbox. Summary: the adapter math, the sharding, and the sync-to-SGLang loop all check out; the RL run proves plumbing and gradient flow but not learning quality (see the caveat at the end). Numerical verification vs dense reference
That last row matters: without it, "grads agree across ranks" could pass trivially on identical partials, so it would not actually test The tagged-param count goes 10 → 14 when sequence parallelism is on, which is exactly the row-parallel Unit tests: 24 passed ( End-to-end RL
392 is the expected count: per layer 6 (q/k/v × A,B) + 2 (o) + 4 (gate/up × A,B) + 2 (down) = 14, over 28 layers. Across 5 optimizer steps: The logprob agreement is the informative number — megatron's LoRA forward under TP2+SP matching SGLang's rollout using the adapter that was shipped over the wire. A wrong export or shard mapping would put that at 0.5+, not 1e-2. Caveat, stated plainlyReward mean stayed 0.0 (Qwen3-0.6B does not solve gsm8k at these lengths), so this run does not demonstrate reward-driven learning — only that the loop runs, the adapter reaches the engine intact, and gradients reach the optimizer. Correctness of the adapter math and the TP grad summation rests on the numerical harness above, which does not depend on reward variance. One loose end I did not chase because it is framework accounting rather than this PR: at one step Not covered by any of the above: PP > 1 (the |
LoRA so far required --megatron-to-hf-mode bridge, so models trained through miles' own model provider (raw mode) could not use it: setup_model_and_optimizer only reaches LoRA through _setup_lora_model_via_bridge, and the raw branch silently built a full-parameter model even with --lora-rank set. miles/backends/megatron_utils/lora_native.py attaches adapters directly to the mcore model the provider builds, before the Float16Module / DDP wrap, so DDP sees an already-frozen base and only allocates grad buffers for adapter params. Adapters are registered per HF projection, so the HF/PEFT export and the SGLang adapter sync are plain name joins. Covered, selected by --target-modules (HF leaf names): q/k/v_proj on the fused linear_qkv, o_proj on linear_proj, gate/up_proj on the fused linear_fc1, and down_proj on linear_fc2. mlp.shared_experts, a plain gated MLP, follows the same MLP targets. Routed MoE experts are out of scope: their adapters need a serving-side layout contract of their own, so a MoE model supplies them through its own provider. Sharding mirrors the wrapped module. On column-parallel modules A is replicated and B row-sharded; on row-parallel ones A is column-sharded and B replicated. A replicated param only sees a partial product per rank, so it is tagged at creation and its grads are summed over TP in finalize_model_grads before the DP reduce-scatter. mcore emits the fused qkv per query group (q1 q2 k1 v1 | ...), so B stays in plain per-projection order and the assembled delta is permuted at forward time, keeping export and load plain row/column gathers. Layouts this generic implementation would silently get wrong are rejected at startup with a pointer to --lora-provider-path, which lets a model plugin supply its own provider: attention_output_gate (Qwen3.5 / Qwen3-Next emit a 4th qkv slice), MLA, layers with no linear_qkv (GDN mixers), and num_query_groups < TP (mcore re-slices a query group across ranks there). Also in this change: * --lora-train-only trains adapters while rollout stays on the frozen base policy. The rollout-facing gate becomes lora_rollout_enabled(), so SGLang's enable_lora, the per-request lora_path, and the adapter sync all follow it. * Adapter params are excluded from the base weight sync and ship through the LoRA path instead; the native exporter's result is assembled across PP so every rank hands the engine one complete adapter. * Native adapter checkpoint shards include the EP rank: ranks sharing (tp, pp) hold different local experts, so the old name collided and lost shards. Every rank creates the save dir, since it may be node-local storage. Bridge-mode LoRA is untouched: it keeps its own model setup, exporter, and PP behavior.
tests/manual/lora/verify_lora_native.py builds a small real megatron-core GPTModel, attaches adapters, and checks the adapter branch against independently computed dense math at TP1 / TP2 / TP2+sequence-parallel: * a fresh adapter (B zero-init) leaves the output bit-identical * the adapter delta equals scale * B @ (A @ x) from the TP-gathered adapter, for both a column-parallel (fc1) and a row-parallel (fc2) module -- the base GEMM is subtracted out, so only our math is under test * exported tensors are identical on every TP rank, and export -> load into a fresh model reproduces both params and outputs * dL/dA == 0 while dL/dB != 0 for a fresh adapter, grads are nonzero once B is randomized, and reduce_marked_lora_grads combined genuinely distinct per-rank partials before the ranks agreed Also register Qwen3-0.6B in the runner script: on 2 GPUs it is the cheapest configuration that still exercises the whole loop end to end.
Registers the MoE checkpoint, adds an --lr knob (it was hardcoded), and guards TP
per model instead of against a single hardcoded query-group count.
Two things learned from running this:
* every one of the 30B-A3B's 48 layers is a pure routed-MoE block (no dense MLP,
no shared expert), so only attention carries adapters there. Pass
--target-modules "q_proj,k_proj,v_proj,o_proj" explicitly: with "all-linear"
the engine would allocate adapter buffers for gate/up/down that the trainer
never sends.
* Qwen3 checkpoints think before answering, so a short cap truncates the response
before it reaches \boxed{} and every sample scores 0. On dapo-math at 2048
tokens, 94-100% of responses truncated and raw_reward sat at 0.07; at 8192 the
truncation rate halves and raw_reward is 0.58. The gsm8k default moves 512 ->
2048 for the same reason.
…serve Turns "these checkpoints need their own provider" from prose into tests, using the flag values from scripts/models/*.sh: * qwen3.5-35B-A3B -- --attention-output-gate, GDN mixer layers, num-query-groups 2 * glm4.7-flash -- --multi-latent-attention (q-lora 768 / kv-lora 512) * kimi-k25_2layer -- --multi-latent-attention (q-lora 1536 / kv-lora 512) * glm5-744B-A40B_4layer -- --multi-latent-attention (q-lora 2048 / kv-lora 512) MLA has no fused qkv to slice per projection -- it is a q/kv down+up projection pair -- so the generic implementation cannot express its adapters at all. Each case asserts both the specific reason and that the message names --lora-provider-path, so someone pointing one of these at --megatron-to-hf-mode raw gets a startup error instead of silently wrong gradients.
MLA (DeepSeek / GLM / Kimi) was rejected outright because it has no fused qkv to slice. It does have a well-defined adapter surface though, so this implements it rather than deferring every MLA checkpoint to a per-model provider: q_a_proj linear_q_down_proj replicated q_b_proj linear_q_up_proj column-parallel q_proj linear_q_proj column-parallel (no q_lora_rank) kv_a_proj_with_mqa linear_kv_down_proj replicated kv_b_proj linear_kv_up_proj column-parallel o_proj linear_proj row-parallel Three things make MLA different from the GQA path: * the down-projections are replicated (TELinear parallel_mode='duplicated'), so both A and B are replicated and their grads only diverge per rank under sequence parallelism. A build that shards them instead is rejected with a pointer to --lora-provider-path rather than silently mis-sharded. * the latent layernorms are separate modules applied before the up-projections, so an up-projection's adapter sees already-normed input and recomputes nothing -- unlike the fused-layernorm linears on the GQA path. * o_proj's input width is heads x v_head_dim, not heads x qk_head_dim. The num_query_groups < TP guard and the missing-linear_qkv guard are both fused-qkv concerns and no longer fire on MLA. Verified with tests/manual/lora/verify_lora_native.py --mla at TP1 and TP2+SP (19 checks, all pass): bit-exact no-op when fresh, kv_b_proj (column-parallel) and kv_a_proj (replicated) deltas within 4e-7 relative of dense reference, TP-identical export, export/load round-trip, dL/dA == 0 under B's zero init, and reduce_marked_lora_grads combining genuinely distinct per-rank partials. The tagged-param count goes 8 -> 20 when sequence parallelism is on, which is exactly the replicated down-projection and row-parallel params joining the TP sum. Dense GQA re-verified unchanged (17 checks). This makes glm4.7-flash, kimi-k25_2layer and glm5-744B-A40B_4layer expressible; qwen3.5-35B-A3B remains out of scope (attention_output_gate plus GDN mixer layers), and the registry tests now pin that split.
It claimed grouped-expert adapters carry an ep tag, which stopped being true when routed experts left this module's scope. State what actually drives the sum -- a replicated A against a per-rank slice of B, plus the sequence-parallel case where a replicated B (or both sides, for MLA's down-projections) diverges per rank -- and note that ep remains an accepted tag for providers whose adapters are expert-parallel.
Zhichenzzz
force-pushed
the
zhichen/lora-native
branch
from
July 25, 2026 06:56
16af989 to
18fd5c7
Compare
Under `--recompute-granularity full` mcore runs each transformer block's forward inside `torch.no_grad()` and recomputes it during backward. Every native adapter param lives inside that block, so with the base frozen the checkpointed region has no grad-requiring input: autograd never enters it, the recompute never runs, and every adapter gradient comes back exactly zero. The run looks healthy the whole way through -- the adapter syncs to SGLang, the sha256 manifest matches, logprob agreement sits at 1e-2 -- because B stays at its zero init and the LoRA delta is a permanent no-op. Measured on GLM-4.7-Flash (MLA, TP2/EP4, colocated SGLang): the adapter branch ran 2820 times with `torch.is_grad_enabled()` false on every call, 470/470 adapter params had a DDP main_grad and all of them were zero, and max|lora_B| stayed 0.000e+00 across optimizer steps. With the fix: 1175 grad-enabled calls, 235/470 params with nonzero grad (the B matrices; dL/dA is still 0 while B is zero, which is the expected fresh-adapter invariant), grad_norm 0.0 -> 0.030, and max|lora_B| 0.000e+00 -> 2.003e-04 after the first optimizer step. The fix registers a forward hook on the embedding that makes the first activation a grad-requiring leaf. PEFT does the same thing as `enable_input_require_grads`, and megatron-bridge ships its own copy in `peft/recompute.py`, which is why bridge-mode LoRA was unaffected. The earlier reference numbers were taken with the frozen zero adapter, so they showed base train/rollout parity rather than LoRA parity. Re-taken at seq 2k, 8 rollouts, 32 prompts x 8 samples, lr 1e-4: | model | logprob_abs_diff | grad_norm | max\|lora_B\| | |---|---|---|---| | Qwen3-0.6B dense (all-linear, 112 modules) | 0.0131 - 0.0154 | 0.074 - 0.588 | 0 -> 1.50e-03 | | Qwen3-30B-A3B MoE (attention-only, 384 tensors) | 0.0155 - 0.0188 | 0.0069 - 0.0352 | 0 -> 1.50e-03 | | GLM-4.7-Flash MLA (470 tensors) | 0.026 - 0.045 | 0.013 - 0.119 | 0 -> 1.49e-02 (lr 1e-3) | No zero grad_norm anywhere, adapter magnitude rising monotonically, adapter sync matching on every step. GLM was also run at lr 1e-3 to grow the adapter 11x, and train/rollout agreement did not degrade -- a wrong slot in SGLang's fused `fused_qkv_a_proj_with_mqa` buffer would scale with |B|. A no-LoRA control on the same config puts GLM's base level at 0.0312-0.0314, i.e. the LoRA path adds nothing. Also in this change: - `--apply-layernorm-1p` (Qwen3.5 / Qwen3-Next) stores `gamma - 1`, so the branch's recomputed RMSNorm has to add the 1 back. Without it the adapter is fed a differently scaled activation than the base GEMM sees. - `run_qwen3_lora_native.py prepare` could never have run: it called `convert_hf_to_torch_dist.py` without sourcing `scripts/models/<type>.sh`, so `--num-layers` was None. Now goes through `U.convert_checkpoint` like every other runner. - `attention_output_gate` (Qwen3.5 / Qwen3-Next) is supported instead of rejected. mcore holds a query group as `[q heads][gate heads]` while HF's q_proj interleaves them per head, and the existing row permutation expresses both. - Layers whose mixer is not a fused qkv (linear-attention / GDN) no longer reject the whole model; they carry no attention adapter and are reported. - Adapter names are read off the checkpoint's own weight index instead of being hardcoded: Qwen3.5 nests the decoder under `model.language_model.layers.`, and DeepSeek / GLM / Kimi spell the shared expert `mlp.shared_experts.`. SGLang keys off `layers.<n>.` plus the leaf name so it tolerated the old spelling, but the export is also a PEFT adapter and `load_lora_adapter_hf` looks names up exactly. - The export log carries max|lora_B|, which is what separates "the sync works" from "the sync carries anything". - `verify_lora_native.py` gains a `--gate` configuration, a fused-qkv delta check taken after mcore's own per-group split, and a DDP stage -- the harness ran on a bare GPTModel, which is exactly why it missed this. All six configurations pass (TP1 / TP2+SP x plain / gated / MLA). - Adapter save/load round-trip (`--lora-adapter-path`) could not have worked, and is the one item the PR listed as uncovered. Three faults: `save_lora_checkpoint` wrote `adapter_model.bin` while `load_lora_adapter_hf` reads `adapter_model.safetensors`; the HF-format export went through `AutoBridge.export_adapter_weights`, which cannot see `NativeLoRAAdapter` modules, so raw mode wrote the wrong contents; and the adapter path was also handed to SGLang as `lora_paths`, so the first weight sync was rejected with "Failed to load LoRA adapter miles_lora because it is already loaded". Raw mode now exports through the native provider, writes safetensors, and leaves engine registration to the first sync (which runs before any rollout). Verified on GLM-4.7-Flash at lr 1e-3: 470 tensors saved with max|lora_B| 3.998e-03, reloaded bit-identically, `adapter sync OK: 470/470`, and the resumed run trains at logprob_abs_diff 0.029-0.038 with no zero grad_norm. - `scripts/run_lora_native.py`: GRPO LoRA recipe for the MoE registries (MLA and the gated Qwen3.5 hybrid).
… names Bridge-mode LoRA shipped a corrupted base weight for MLA's fused down-projection, so the rollout engine's forward diverged from Megatron's. On Kimi-K2.5 this showed up as train_rollout_logprob_abs_diff 0.53 against a 0.0122 no-LoRA baseline, and `--check-weight-update-equal` flagged `...self_attn.fused_qkv_a_proj_with_mqa.weight` (2112x7168, 2112 = q_a_proj 1536 + kv_a_proj_with_mqa 576) with max_abs_err 2.64 in the LoRA run and 0 without LoRA. megatron-bridge PEFT reports a wrapped module's *base* weight under the wrapper name `<module>.to_wrap.weight`: it unwraps for HF-name matching but yields the wrapped megatron name in the (hf_name, weight, megatron_name) triple. Miles then groups the MLA pair in `_stream_atomic_units` via `megatron_name.endswith(suffix)` against unwrapped suffixes, so under LoRA the match fails and the pair is never grouped. SGLang fuses the two into `fused_qkv_a_proj_with_mqa`, so once the halves land in different chunks a fused write leaves the rest of the tensor stale. The failure is silent because `assert not pending` cannot fire -- neither param ever enters `pending`. The distributed iterator already strips this segment immediately before the same grouping call (`_get_weight_transfer_update_units`); only the bridge/from-tensor path missed it. Stripping it also restores the fp8/nvfp4/mxfp8 quantizer regexes, which dispatch on the megatron name too (DeepSeek-V4 FP8; Kimi's compressed-tensors path keys off the HF name and was unaffected). Verified on Kimi-K2.5-2layer bridge LoRA, TP8/EP8, seq 2k: logprob_abs_diff 0.530 -> 0.01227, matching the no-LoRA baseline, with adapter sync 20/20 sha256 on all 8 ranks. Both halves of the mechanism are covered by the new test: plain names group into one unit, `.to_wrap.` names must too. 267 megatron_utils fast tests pass.
Deletions only, no code change. The explanations these carried are already in the module and function docstrings (`lora_native.py`'s projection tables and `save_lora_checkpoint`'s two-format description), so inline they were duplication. Section banners stay: all three files already use them (16 / 4 / 5 pre-existing), so dropping only the new ones would leave the added sections as the only unlabelled ones in each file.
Zhichenzzz
changed the title
Dissolve spec/inkling.py into the role files so each family's specs sit next to the base classes future architectures inherit from: InklingAttentionSpec -> spec/attention.py, InklingDenseMLPSpec -> spec/mlp.py, InklingMoESpec -> spec/moe.py, InklingExtrasSpec -> new spec/extras.py (model-level adapters). Registry imports hoist to top level; drop a dead parallel_state import.
….lm_head The spec is only used by registry assembly, so it lives next to _inkling_arch_spec instead of its own file; the slot renames extras -> lm_head.
Inkling-Small (full 42-layer) 4-node LoRA run on 32xH200 (TP4/EP4/PP8, 16-GPU engines): train_rollout_logprob_abs_diff 0.0091-0.0123 (mean 0.0107) over 10 rollouts, reward stable on dapo-math.
Qwen3.5-9B (dense hybrid GDN) 20-rollout native run: logprob_abs_diff 0.0045-0.0090 (mean 0.0063), grad_norm <= 0.032 throughout — the recorded raw-mode backward divergence does not reproduce, matching the earlier Qwen3.5-35B-A3B result. qwen3_6*/qwen3_next stay UNSTABLE until run.
MODEL_SPECS maps model_type directly to LoRAArchSpec; validation evidence lives in the PR description, not in code.
Nothing on main references miles.backends.megatron_utils.lora_native; the dotted path only ever served this branch's own history. Also retire a docstring reference to the removed Inkling provider dispatch.
…dead code
LoRAArchSpec gains sglang_lora_target_modules; the Inkling entry declares
('all',) and serving.sglang_target_modules() answers launch/sync target
lists for every native run. Core loses the inkling-aware
sglang_lora_target_all_sentinel and its chat-template import, and
sglang_engine's two sentinel branches collapse back to one call each.
Also delete the superseded private _pp_assemble_full_adapter copy in
update_weight_from_tensor (call sites use the shared lora_utils one) and
a duplicated tuple element in chat_template_utils/inkling.
Module docstrings shrink to one line; comment blocks compress to the constraint they state; Unsupported/TODO essays leave (the asserts carry the boundaries).
# Conflicts: # tests/ci/labels.py
…ayout Per review, run_lora_native.py moves to examples/lora/ together with its siblings run_qwen3_lora_native.py and run_glm5_2_744b_a40b_lora_native.py (CI imports and usage strings updated). Test files now mirror the module they test: test_lora.py splits into test_modules/test_hf_adapter/test_specs under the plugin dir, the lora_utils seam classes join test_lora_utils.py, and the rollout gate moves to tests/fast/utils/test_lora.py. Renames: test_native_dist_checkpoint -> test_checkpointing, test_native_lora_dist_ckpt_optimizer -> test_lora_checkpointing, test_bridge_atomic_group_unwrap -> test_hf_weight_iterator_bridge.
main's compute-server-args tests drive the LoRA branch with a fake --hf-checkpoint; only a readable config.json routes through the plugin's spec-aware target resolution, everything else keeps the plain HF name conversion. resolve_model_spec still fails closed at model build.
yushengsu-thu
approved these changes
Aug 4, 2026
Comment on lines
104
to
+156
| @@ -116,14 +145,15 @@ def _stream_atomic_units(items, atomic_update_groups): | |||
| None, | |||
| ) | |||
| if match is None: | |||
| yield [(hf_name, weight)] | |||
| yield tensors | |||
| continue | |||
| group, idx, suffix = match | |||
| prefix = megatron_name[: -len(suffix)] | |||
| slots = pending.setdefault((prefix, group.key), [None] * len(group.suffixes)) | |||
| slots[idx] = (hf_name, weight) | |||
| assert slots[idx] is None, f"duplicate {suffix} for {prefix} in atomic group {group.key}" | |||
| slots[idx] = tensors | |||
| if None not in slots: | |||
| yield list(slots) | |||
| yield [pair for slot in slots for pair in slot] | |||
Collaborator
There was a problem hiding this comment.
if you modified this, could it still pass lora e2e case in megatron-bridge path?
# Conflicts: # docs/advanced/lora.md # miles/utils/arguments.py
mkdir/hf-download prep goes through exec_command_cpu and the kimi INT4 dequant through exec_command_gpu, matching scripts/run_kimi_k25.py.
…to pure mapping convert_target_modules_to_hf and its tables move to miles_plugins/lora/hf_adapter.py (core lora_utils re-exports them), so the plugin no longer imports core for name normalization. InklingLMHeadSpec moves out of registry.py into spec/lm_head.py, and the empty kernel/ placeholder goes away.
…rt/unused-import fixes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
LoRA in miles requires
--megatron-to-hf-mode bridge, so models trained through miles' own model provider (raw mode) cannot use it at all:setup_model_and_optimizeronly reaches the LoRA setup through_setup_lora_model_via_bridge, and the raw branch silently builds a full-parameter model even when--lora-rankis set.This PR adds a native LoRA path for raw mode, shipped as a plugin (
miles_plugins/lora/). It attaches adapters directly to the mcore model the provider builds, before the Float16Module / DDP wrap, so DDP sees an already-frozen base and only allocates grad buffers for adapter params. Bridge-mode LoRA is untouched — it keeps its own model setup, exporter, and PP behavior.The branch absorbed three follow-ups as it matured: #2017 (extraction into the plugin), #2072 (function-level redesign + fixes from validating four models end-to-end), and — after merging main's Inkling LoRA (#2122) — a rewrite of Inkling onto the plugin architecture, replacing the model-specific monolith.
Layout
Design
Classes carry the taxonomy; every fact is declared once and derived everywhere else.
spec/layout.py): an architecture spec is a class holding aModuleLayouttable — which projections exist, on which physical linear, with which shard geometry. One sharedattach_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.LayoutSpec→AttentionSpecBase→GQA/MLA/GDN/Inkling,HybridGQAGDN(GQAAttentionSpec),SharedOuterExpertMoESpec(GeneralExpertMoESpec)):supported_targets, the canonical--target-modulesorder, 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/AttentionFamilystr-enums replace bare strings; no parallel name tables in the exporters.num_query_groups < TP, where mcore re-gathers and re-slices qkv) fail at startup with a pointer to bridge mode instead of mis-sharding silently.exports() -> ProjectionExport(withSGLangFusedGroupmetadata);sglang_adapter.pyconsumes only that API.export_plan(gather)/load_plan_custom(take);hf_adapter.pystays generic.ParallelGatherbatches async all-gathers over the TP or EP group behind onerequest(tensor, dim, group=...)call.ModuleLayout.hf_block_prefixandProjectionBinding.adapter_classoverride HF naming and the module class per binding;LoRAArchSpeccarriesmoe.attach(...)and an optionallm_headspec so a fully custom family still routes through the one orchestrator inlora.py.default_target_modules(hf_checkpoint)andpreflight_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.Two numerics details worth calling out:
q1 q2 k1 v1 | q3 q4 k2 v2 | ...). The split adapter keepsBin plain per-projection order and permutes the delta at forward time, so export/load stay simple row gathers.Ais replicated across TP while each rank holds only its slice ofB, so every rank computes a partialdL/dAand the true gradient is their sum; the same applies to replicatedB(row-parallel) and to MLA's replicated down-projections once sequence parallelism shards the sequence. Tagged at creation and summed infinalize_model_gradsbefore the DP reduce-scatter; a no-op when nothing is tagged.Architecture coverage
llama,qwen2,qwen2_moe,qwen3,qwen3_moe,mimo,glm4,glm4_moeqwen3_5,qwen3_5_moe,qwen3_6,qwen3_6_moe,qwen3_nextdeepseek_v3,deepseek_v32,glm4_moe_lite,glm_moe_dsa,kimi_k2,kimi_k25,joyai_llm_flashinkling_mm_modelWhich of these have an end-to-end run on record is the Validation section below.
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 (InklingAttentionSpecinspec/attention.py,InklingDenseMLPSpecinspec/mlp.py,InklingMoESpecinspec/moe.py,InklingLMHeadSpecinspec/lm_head.py):[q;k;v;r]declared as oneFusedAttach— 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 theProjectionSpecs.gate_up_projprojection via anadapter_classoverride whoseexport_plangathers the gate/up halves separately (a plain dim-0 TP gather would interleave ranks).LoRAGroupedFC1/FC2— shared-A, per-expert-B grouped-GEMM deltas (F.grouped_mmover 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.LoRAOutputHeadwith muP logit scaling and unpadded-vocab trim.sglang_lora_target_modules=("all",)andsglang_adapter.sglang_target_modules()answers launch/sync target lists for every native run — core carries no model-specific serving logic.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 ownKimiK25VLbridge and never converts). ThinDeepseekV3Bridgesubclass:text_configdescent +language_model.prefix remap.megatron_to_hfkimi dispatch:_convert_to_hf_corematched onlykimi_k25, but the direct weight iterator passeskimik25config(dead branch); the converter also lacked the raw text-only provider's unprefixed top-level weight names.convert_kimi_int4_to_bf16.pystripsquantization_config: a surviving one sends SGLang through its CompressedTensors path, which serves the BF16 checkpoint with a degenerate context-free forward → garbage rollouts, train/rolloutlogprob_abs_diff~2.2. Verified by a megatron-vs-HF-vs-SGLang three-way logprob comparison.torch_dsa_topkclamps k to the key length: short sequences (<index_topkkeys) crashedtorch.topkin 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).Also in this change
--debug-lora-train-only: train adapters in Megatron while rollout stays on the frozen base policy. The rollout-facing gate becomeslora_rollout_enabled(), so SGLang'senable_lora, the per-requestlora_path, and the adapter sync all follow it. Useful for isolating training-side numerics.--check-lora-weight-equal: sha256 comparison of every adapter tensor between Megatron's export and the engine's received copy, verified to fail (not pass vacuously) under single-tensor fault injection.CI
tests/e2e/lora_native/(labellora-native, triggerrun-ci-lora-native): qwen3, qwen3.5, glm4.7, glm5.2-5layer, kimi-2.5-2layer, and inkling-small-4layer — native runs checking logprob parity and weight-update equality (vision towers and engine-transformed params skip-listed).Validation
Numerics vs dense reference (
tests/manual/lora/verify_lora_native.py): fresh-adapter no-op exact 0.0; column/row/MLA deltas at 1e-7 relative error; export identical on every TP rank; export→load round-trip; replicated-param grad sums exercised on distinct per-rank partials — across TP1, TP2, TP2+SP, ± MLA. The plugin extraction was additionally checked side-by-side against the pre-plugin implementation: bit-identical adapter params, forward outputs, backward grads, and HF exports.End-to-end: 20-rollout native GRPO runs (wandb
ch271828n-team/miles_lora_native), acceptance ≈ 0.01 train/rolloutlogprob_abs_diff:The Inkling run exercises the whole custom stack at multi-node scale: 4-way fused [q;k;v;r] export, grouped-GEMM expert adapters under EP4, shared-experts hooks, the muP lm_head, and SGLang serving via the
["all"]sentinel. (SGLang note: LoRA-on-MoE serving currently supports only the triton MoE runner —flashinfer_trtllm_routedfails engine init with NotImplementedError.)tests/fast/miles_plugins/lora: 64 passed (layout locks, target normalization, guard behavior, Inkling spec). Adapter attribute names are pinned, so checkpoint keys are unchanged.