feat(training): default fp32 master weights for full-param torch.optim training

[yuhezhang-ai]

Summary

Rebased onto main now that the dtype bug fixes have merged (#2419). This PR keeps only the default-dtype behavior change: full-parameter training with a torch.optim optimizer (e.g. torch.optim.Adam/AdamW) now defaults model.torch_dtype to float32, so the resident parameters act as the fp32 master copy and the Adam state stays full precision, while forward/backward stay bf16 through FSDP2 mixed precision. Previously, leaving model.torch_dtype unset silently gave bf16 master weights + bf16 Adam state (lower-precision training than DeepSpeed/Accelerate/torchtitan defaults).

The resolver itself (resolve_storage_dtype), the HF load-dtype promotion, the per-param compute-dtype FSDP grouping, the pipeline-dtype default, and the frozen-module cast all landed in #2419 (where resolve_storage_dtype was merged but left unwired). This PR turns the default on and adds the matching config/doc changes.

What changed

Wire the fp32-master default into the recipes. Calls resolve_storage_dtype() before model construction in the LLM SFT (train_ft), VLM (finetune), sequence-classification (train_seq_cls), and bi-encoder retrieval (train_bi_encoder) recipes. No-op for PEFT/LoRA, TE FusedAdam (keeps its own fp32 master), and any explicit model.torch_dtype.

EAGLE draft fp32 master. The DDP-trained EAGLE-1/EAGLE-3 draft is stored in fp32 (storage_dtype) and runs its forward under a bf16 autocast context (_compute_autocast), matching the reference EAGLE setup (accelerate --mixed_precision=bf16). Autocast is applied in the train/val loops, the shared _forward_batch helper, and the P-EAGLE partitioned segment step; backward() stays inside the DDP no_sync context. CPU stays fully fp32.

Diffusion configs. Full-param diffusion finetune configs (flux / qwen_image / hunyuan / wan2.1 t2v multinode) pinned to torch_dtype: float32 + compute_dtype: bfloat16, matching wan2_1_t2v_flow. Benchmark/pretrain configs need no edits — they already resolve to fp32 master via the new default.

Docs. Mixed-precision and finetune guides updated to document the new default (Default Behavior section), restore the TE FusedAdam bf16 Adam-state mode (fp32 master + bf16 EMA, as used by DeepSeek-V3), reflect that bf16 storage is now opt-in (Risky Pattern), and explain how to opt out.

Validation notes

• resolve_storage_dtype is unit-tested in fix(precision): dtype contract bug fixes for FSDP2 mixed-dtype loads #2419; recipe wiring follows the existing convention (no dedicated setup test).
• Memory impact: fp32 master weights increase optimizer memory; several example configs may now OOM at their current batch size / GPU count and will need tuning. GPU smoke tests (especially the EAGLE recipes, which can't be sandbox-validated) are planned before marking ready.

Test plan

• GPU smoke test: LLM full-param SFT (torch AdamW) — confirm fp32 master, bf16 compute
• GPU smoke test: EAGLE-1 and EAGLE-3 draft training (incl. P-EAGLE partitioned path)
• GPU smoke test: a hybrid model (NemotronH / Qwen3.5) — confirm fp32-pinned params compute in fp32
• Re-tune example configs that OOM under fp32 master
• Confirm PEFT / TE FusedAdam paths are untouched

[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.