feat(fsdp): ref model refactor, frozen ref/teacher models, EMA optimizer, pinned offload

[Rockdu]

What

• Independent frozen reference and teacher models. --ref-mode ref --ref-load <ckpt> and --teacher-load <ckpt> load their own frozen FSDP models through the actor's loader (same components, mesh and input prep), each optionally with a pretrained PEFT adapter (--ref-lora-adapter-path, --teacher-lora-adapter-path) and native FSDP CPU offload (--ref-cpu-offload, --teacher-cpu-offload). Teacher models are exposed to custom losses as ctx.teacher_models; no loss consumes them yet.
• EMA is a checkpointed torch.optim.Optimizer. EMAOptimizer keeps state[param]["ema"] next to its parameter, is saved/loaded through DCP under ema/ (independent of --no-save-optim), cold-starts from the loaded weights on older checkpoints, and updates once after each rollout's training instead of on every weight publication. use_weights() swaps EMA values into the actor's shards in place (no full copy of the trainable weights).
• Phase offload uses pinned CPU memory. --offload-train sleep/wake now moves the actor, reference, teacher, AdamW and EMA state between GPU and pinned host memory, preserving Parameter objects, optimizer bindings, FSDP shard views/padding and shared buffers.
• Pretrained actor LoRA. --lora-adapter-path resumes training from a saved PEFT adapter; --lora-rank/--lora-alpha must match it, and LoRA IPC weight sync reads rank/alpha from the loaded adapter config.
• --ref-mode lora_base is kept (adapter disabled on the actor, no second base model).

Why

• DMD2-style training needs a teacher and a reference that are real models, not views of the actor (EMA / adapter-disabled base).
• EMA as optimizer state follows torchtitan: it gets resharding-safe DCP checkpointing for free and survives resume.
• Pinned host memory makes sleep/wake copies asynchronous DMA instead of pageable staging.

Three non-obvious traps handled here:

• PEFT + FSDP2 drops adapter gradients. With reshard_after_forward=False, disable_adapter() re-enables adapter requires_grad on exit while the gathered parameters are registered, leaving the shards frozen; the next training step then has no LoRA gradients. The lora_base forward reshards before leaving the context (regression test: test_lora_base_reference_keeps_adapter_gradients_on_fsdp_shards).
• FSDP2 lazy init with nothing trainable. FSDP2 fixes each group's gradient dtype at its lazy init from the parameters that require grad, and that init normally runs at the first forward. If the first forward is the lora_base reference (adapters disabled, base frozen), the dtype is left unset and fp16/bf16 gradients are assigned to the fp32 shards (attempting to assign a gradient with dtype 'c10::Half' to a tensor with grad_dtype 'Float', seen in test_sd3_ocr_grpo_2xGPU). load_fsdp_models now finishes the lazy init at load time, while the trainable set is final, so the order of reference/teacher/actor forwards no longer matters (regression test: test_first_forward_with_adapters_disabled_keeps_master_dtype_gradients). The forward order itself is kept as on main (actor first).
• Stale gathered weights. FSDP2 keeps the root's gathered parameters after forward, so an EMA forward would leak EMA values into the next actor forward. use_weights() reshards before swapping in and before swapping back.

Validation

Pinned vs pageable phase offload — Qwen-Image Flow-GRPO recipe (run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py), 5×H200, LoRA, colocate, --deterministic-mode, steady-state rollouts 1–3 (measured on an earlier revision of this branch's offload code):

per rollout	pinned	pageable
step time (s)	304.9 / 297.3 / 297.2	312.0 / 304.7 / 304.1
sleep (s)	4.45 / 4.00 / 3.98	10.05 / 10.10 / 9.87
wake_up (s)	0.81 / 0.82 / 0.82	3.86 / 3.92 / 3.99
actor_train (s)	135.9 / 130.4 / 130.5	135.5 / 129.6 / 129.9

Current offload implementation — standalone 4-rank FSDP2 Qwen-Image + recipe LoRA (3373 tensors, 20.9 GB/rank), steady state: offload 0.58 s → 0.40 s, onload 0.57 s → 0.40 s vs the previous revision. Raw pinned D2H is 42–50 GB/s per rank with 4 ranks copying concurrently, so the copies are near hardware speed.

Known cost not addressed here: under torch.use_deterministic_algorithms(True) (i.e. --deterministic-mode), PyTorch NaN-fills every new pinned buffer (torch.utils.deterministic.fill_uninitialized_memory), which adds ~2.3 s to each sleep on this recipe. Without deterministic mode the offload is ~0.6 s.

E2E standard change: test_sd3_nft_pickscore_3xGPU (re-recorded in fix(fsdp): update EMA once after training ...)

The SD3 NFT e2e (--ref-mode ema --use-ema --ema-rollout-policy ema, --ema-decay-flat-steps 0, so the decay changes with every EMA update) drifts from step 2 on, in 8 of its 10 series; steps 0–1, raw_num_samples and nft_t_mean are unchanged. Cause:

• On main, EMA updates inside update_weights, which also runs once before training starts, so after rollout k the update count is k+2; it now updates only after training, so the count is k+1 and every decay is one step earlier.
• The update is now torch._foreach_lerp_ instead of mul_(d).add_(x, alpha=1-d) (same math, different rounding).

Step 1's fp16 grad-scaler overflow leaves the weights unchanged, so EMA equals the actor until step 2 whatever the decay; hence the drift starts there.

Attribution, recorded with record-e2e-standards.yml on the verifying ["h200", "3gpu"] runner class: this branch plus a throwaway commit restoring main's update cadence and formula reproduces the committed main standard bit for bit, on both the pre-#257 and the current main base. This branch's head reproduces the values the PR CI reported (e.g. rollout/reward/raw_mean[2] = 0.7545289993286133); that recording is the new standard. No other standard changes.

Correctness

• 4×H200: offload/onload round trip preserves shape, stride, dtype, requires_grad, DTensor spec, Parameter identity, channels_last / non-contiguous buffers and shared-buffer aliasing; everything is pinned after offload.
• 2×H200: tests/fast-gpu/backends/fsdp_utils/test_offload.py and test_reference_models.py passed on the offload revision; not re-run after the last two changes (lora_base / use_weights rewrite, rollout-default revert).
• CPU: every commit passes its focused tests on its own (last commit: 56 passed across test_offload, test_ema_optimizer, test_reference_models, test_reference_arguments, test_loss_hub_nft, test_sglang_parallel_defaults).
• tests/fast-gpu/.../test_reference_models.py needs PYTHONPATH=<repo> for its worker to import tests.*; that also fails the same way without this PR.

Files

• miles/backends/fsdp_utils/model_loader.py — new: shared FSDP model loader (mechanical extraction from actor.py, then frozen roles + adapters).
• miles/backends/fsdp_utils/offload.py — new: pinned-memory model/optimizer phase offload.
• miles/backends/fsdp_utils/ema.py — EMAOptimizer, use_weights, reshard_model.
• miles/backends/fsdp_utils/actor.py — load ref/teacher, EMA wiring, sleep/wake, reference forwards.
• miles/backends/fsdp_utils/checkpoint.py — save/load EMA under ema/.
• miles/backends/fsdp_utils/model_backend.py, models/ltx/loading.py, models/package.py — load_component takes an explicit checkpoint_path.
• miles/backends/fsdp_utils/diffusion_update_weight_utils.py — IPC LoRA metadata from the loaded adapter.
• miles/backends/fsdp_utils/loss_hub/types.py, loss_hub/flow_grpo.py — expose reference/teacher models; KL error message.
• miles/utils/arguments.py — new flags and validation.
• docs/fsdp_model_roles.md — model roles and offload lifecycles.
• tests/fast/..., tests/fast-gpu/... — CPU and 2-GPU coverage for the above.

Commits

1. refactor(fsdp): [mechanical] extract actor model loading — pure move, reviewable separately.
2. feat(fsdp): parameterize model loading for frozen roles
3. feat(fsdp): offload model and optimizer tensors to pinned CPU memory
4. feat(fsdp): load pretrained actor LoRA adapters and publish their metadata
5. fix(fsdp): isolate base-reference inference from actor training state
6. feat(fsdp): load independent frozen reference and teacher models
7. feat(fsdp): store EMA in a checkpointed optimizer
8. fix(fsdp): update EMA once after training instead of on every weight publication — also re-records the SD3 NFT e2e standard (see above)
9. docs(fsdp): explain model roles and offload lifecycles

Not in this PR: renaming --ema-rollout-policy to --rollout-weights (stacked follow-up PR).

Checklist

• pre-commit run --all-files passes — only run on this branch's changed files (--from-ref origin/main), which pass
• Added/updated tests for new behaviour
• pytest -x is green — only the focused CPU suites listed above were run (local env has no sglang); 2-GPU suites not re-run after the last two changes
• If launch flags changed, python3 train.py --help still parses — not run
• If a public flag was added, it appears in the CLI reference docs — not yet: the new --ref-load, --teacher-load, --ref-lora-adapter-path, --teacher-lora-adapter-path, --ref-cpu-offload, --teacher-cpu-offload, --lora-adapter-path and --ref-mode ref are missing from docs/user-guide/cli-reference.md
• If an example was added, it has a real walkthrough — n/a, no example added