NVIDIA / NVIDIA/Megatron-LM

[BUG] Distributed optimizer materializes HybridDeviceOptimizer twice

Open
#7,232 0 comments 0 reactions 0 assignees View on GitHub
community-request waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Describe the bug**

Tagging @NVIDIA/mcore-oncall.

Enabling `optimizer_cpu_offload` together with `use_distributed_optimizer` materializes
`HybridDeviceOptimizer` twice.

The optimizer factory first constructs HDO from the unsharded parameter groups. HDO's
constructor immediately calls `_init_sub_optimizers()`, which copies offloaded parameters
to CPU and may create FP32 master parameters. `DistributedOptimizer` later builds its
parameter shards, discards that HDO, and constructs a second HDO from the sharded groups.
The first set of allocations is never used.

The current flow on main is visible at:

- initial HDO construction:
https://github.com/NVIDIA/Megatron-LM/blob/5bbf72793092223400b684ed95fbd85fa1d48d4c/megatron/core/optimizer/__init__.py#L540-L550
- eager sub-optimizer initialization:
https://github.com/NVIDIA/Megatron-LM/blob/5bbf72793092223400b684ed95fbd85fa1d48d4c/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py#L72-L81
- replacement after shard construction:
https://github.com/NVIDIA/Megatron-LM/blob/5bbf72793092223400b684ed95fbd85fa1d48d4c/megatron/core/optimizer/distrib_optimizer.py#L785-L795

For large models this substantially increases optimizer startup time and transient host
memory. At high offload fractions, the throwaway instance clones most selected BF16/FP16
parameters to CPU and then creates FP32 master copies before the distributed optimizer
has narrowed them to the rank-local shards.

This is separate from #6212/#6982: those fixed non-leaf FP32 views passed into the HDO
that is reconstructed after sharding, while this report concerns redundant eager
materialization of the earlier HDO before distributed sharding.

It is also independent of #5083 (precision-aware FP8 placeholders) and #5071
(refreshing HDO's internal copies after checkpoint-only model reloads). None of those
changes removes the first, unsharded HDO materialization.

**Steps/Code to reproduce bug**

Construct a Megatron optimizer with both options enabled and count calls to
`HybridDeviceOptimizer._init_sub_optimizers`:

```python
optimizer_config = OptimizerConfig(
optimizer="adam",
lr=1e-3,
bf16=True,
use_distributed_optimizer=True,
optimizer_cpu_offload=True,
optimizer_offload_fraction=1.0,
)
optimizer = get_megatron_optimizer(optimizer_config, [ddp_model])
```

On current main, `_init_sub_optimizers` is called twice for one resulting optimizer:
once for the full parameter groups and once for the sharded replacement. The existing
`test_distributed_optimizer_with_cpu_offload_and_fp32_marked_param` is a small reproducer
when the method is wrapped with a call counter.

**Expected behavior**

Only the final HDO built from distributed optimizer shards should materialize CPU/GPU
sub-optimizers and parameter copies. Ordinary, non-distributed HDO construction should
remain eager.

**Additional context**

A narrow fix is available that defers sub-optimizer initialization only for the
throwaway HDO. The deferral flag is not copied into optimizer `defaults`, so the sharded
replacement initializes normally. A one-shot guard also preserves safe behavior if a
deferred instance is used before replacement.

Contributor guide

Open the contributing guide

Research direction

Start at the optimizer factory in megatron/core/optimizer/__init__.py, then trace HybridDeviceOptimizer initialization in megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py and replacement in megatron/core/optimizer/distrib_optimizer.py. Use test_distributed_optimizer_with_cpu_offload_and_fp32_marked_param as the reproducer, counting _init_sub_optimizers calls; done means only the final sharded HDO materializes sub-optimizers while ordinary HDO construction remains eager.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
76/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.