feat(peft): LoRA support for SFT

[NancyFyong]

First of a small series adding LoRA. This one covers the model wrapper and the SFT
launcher; RL comes next (see the note at the bottom). Draft because the split is worth
agreeing on before I push the rest.

Why it is this small

AutoModel already does all of it: from_pretrained(..., peft_config=PeftConfig(...))
injects the adapters before FSDP2 shards, then freezes every non-lora_ parameter
after parallelization (so params materialized during sharding are caught too). Molt
already builds the optimizer from requires_grad params, so nothing else needs to opt
out of training the base. This PR is mostly flag plumbing.

What is here

• add_lora_args(parser, prefix) — lora.{rank,alpha,dropout,target_modules}, off
  unless rank > 0, alongside the other shared blocks. Prefixed like
  add_optimizer_args so the RL launcher can reuse it unchanged.
• _lora_peft_config() in base.py, next to the existing single-call-site helpers
  (_mtp_off_kwargs, _validate_attn_implementation), which keeps the mapping and the
  guard unit-testable without building a model.
• Reject custom-MoE + --fsdp.tp_size > 1. AutoModel's safe MoE-TP path replicates
  non-expert modules and rejects PEFT, but only in a post-shard validator — i.e. after
  the whole checkpoint is loaded. Failing at flag-parse time turns minutes into
  milliseconds. Dense + TP stays allowed (it only costs the Triton LoRA kernel).
• CheckpointingConfig.is_peft derived from the model instead of hard-coded False.
  All three _build_checkpointer call sites already pass the model, so save, resume and
  consolidated export agree on the adapter-only format rather than writing base weights
  with lora_ keys mixed in.
• LoRA adapters are routed to AdamW even under --optim muon: Muon's orthogonalized
  update assumes a full-rank layer, not a rank-r factor pair.
• A trainable-share line at startup — the cheap way to notice target_modules matched
  nothing (share stays ~100%).

Default target_modules is *_proj, matching AutoModel's own implicit default. Patterns
are anchored fullmatches, so it covers dense attention/MLP linears but not custom-MoE
grouped experts, which are named *_projs — those need an explicit pattern such as *.
The docs/help text say so.

Not here, on purpose

RL. Refit pushes model.state_dict() by HF name, and vLLM's AutoWeightsLoader raises
on lora_A/lora_B; skipping them instead ships frozen base weights, so the rollout
silently stays on the initial policy. That needs merge-on-refit, which is its own PR with
its own tests. The --actor.lora.* flags are still registered so the RL launcher gives a
clear error, but it fails fast on rank > 0 rather than silently serving the base
policy.

Also deliberately absent: no use_dora flag (nothing can trigger it, and
materialize_effective_weight does not support DoRA), and no TE-experts guard —
BackendConfig.experts defaults to torch_mm, so that path needs an explicit
MOLT_MOE_EXPERTS=te, and AutoModel's own error already names TE.

Testing

tests/unit/test_lora_config.py (7 tests): defaults keep LoRA off, flags map onto
PeftConfig (rank→dim), the default/explicit target_modules, and the MoE+TP guard
including the two combinations that must stay allowed (MoE+EP-only, dense+TP).
tests/unit/test_muon_param_classify.py covers the adapters→AdamW routing.

End-to-end SFT validation

Two MoE models, LoRA r16/α32, default *_proj target (attention/dense linears; experts
stay frozen), expert-parallel EP=2 / TP=1, --model.gradient_checkpoint full,
attn=te, 2 epochs. Both drive the loss down with 0 CheckpointError and write an
adapter-only checkpoint:

Model	Base params	Trainable (LoRA)	sft_loss start → end	CheckpointError
Qwen3.6-35B-A3B (custom MoE)	35.1B	11.3M (0.03%)	2.07 → 0.011	0
Qwen3-30B-A3B-Base (HF Qwen3MoeForCausalLM → AutoModel custom MoE)	30.5B	13.4M (0.04%)	6.54 → 0.319	0

Note: MoE + EP + ac=full needs --data.pad_to_max_len (shipped separately in #61) to
keep the checkpoint recompute deterministic; LoRA itself is independent of that flag.

pytest -q clean, pre-commit clean, --help verified.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hijkzzz]

also cc @HuiyingLi

[NancyFyong]

Hi @HuiyingLi, can you take a look. Thank you!

[HuiyingLi]

Thanks for the focused LoRA plumbing—the overall SFT-only split makes sense. I tested this against the pinned AutoModel checkout and found a few issues to address before approval:

1. PEFT+EP optimizer resume is not currently correct. Saving uses AutoModel’s PEFT-aware native optimizer format, but loading constructs a different state representation. Even with is_peft=True, a fresh AdamW has no moment tensors for DCP to populate, so loading can complete while silently restoring zero optimizer states.

   Please initialize the optimizer-state schema before loading (for example with torch.distributed.checkpoint.state_dict._init_optim_state) and add a real regression covering: optimizer step → checkpoint save → fresh model/optimizer load → exact restoration of exp_avg and exp_avg_sq.

2. The four new MoE recipes depend on functionality from feat(sft): add --data.pad_to_max_len to fix MoE+EP gradient-checkpoint crash (AutoModel#3325) #61. They enable full activation checkpointing without --data.pad_to_max_len, although the PR description says this combination raises CheckpointError. That flag does not exist on this PR’s base. @hijkzzz I'd suggest merge 61 as a WAR now. This part will be refactored into AM in the future. For now it let users run molt without err. wdyt?

3. The grouped-expert target guidance is incorrect. AutoModel matches grouped expert module paths such as *.experts; *_projs describes parameter names and does not match the grouped module. To adapt both dense projections and experts, the example should use:
   --model.lora.target_modules "*_proj" "*.experts".

[NancyFyong]

Thanks for the focused LoRA plumbing—the overall SFT-only split makes sense. I tested this against the pinned AutoModel checkout and found a few issues to address before approval:

1. PEFT+EP optimizer resume is not currently correct. Saving uses AutoModel’s PEFT-aware native optimizer format, but loading constructs a different state representation. Even with is_peft=True, a fresh AdamW has no moment tensors for DCP to populate, so loading can complete while silently restoring zero optimizer states.
   Please initialize the optimizer-state schema before loading (for example with torch.distributed.checkpoint.state_dict._init_optim_state) and add a real regression covering: optimizer step → checkpoint save → fresh model/optimizer load → exact restoration of exp_avg and exp_avg_sq.
2. The four new MoE recipes depend on functionality from feat(sft): add --data.pad_to_max_len to fix MoE+EP gradient-checkpoint crash (AutoModel#3325) #61. They enable full activation checkpointing without --data.pad_to_max_len, although the PR description says this combination raises CheckpointError. That flag does not exist on this PR’s base. @hijkzzz I'd suggest merge 61 as a WAR now. This part will be refactored into AM in the future. For now it let users run molt without err. wdyt?
3. The grouped-expert target guidance is incorrect. AutoModel matches grouped expert module paths such as *.experts; *_projs describes parameter names and does not match the grouped module. To adapt both dense projections and experts, the example should use:
   --model.lora.target_modules "*_proj" "*.experts".

Thank you for review, I'll fix it quickly.

[NancyFyong]

Hi @HuiyingLi, addressed issues 1 and 3. Issue 2 (recipe dependency on #61) I'd like to leave for a follow-up — see below.

Issue 1 — PEFT+EP optimizer resume

Root cause was deeper than the schema not being materialized: save_ckpt routes through Checkpointer.save_optimizer, which builds OptimizerState(is_peft=True, has_expert_parallelism=True) and writes the native AdamW state dict; the old load_ckpt hand-rolled its own OptimizerState(model, optimizer, scheduler) without those flags, so it requested the DCP FQN shape against native-keyed data. allow_partial_load=True then silently returned zero moments. Materializing the optimizer schema before load would still leave the shape mismatch, so I routed load through the symmetric Checkpointer.load_optimizer(...) instead — both sides now build OptimizerState under the same gates, and OptimizerState.state_dict() internally materializes the moment schema (_materialize_missing_adam_state) before DCP populates it. Also dropped the allow_partial_load=True flag that kept the bug silent.

Commit: 595c54e

Real-GPU verification (8×H20, Qwen3-30B-A3B, LoRA rank 16, EP=8, MOLT_MOE_DISPATCHER=deepep):

stage	q_proj.lora_A.weight state
Run 1 save at step 2	step=2.0 exp_avg.norm=1.128390e-04 exp_avg_sq.norm=6.385179e-09
Run 2 resume from global_step4	step=4.0 exp_avg.norm=4.056379e-04 exp_avg_sq.norm=4.188193e-08
trainer log after load	Loaded the checkpoint: consumed_samples: 32

Non-zero moments, AdamW step counter restored, consumed_samples restored. Pre-fix, the same run would have loaded silently-zero moments and continued training with an effectively cold AdamW.

Regression test: tests/unit/test_checkpoint_optim_resume.py patches _build_checkpointer with a spy and asserts load_ckpt invokes load_optimizer with the same model/optimizer/scheduler/weights_path (pins the delegation invariant against future regressions).

Issue 3 — grouped-expert target guidance

You're right — AutoModel's ModuleMatcher runs an anchored fullmatch against the full dotted module path. Grouped experts live at <...>.experts (one module per MoE layer wrapping all experts as GroupedExperts/GroupedExpertsDeepEP), so the pattern is *.experts. *_projs describes the parameter names inside that module (gate_and_up_projs, down_projs) and matches no module path — the prior guidance would have adapted nothing when a user tried it.

Fixed the docstring in _lora_peft_config, the --model.lora.target_modules help text, all four recipe headers (slurm/quick_start × qwen3_30b_a3b_lora/qwen3_5_35b_lora), and the codified misclaim in tests/unit/test_lora_config.py. test_explicit_target_modules_are_forwarded now exercises --model.lora.target_modules "*_proj" "*.experts" as the canonical grouped-expert recipe.

Commit: 1592955

Issue 2 — recipe dependency on #61

Confirmed AutoModel#3325 is still open (waiting-on-maintainers) and the fuller upstream fix (agree on a common length per batch, add only the padding HybridEP needs, remove it afterward) is not yet in flight in AutoModel — the current maintainer recommendation on that issue is exactly what #61 does (fixed-length padding). So @hijkzzz's suggestion to merge #61 as a WAR still applies. I'd prefer to leave the recipe-side change out of this PR and handle it in the same commit that lands the AutoModel-side fix (or the AC-off variant if we go that direction), rather than churning these recipes twice — let me know if you'd rather I add --data.pad_to_max_len conditionally here.