[listar2000]

Merge with rebase, not squash (three fork ids: hook-006, ext-030, fix-044).

Outcome

DeepSeek-V4.1 can train small adapter weights (LoRA), resume them correctly, and fit a 256k-token batch on one trainer node instead of two, with the model-version and sample checks still on.

Full-length V4.1 batch (64 samples, mean 139k tokens)	Full fine-tune, 2 nodes	LoRA, 1 node
Peak GPU memory per GPU	near the limit	206 GiB of ~255 usable
Train step	46 min	27 min, then 18 min
Checkpoint save	~53 min	2 s

Terms: LoRA trains small adapter matrices while keeping the base weights frozen. Fully async training overlaps sample generation with training. Staleness measures how far a sample's model version trails training; sample ownership checks track sample consumption and drops.

	Before	After
LoRA on DeepSeek-V4.1	not possible (no V4.1 description, and the inference engine has no V4.1 LoRA)	a recipe lora: block turns it on
What the inference engines receive	adapter weights, through engine-side LoRA	the full weights with the adapter merged in, like a full fine-tune
Weight-version, staleness and sample-ownership checks	off under LoRA	on, as for full fine-tuning
Resume of a LoRA run	data restarted at rollout 0 and the optimizer was reset	point --load at the run's own saves: adapter, optimizer, rollout position, sample counts and weight version all continue

GPU check on a 4-layer V4.1 slice (1x B300:8 and 1x H200:8; W&B mobiuslab/miles-smoke):

Check	Result
Engine serves the merged weights	trainer/engine log-prob gap 0.081 with large random adapters vs 2.70 when the sync is skipped
Version per sync	engine version rises by exactly 1 per sync
Exact bytes	first-sync bit check equal on all 4 engine TP ranks (BF16 engine)
Layouts	TP2/PP2/CP2 vs TP4: gap 2.7099 vs 2.7089, grad norm within 0.18%
Recompute	every adapter B moves on every rank; grad norm with and without recompute 0.1194963 vs 0.1194966
Release format (FP8 + MXFP4 engine)	merged sync works, version +1 per sync; gap 0.125 at zero adapters, 0.187 with large ones
Learning on a fixed replayed batch (5 steps, lr 1e-4)	SFT loss 4.2687 → 4.2239 every step down; policy loss 0.0067 → −0.0273 every step down
Resume from the run's own saves	resumes at iteration 1; optimizer restored on 4/4 ranks; first rollout after resume is 2 with 0/16 prompts repeated (old resume: 16/16); engine version continues at 2 (on B200:8)
One trainer node, full model, full-length batch	a saved batch from a 2-node full-parameter V4.1 run (64 samples, mean 139k tokens, 256k max), trainer only, on 1x B300:8 (TP4 CP2 PP1 EP8), rank 32, optimizer on: 2 steps, no OOM, peak reserved 206 GiB, host RAM 780 GiB, adapter save 2.1 s
End-to-end RL on Qwen3.5-9B (agentic, Sail sandboxes)	10 steps: train/rollout log-prob gap 0.0072–0.0082 vs 0.0075–0.0088 for full-param on the same recipe; KL 3e-4 both; TIS clip 0 (W&B mobiuslab/miles-lora-parity)
Sync memory	+7.4 GB reserved on the first sync, same as full-parameter (+7.15 GB)

Adapters cover V4.1's attention projections (wq_a, wq_b, wkv, wo_b) and the shared expert (w1, w3, w2); routed experts stay frozen. Nothing changes for runs without a lora: block.

Before the first real LoRA run (a result anyone relies on): a short end-to-end V4.1 LoRA RL run on the full model, and the owner's go for that run. Merging this PR changes nothing for existing runs.

Contract

• New recipe block lora: {rank: 0, alpha: 32, target_modules: null}. rank: 0 (default) is off. A positive rank needs the megatron backend in fully_async mode and refuses optimizer.offload and rollout.expert_qat.
• New Miles flags: --lora-merged-sync and --custom-lora-spec-path <module>.
• Preset deepseek-v4.1-flash gains lora_spec_path = "miles_ext.lora.dsv41".
• LoRA checkpoints carry training_state["model_companions"]; older checkpoints load as before.
• Resume a LoRA run with --load <the run's --save dir> (picks the latest iteration that has both iter_N/adapter and rollout/N). --lora-adapter-path alone seeds a new run from an adapter.

Side-effects

Upstream files touched, each spot marked [fork hook-006]: miles/utils/lora/arguments.py, miles/utils/lora/utils.py, miles/utils/weight_version.py, miles/utils/arguments.py, miles/backends/megatron_utils/actor.py, miles/backends/megatron_utils/model.py, miles/backends/megatron_utils/lora/utils.py, miles/backends/megatron_utils/megatron_config.py, miles/ray/rollout/rollout_executor.py, miles/rollout/session/config.py, tests/fast/rollout/session/test_config.py, tests/fast/fixtures/session_fixtures.py, plus the new test tests/fast/backends/megatron_utils/test_lora_merged_sync.py. Marked [fork fix-044]: miles/backends/megatron_utils/checkpoint.py and tests/fast/backends/megatron_utils/test_checkpoint.py.

Without --lora-merged-sync, existing LoRA serving and weight-sync selection remain unchanged. LoRA training checkpoints also save and restore model companions regardless of that flag; older checkpoints without the key remain readable.

Found on the way, not caused by this PR: with --check-weight-update-equal on a release-format (MXFP4) V4.1 engine, the first sync crashes in SGLang's mxfp4_flashinfer_trtllm_moe.py:248 (restore_weights_before_loading, CUDA index out of bounds) right after the checker's reset_tensors. A full-parameter run crashes the same way; with the checker off, LoRA and full-parameter syncs work. Keep the bit checker off on MXFP4 engines. Tracked in Linear MOB-189 (https://linear.app/mobius-lab/issue/MOB-189).

Patch page: fork/patches/hook-006-lora-merged-sync.md (includes two sync risks: the hook writes upstream's private _HF_LORA_MODELS, and every LoRA checkpoint gains a model_companions key).

Done

• hook-006 --lora-merged-sync: the weight sync sends merge_lora_into_weights(model, weights) (upstream's fp32 merge from the frozen base) through the full-parameter path. lora_rollout_enabled is false, so engine-side LoRA stays off. The version checks, the CI engine-version check and the sample-ownership checker stay on. The weight version comes from the model companion (bumped every optimizer step) and is saved in the LoRA checkpoint. When --load holds the run's own LoRA saves, startup loads the base from --ref-load, the adapter and optimizer from the latest fully saved iteration, and the rollout side from the same directory, so the full companion state (sample counts too) is restored; an --lora-adapter-path seed restores the weight version only. The session server's config carries lora_merged_sync, so agentic requests never name an adapter the engines don't have. Bridge mode, --lora-train-only and multi-LoRA are refused at startup; --colocate and --offload-train (also when set through --offload or --use-critic) are refused after Megatron's validation, once those flags are resolved.
• hook-006 --custom-lora-spec-path: imports a module that registers an architecture in the native registry and the HF target table; it never replaces a shipped spec.
• ext-030 miles_ext/lora/dsv41.py: the V4.1 spec (attention runs full-sequence; wq_a and wkv replicated, guarded by full width as in upstream's MLA spec; wq_b column; wo_b row; fused shared-expert FC1 split into w1/w3; w2 row).
• fix-044: the raw-mode mbridge HF load (Qwen3.5 starts from HF) hides NativeLoRAAdapter children, which mbridge otherwise rejects as unknown parameters (lora_fc1_adapter.gate_A). V4.1 loads torch_dist and was unaffected.
• ext-030 recipe lora: block emitting --lora-rank --lora-alpha --lora-merged-sync [--target-modules] --custom-lora-spec-path.

Verification

• python3 fork/tools/ledger.py check: ok.
• Run by the author (the tests need torch): full CPU suite (pytest tests/fast miles_ext/tests, in the miles-mobius image): 3384 passed, 30 skipped; 2 failed + 1 error are test_session_server_benchmark[v1,v2] and test_agentic_rollout_benchmark, which fail the same way on main. With main's production files, both new test files fail (No module named 'miles_ext.lora', cannot import name '_validate_lora_merged_sync').
• test_targets_default_to_attention_and_the_shared_expert lists the seven targets literally; dropping attn.wo_b from the spec makes it fail.
• Megatron parse_args on 1x L4 with the recipe's LoRA argv: passes; adding --offload is refused with the merged-sync message.
• pre-commit on the changed files: clean except miles_ext/recipes/recipe.py, whose isort/black findings already fail on main.
• New tests: zero-init adapters leave every output bit-identical; merged weights reproduce the adapter forward; on FP8-grid base weights the merged sync changes exactly the adapted tensors (w3 half untouched); a TP-sharded wkv is refused; routed experts and wo_a are refused as targets; the recipe's LoRA argv passes upstream's native-LoRA startup checks; merged mode turns the weight-version gate (lora_without_weight_versions) off, which keeps the staleness assertion and lets the sample-ownership checker run (the engine-version check itself is exercised by the GPU slice, not a CPU test); a LoRA --load resumes the latest fully saved iteration and restores sample counts, while an adapter seed keeps them fresh; the session server requests no adapter under merged sync; the HF load hides native adapters.
• Red/green, run by the author in the miles-mobius image on CPU: these four tests pass on the head (82 passed with test_lora_merged_sync.py, test_config.py and the new checkpoint test). With the production files from before each fix, each fails on its own assertion, not on an import: the restore test (empty sample_consumptions instead of [[0, 0, 1, 1, 0]]), the --load test (adapter path /seed-adapter instead of iter_0000000/adapter), the session test (lora_rollout_enabled true) and the HF-load test (lora_fc1_adapter.gate_A reaches the loader).