NVIDIA-NeMo / NVIDIA-NeMo/RL

[Automodel]: Hard-coded fp32 model weights doubles memory with FusedAdam

Open
#2,865 3 comments 0 reactions 0 assignees View on GitHub
bug Memory
Dominant language
Python
Stars
2k
Forks
562
Avg merge
4d 1h
Merged PRs (30d)
150

Description

## Summary

The automodel backend hard-codes the model load dtype to **fp32** in order to support master weights [here](https://github.com/NVIDIA-NeMo/RL/blob/1ff6e117a41e37fdd80401ae375ec9f23e8d28ea/nemo_rl/models/automodel/setup.py#L305).
However, the optimizer is configurable and supports both `torch.optim.AdamW` and TE `FusedAdam`. When **FusedAdam** is used, the fp32 load is unnecessary and causes **~2× optimizer master-weight memory** (and ~2× model-weight memory), because FusedAdam already maintains its own fp32 master internally.

## Details

With the TE FusedAdam optimizer, the fp32 master is kept **inside the optimizer**, via:

```yaml
optimizer:
name: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
kwargs:
master_weights: true
store_param_remainders: true
```

`store_param_remainders: true` lets FusedAdam split the fp32 master into **bf16 (high 16 bits) + int16 remainder (low 16 bits)** — `16 + 16 = 32 = fp32`. This depends on the parameters being **bf16**: once a parameter is fp32, TE detects `param.dtype != bf16` and **skips the remainder path, storing a separate full fp32 master** instead. Ref: [https://github.com/NVIDIA/TransformerEngine/blob/42b840051647eef89761a16dfdff87e82bb253ab/transformer_engine/pytorch/optimizers/fused_adam.py#L97](https://github.com/NVIDIA/TransformerEngine/blob/42b840051647eef89761a16dfdff87e82bb253ab/transformer_engine/pytorch/optimizers/fused_adam.py#L97)

So when the load dtype is hard-coded to fp32, with FusedAdam:

- **Model weights**: bf16 → fp32 = **2× memory**.
- **Optimizer master**: `bf16 + int16 remainder` → **full fp32** = **2× memory**.

## Reproduction

Run a grpo training, toggling only the model load dtype (`AutoConfig(torch_dtype=...)` in `nemo_rl/models/automodel/setup.py`) between
`float32` and `bfloat16`, and sum the `optimizer.state` tensor bytes after the first `optimizer.step()`.

## Observed (rank 0)

| metric | fp32 load (current) | modified to bf16 load |
| ------------------------- | -------------------------------- | -------------------- |
| `master_param` | **fp32 7.32 GB** + int16 0.60 GB | int16 — 4.26 GB |
| optimizer state /rank | **16.43 GB** | 12.77 GB |
| optimizer state all-ranks | 131.5 GB | 102.2 GB |
| model params all-ranks | 66.6 GB | 35.7 GB |

`exp_avg` / `exp_avg_sq` are unchanged (bf16); the entire increase is the master copy.

## Suggested fix

Choose the model load dtype based on the optimizer (AdamW vs FusedAdam) instead of hard-coding fp32: load fp32 only for optimizers without an internal master, and keep the compute dtype (bf16) when FusedAdam holds the master.

For reference, I have a candidate fix in: [https://github.com/NVIDIA-NeMo/RL/commit/e968d90f72925637e8834f81510b5f7a181b4093](https://github.com/NVIDIA-NeMo/RL/commit/e968d90f72925637e8834f81510b5f7a181b4093)

Contributor guide

Open the contributing guide

Research direction

Start in nemo_rl/models/automodel/setup.py at the AutoConfig(torch_dtype=...) model-loading logic, then inspect how the configured AdamW and Transformer Engine FusedAdam optimizers are selected. Reproduce the issue by comparing optimizer.state tensor bytes after the first optimizer.step() with float32 and bfloat16 loads; done means the dtype follows the optimizer and avoids the unnecessary memory increase with FusedAdam.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
66/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.