fix(fp16): filter requires_grad in FP16 optimizer flat buffer init (#8029)
## Description
`FP16_Optimizer` (fused) and `FP16_UnfusedOptimizer` allocate bf16 and
fp32 flat buffers for **every** parameter in a group, regardless of
`requires_grad`. With LoRA/PEFT configurations, where ~99.9% of
parameters are frozen, both buffers are allocated for billions of frozen
weights that will never receive a gradient update — causing CUDA OOM
before training begins.
This is a regression introduced in PR #7515 ("Enable non-ZeRO mode"),
which routed `bf16 + ZeRO stage=0` to `FP16_Optimizer` instead of
`BF16_Optimizer`. `BF16_Optimizer` has always filtered to trainable
parameters when building its flat buffers
([`bf16_optimizer.py#L147-148`](https://github.com/deepspeedai/DeepSpeed/blob/510ebe58e4e4148bbf49350ad18fe291a0afb381/deepspeed/runtime/bf16_optimizer.py#L147-L148));
that filter was not carried over:
```python
trainable_parameters = [param for param in param_group['params'] if param.requires_grad]
self.bf16_groups.append(trainable_parameters)
```
## Root cause
`fused_optimizer.py` and `unfused_optimizer.py` both do:
```python
self.fp16_groups.append(param_group['params']) # all params, frozen included
fp32_group = [p.clone().float().detach() for p in param_group['params']] # clones all
```
([`fused_optimizer.py#L89`](https://github.com/deepspeedai/DeepSpeed/blob/510ebe58e4e4148bbf49350ad18fe291a0afb381/deepspeed/runtime/fp16/fused_optimizer.py#L89),
[`unfused_optimizer.py#L72,L76`](https://github.com/deepspeedai/DeepSpeed/blob/510ebe58e4e4148bbf49350ad18fe291a0afb381/deepspeed/runtime/fp16/unfused_optimizer.py#L72-L76)
at `510ebe58`)
## Fix
Apply the same `requires_grad` filter that `BF16_Optimizer` already
uses:
```diff
- self.fp16_groups.append(param_group['params'])
+ trainable = [p for p in param_group['params'] if p.requires_grad]
+ self.fp16_groups.append(trainable)
```
([`fused_optimizer.py#L89-90`](https://github.com/avicooper1/DeepSpeed/blob/2fa169746b140d575f7463e71a3e1e890a34eecd/deepspeed/runtime/fp16/fused_optimizer.py#L89-L90),
[`unfused_optimizer.py#L72-73`](https://github.com/avicooper1/DeepSpeed/blob/2fa169746b140d575f7463e71a3e1e890a34eecd/deepspeed/runtime/fp16/unfused_optimizer.py#L72-L73))
Frozen parameters are never updated and never accumulate gradients, so
excluding them from the flat buffers is both correct and safe.
`allreduce_gradients()` iterates `module.named_parameters()` directly
and already skips frozen params at
[`engine.py#L3084-3085`](https://github.com/deepspeedai/DeepSpeed/blob/510ebe58e4e4148bbf49350ad18fe291a0afb381/deepspeed/runtime/engine.py#L3084-L3085),
so this change has no correctness impact on gradient synchronization.
## Memory impact (Qwen3-8B + LoRA r=16, A100-40GB)
| | GPU allocated |
|---|---|
| Without fix | 37.90 GiB → **CUDA OOM** |
| With fix | **15.36 GiB** → training proceeds normally |
## Testing
Tested on A100-40GB with Qwen3-8B + LoRA (trainable: 7.67M / 8.20B
params, 0.0935%):
```python
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", torch_dtype=torch.bfloat16)
model = get_peft_model(model, LoraConfig(r=16, target_modules=["q_proj", "v_proj"]))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
ds_model, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config={
"bf16": {"enabled": True}, "zero_optimization": {"stage": 0},
"train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 1
})
# SUCCESS: 15.36 GiB allocated, fp16_groups contains only 144 trainable params
```
## Notes on empty trainable groups
Codex flagged that filtering to `requires_grad` can produce an empty
`trainable` list when an entire param group is frozen, and that
`_flatten_dense_tensors([])` would crash in that case.
This gap exists identically in `BF16_Optimizer`
([`bf16_optimizer.py#L147-152`](https://github.com/deepspeedai/DeepSpeed/blob/510ebe58e4e4148bbf49350ad18fe291a0afb381/deepspeed/runtime/bf16_optimizer.py#L147-L152)),
which applies the same filter with no empty-group guard before
flattening. Since this PR is porting that behavior faithfully, the gap
is left as-is to stay consistent. If maintainers prefer, we can add an
empty-group guard in both optimizers together.
## Related
- PR #7515 — introduced `FP16_Optimizer` for `bf16 + stage=0` (source of
the regression)
- PR #7839 — fixed loss scaling and `zero_grad` bugs in the same code
path
Signed-off-by: Avi Cooper <avicooper007@gmail.com>
DeepSpeed
01e3a28b - fix(fp16): filter requires_grad in FP16 optimizer flat buffer init (#8029)
