NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Use explicit FP32 storage for full-parameter Adam/AdamW training
- Dominant language
- Python
- Stars
- 963
- Forks
- 318
- Avg merge
- 3d 20h
- Merged PRs (30d)
- 143
Description
Migrate applicable full-parameter training examples to explicit FP32 model storage with torch.optim.Adam / AdamW, while retaining BF16 compute where supported.
Decision and rationale
For these optimizers, resident model parameters serve as master weights and Adam moment buffers follow their dtype. BF16 storage therefore gives BF16 updates and moments; an FSDP mixed-precision compute policy does not provide a separate FP32 master copy.
Select storage precision explicitly in each YAML instead of inferring it from the optimizer. This replaces the original proposal to enable resolve_storage_dtype: explicit settings keep model construction independent of optimizer configuration, expose the memory tradeoff, and avoid silently changing existing workloads. PEFT and TE FusedAdam need separate precision choices; the BF16 Adam/AdamW warning remains useful.
Earlier mixed-dtype fixes landed in NVIDIA-NeMo/Automodel#2419. This issue tracks the remaining config migration and its validation.
First: loader compatibility — PR #2379
NVIDIA-NeMo/Automodel#2379 is a draft prerequisite for using canonical model.dtype in the planned migration. Transformers introduced dtype as the preferred spelling, but NeMo's wrapper did not handle it consistently: a Python or YAML request for dtype="float32" could still construct BF16 parameters through both custom and HF loading paths.
The PR normalizes the argument before config overrides consume it, supports both dtype and legacy torch_dtype, and preserves the requested precision through initialization, loading, and retries. It also removes the optimizer-driven resolver and updates precision documentation and authoring skills while retaining the warning.
Existing model.torch_dtype: float32 remains supported; the prerequisite specifically makes the newer spelling reliable. No example YAMLs change in this PR. Local validation passed 331 CPU tests, including 77 dtype cases. The loader tests bypass CUDA/distributed infrastructure, so this is not training-config validation.
Follow-up: migrate and validate configs
Set model.dtype: float32 explicitly for applicable NeMoAutoModel full-parameter Adam/AdamW examples. Preserve the intended compute policy. Review PEFT and TE separately; diffusion uses its own model.torch_dtype / model.compute_dtype interface.
FP32 parameters, gradients, and optimizer state can increase memory enough to require batch-size or parallelism adjustments. Migrate in manageable batches and validate actual storage/state/compute dtypes, peak memory, loss and gradient behavior, and checkpoint/resume where applicable. Use targeted GPU runs and scoped CI for changed configurations; broaden coverage when warranted rather than requiring a release-wide sweep for every batch.
See the mixed-precision guide for the precision patterns.
Completion criteria
- Land PR #2379's dtype compatibility fix and resolver removal.
- Audit and update applicable full-parameter Adam/AdamW configs with explicit storage precision.
- Record validation, resource adjustments, and justified exceptions for migrated configs.
This issue remains open after PR #2379 merges, until config migration and validation are complete.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.