Distributed optimizer + CPU offload (HybridDeviceOptimizer) crashes with FP8 params (fp8_param=True): "optimizer can only optimize Tensors ... params is NoneType"
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
# Distributed optimizer + CPU offload (`HybridDeviceOptimizer`) crashes with FP8 params (`fp8_param=True`): `optimizer can only optimize Tensors, but one of the params is NoneType`
## Describe the bug
When training with **native FP8 parameters** (`--fp8-param-gather` / `fp8_param=True`, i.e. model weights stored as `Float8Tensor`) **together with the precision-aware distributed optimizer CPU offload** (`--optimizer-cpu-offload` + `--use-precision-aware-optimizer`, which uses `HybridDeviceOptimizer`), optimizer construction crashes with:
```
File ".../megatron/core/optimizer/distrib_optimizer.py", line 599, in __init__
self.optimizer = HybridDeviceOptimizer(
File ".../megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py", line 57, in __init__
super(HybridDeviceOptimizer, self).__init__(...)
File ".../torch/optim/optimizer.py", line 1115, in add_param_group
raise TypeError(
TypeError: optimizer can only optimize Tensors, but one of the params is NoneType
```
This happens for a MoE model (Grouped-GEMM experts) on the very first optimizer build, before any training step.
## Root cause (analysis)
The two precision/offload code paths assume incompatible "main-param" ownership models:
1. `OptimizerConfig.use_precision_aware_optimizer_no_fp8_or_ds_fp8` is **forced `True` whenever `optimizer_cpu_offload` is set** (`optimizer_config.py`):
```python
self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = (
self.use_precision_aware_optimizer
and (
self.main_params_dtype != torch.float32
or (self.fp8_recipe is None or self.fp8_recipe == "delayed")
or self.optimizer_cpu_offload # <-- forced True with offload
)
)
```
2. With that flag `True`, `_build_model_and_main_param_groups` (`distrib_optimizer.py`) does **not** build a real FP32 master shard for FP8 params; instead it puts `None` placeholders into the group:
```python
# shard_model_param becomes None for quantized params
if is_float8tensor(model_param) and config.fp8_recipe != "delayed":
shard_model_param = None # grouped/MXFP8/Blockwise can't view(-1)
...
else: # precision-aware branch
shard_main_param = None # "main params are held by FusedAdam"
...
group_range["orig_group"]["params"] = [*shard_fp32_params, *shard_float16_params] # contains None
```
The intent is that the **precision-aware fused optimizer manages the FP8 masters internally** (via the grad/param buffers), and the `None` entries in `param_groups` are just placeholders.
3. The non-offload path tolerates these `None` placeholders because it assigns `param_groups` **directly**:
```python
self.optimizer.param_groups = [g["orig_group"] for g in self.opt_group_ranges]
self.optimizer.load_state_dict(self.optimizer.state_dict())
```
But the **offload path reconstructs** the optimizer through `__init__` → `add_param_group`, which **rejects `None`**:
```python
if isinstance(self.optimizer, HybridDeviceOptimizer):
self.optimizer = HybridDeviceOptimizer(
params=[g["orig_group"] for g in self.opt_group_ranges], **self.optimizer.defaults
)
```
4. `HybridDeviceOptimizer` additionally builds its **own** FP32 master from the passed tensors (`_get_sub_optimizer_param_groups`: `param.detach().clone().float()`), so even if `None` were filtered out it has no way to recover/optimize the FP8 master that the precision-aware path expects to manage internally. The two master-ownership models are mutually exclusive.
### Diagnostic
Instrumenting right before the `HybridDeviceOptimizer` reconstruction (MoE, 256 grouped experts, TP2):
```
use_precision_aware_optimizer_no_fp8_or_ds_fp8 = True
group 0: n_params=520, none_count=518, types={'Tensor': 2, 'None': 518} # FP8 expert shards = None
group 1: n_params=5, types={'Tensor': 5}
```
The 518 `None` entries are exactly the FP8 (Grouped-GEMM) expert parameter shards.
## Variants tried (all crash identically)
| `fp8_recipe` | `fp8_param_gather` | `optimizer_cpu_offload` | result |
|---|---|---|---|
| `tensorwise` | off | **on** | ❌ NoneType |
| `delayed` | off | **on** | ❌ NoneType |
| `delayed` | **on** | **on** | ❌ NoneType |
| `tensorwise` | off | off | ✅ trains (but optimizer states no longer offloaded) |
`first_last_layers_bf16=False` in all FP8 runs. The non-offload run loads a BF16 checkpoint into FP8 params, reshards to the inference engine, runs generation and reaches the optimizer step successfully — so the FP8 model/param path itself is fine; only the **CPU-offload optimizer** is incompatible.
> `main` branch (checked) keeps the same logic — the `None` assignment now explicitly lists *"grouped quantized tensors"* among the cases that can't `view(-1)`, and the `HybridDeviceOptimizer` reconstruction is unchanged — so this is not fixed on `main`.
## Expected behavior
Either:
- **Support** `fp8_param` together with `optimizer_cpu_offload` (have `HybridDeviceOptimizer` cooperate with the precision-aware FP8 master management, e.g. by building real FP32 master shards from the dequantized FP8 params and copying back/requantizing on `step()`), **or**
- raise an **early, explicit error** (`assert`) that the combination is unsupported, instead of a cryptic `NoneType` deep inside `torch.optim`.
## Why it matters
For very large MoE models, FP8 weight storage (halves resident weight memory) and optimizer-state CPU offload (FP32 master + Adam moments don't fit on GPU) are **both required simultaneously**. Today they are mutually exclusive, which blocks memory-efficient FP8 RL/training of large MoE models on limited GPU counts.
## Environment
- Megatron-core: **0.15.0** (same code path confirmed on `main`)
- Transformer Engine: **2.10.0**
- PyTorch: **2.9.0+cu129**, CUDA 12.9
- Model: large MoE, Grouped-GEMM experts (`moe_grouped_gemm=True`), `moe_token_dispatcher_type=alltoall`
- Parallelism in repro: TP2/PP1/EP1 (also reproduces at TP4/PP4/EP8)
## Minimal config to reproduce
```
--use-distributed-optimizer
--use-precision-aware-optimizer
--optimizer-cpu-offload --optimizer-offload-fraction 1.0
--fp8-format e4m3 --fp8-recipe tensorwise # also delayed
--fp8-param-gather # (also reproduces without it, via fp8_param=True)
# MoE with grouped GEMM experts
```
Contributor guide
Assessment
This issue has not been assessed yet.