UX feat: default fp32 master weights for full-parameter training with torch.optim

[yuhezhang-ai]

Background

The dtype-handling bug fixes are now merged (#2419 and related). The framework now correctly wires model.torch_dtype through to optimizer state, sharding, and checkpointing. What remains is a UX / default-behavior problem, tracked here separately because it needs broad validation.

Problem

For long full-parameter training (pre-training or extended fine-tuning), the safe pattern is fp32 master weights + fp32 optimizer state. But today, if a user leaves model.torch_dtype unset (or auto/bfloat16) and uses torch.optim.AdamW, the resident params and the AdamW EMA buffers (exp_avg, exp_avg_sq) end up in bf16. This is silently fragile: bf16 EMA quantization can slow convergence (higher final loss) or, in worse cases, cause unstable grad_norm / loss spikes / divergence. Users have to know to set torch_dtype: float32 themselves.

See the mixed-precision guide for the full explanation.

Proposed change

Default model.torch_dtype to float32 for full-parameter training when the optimizer is a torch.optim.* optimizer (TE FusedAdam and PEFT/LoRA are unaffected). The resolver (resolve_storage_dtype) is already implemented but left dormant/unwired from the bug-fix PR — this issue is about turning it on.

Why this is longer-term

Enabling the fp32 default increases resident memory, so a number of existing example configs will likely OOM and need retuning (lower micro-batch, more GPUs, or activation checkpointing). This requires GPU validation across the affected configs before flipping the default.

Scope / tasks

• Wire resolve_storage_dtype into the recipes (train_ft, vlm/finetune, train_seq_cls, retrieval/train_bi_encoder).
• Audit and update example configs that OOM under the new default.
• Ensure benchmark/generation paths opt out (throughput/inference don't want the fp32 default).
• Update docs to describe the new default behavior.

[stanley1208]

Hey @yuhezhang-ai, mind if I grab this one? I saw the resolver is already in from #2419 so I'd mainly be wiring it in and pinning the benchmark configs.
But a quick question here, should the regular examples get pinned too or just let them inherit the default?

[yuhezhang-ai]

Thanks @stanley1208, really appreciate you offering to help here!

I actually already have a draft PR up for this in #2379, but I haven’t had time to fully verify it yet. The code change itself is fairly small now that the resolver is in from #2419, but the important part is validation: changing the default can increase memory usage for existing configs, so we need a full pass to catch and adjust any configs that start to OOM.

Because of that, I think it’s better if I or another internal member drives this one, since we can run the broader internal validation sweep.

Thanks again for jumping in, and I really appreciate the offer!

[stanley1208]

Got it, thank you so much for the reply!