feat(ltx): consume rollout sigmas exactly and train with fp32 master weights

[Rockdu]

Stacked on #91 (feat/fsdp-precision-control). Review the last 4 commits (fdd65ee and up); everything below them is #91's diff and lands there. fdd65ee fixes a #91 wrap-plan bug this branch needs to run and can be absorbed into #91 directly.

What

• Carry the rollout scheduler sigmas alongside the trajectory timesteps into every DiT forward: PreparedBatch gains a sigmas field (exact-matched from the rollout snapshot) and compute_noise_pred receives both timesteps_input and sigmas_input; each family reads the domain its model consumes (sd3/wan → t, ltx/qwen → σ). The needs_timestep_scaling knob and its rescaling branches are deleted.
• Switch the LTX-2.3 recipe to --fsdp-master-dtype fp32 and re-record its e2e standard.
• Fix refactor(fsdp)(4/5): per-parameter weight-precision control for FSDP2 #91's wrap plan on PEFT-wrapped models: the plan is compiled on the raw component, so block-dtype lookups must resolve through the base_model.model. prefix (KeyError on every LoRA run otherwise).

Why

• The trajectory t_host and the σ the rollout DiT consumed (t_device) are computed separately in sglang-d and sit 1–2 fp32 ULPs apart, so no rescaling of one recovers the other — each domain's exact original is a different array. Dividing t back down was the last known timestep gap (LTX AdaLN input rel 1.8e-3 → now bit-exact).
• Requires the sigmas snapshot in dit_trajectory (sglang #32683, miles refactor(diffusion): require rollout scheduler sigmas, drop timesteps-derived fallbacks #92) — present on both mains.

Validation

Paired-dump alignment (LTX-2.3, 4×H200, latent bit-exact pairing, 2580 tensor pairs):

metric	before	after
AdaLN σ input	rel 1.8e-3	bit-exact
bit-exact pairs	98/2580	675/2580
pairs ≥5e-2	379 (peak 1.5e-1)	75 (deep-block attn2)
train/log_prob_mean_abs_diff	1.3e-6	4.2e-7

With attention kernels pinned to the same flash path on both sides (--fsdp-attention-backend sdpa_flash + the rollout patches here) the forward reaches 2580/2580 bit-exact pairs — full train↔rollout forward parity. The production recipe keeps sdpa_math (deterministic backward), whose residual is pure MATH-vs-flash kernel noise. Two rollout parity patches (ltx patch group): drop all-true attention masks (a kept trivial mask disqualifies flash and lands on cuDNN — torch's Hopper default priority is cuDNN-first) and run the output LayerNorm tail in fp32 with a single rounding at proj_out, matching autocast's rounding points.

Train metrics reproduce bitwise across two machines and sglang builds (model_output_mean_abs_diff=4.362646e-03, log_prob_mean_abs_diff=4.172325e-07). The re-recorded standard shows the same shift in CI: step-0 rollout reward identical, log_prob_mean_abs_diff 1.0e-6 → 3.6e-7; grad_norm moves with the fp32 optimizer path.

Files

• miles/backends/fsdp_utils/loss_hub/types.py — PreparedBatch.timesteps/sigmas, one domain each
• miles/backends/fsdp_utils/loss_hub/flow_grpo.py — _sigmas_for_timesteps exact-match lookup; branch removed
• miles/backends/fsdp_utils/loss_hub/nft.py — same rename; field-domain mixup fixed (timesteps carried σ)
• miles/backends/fsdp_utils/configs/train_pipeline_config.py — compute_noise_pred(..., sigmas_input); knob deleted
• miles/backends/fsdp_utils/configs/{ltx,qwen_image,sd3,wan2_2}.py — families read their model's domain
• miles/backends/fsdp_utils/actor.py — passes both through the input dtype policy
• miles/backends/fsdp_utils/precision.py — wrap-plan LoRA-prefix fix (refactor(fsdp)(4/5): per-parameter weight-precision control for FSDP2 #91 bug)
• scripts/run-diffusion-grpo-ltx23-sglang.sh — fp32 master weights
• tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json — re-recorded on the h200-5gpu runner
• tests/fast/backends/fsdp_utils/test_flow_grpo_sigma_lookup.py, test_precision_plan.py, test_input_dtype_policy.py, configs/test_train_pipeline_config_registry.py — new/updated CPU tests (stage-a-cpu)

Checklist

• pre-commit run passes on the touched files (black/ruff/isort/autoflake, dev image)
• Added/updated tests for new behaviour
• pytest tests/fast/backends/fsdp_utils tests/fast/utils -q green (121 passed; precision tests re-run after the final rebase)
• If launch flags changed, python3 train.py --help still parses — no flag changes (recipe value only)
• If a public flag was added, it appears in the CLI reference docs — no new flags
• If an example was added, it has a real walkthrough — no new examples