[DeepCompile] Fix KeyError on frozen parameters in ZeRO-3 (#8214)
Under ZeRO-3 with `"compile": {"deepcompile": true}`, `engine.compile()`
fails with a `KeyError` if the model has any parameter with
`requires_grad=False`. LoRA/PEFT and partial-freeze runs cannot use
DeepCompile at all.
`init_z3()` looks up a grad partition for every module parameter, but
the stage-3 optimizer builds that map only for the parameters it owns.
`_get_trainable_parameter_groups()` drops `requires_grad=False` params
(`stage3.py` L651), and the map is filled from the resulting
`fp16_groups` (`stage3.py` L706-720), so the first frozen parameter is
simply not there.
### Repro
```python
# repro.py, run with: torchrun --nproc_per_node=2 repro.py
import torch
import deepspeed
class Net(torch.nn.Module):
def __init__(self, dim=128):
super().__init__()
self.frozen = torch.nn.Linear(dim, dim)
self.trainable = torch.nn.Linear(dim, dim)
self.frozen.requires_grad_(False)
def forward(self, x):
return self.trainable(self.frozen(x)).sum()
config = {
"train_batch_size": 2,
"train_micro_batch_size_per_gpu": 1,
"optimizer": {"type": "Adam", "params": {"lr": 1e-4}},
"zero_optimization": {"stage": 3},
"bf16": {"enabled": True},
"compile": {"deepcompile": True},
}
model = Net()
trainable = [p for p in model.parameters() if p.requires_grad]
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=trainable, config=config)
engine.compile()
x = torch.randn(1, 128, device=engine.device, dtype=torch.bfloat16)
loss = engine(x)
engine.backward(loss)
engine.step()
print("ok")
```
```
File "deepspeed/compile/init_z3.py", line 131, in init_z3
grad_buffer = optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[p.ds_id]
KeyError: 0
```
`ds_id 0` is the frozen layer's weight. We originally hit this through
the HF Trainer path with a LoRA adapter, where the crash happens inside
`accelerator.prepare()` when accelerate calls `engine.compile()`. Also
reproduces on the released 0.19.2 and 0.19.3.
### Fix
Skip the lookup for frozen parameters and leave them with the empty
buffer that the optimizer-less path (`use_opt == False`) already uses
for every parameter.
Nothing reads that buffer for a frozen parameter.
`add_gather_and_reduce()` skips parameters whose grad node is `None`
(`passes/zero3_compile.py` L130), and the registered buffer is only
consumed by `flushReduceBucket` through those reduce ops
(`csrc/compile/z3.cpp` L261, L282, L311). `set_grad_buffer()` a few
lines below already applies the same `requires_grad` guard against the
same map. Frozen parameters are still registered with the native handle,
since they are partitioned and need gather/release ops in the forward
graph.
### Test
`TestDeepCompile::test_frozen_params` runs the existing
`SimpleFrozenModel` for 10 steps on 2 ranks with ZeRO-3 + deepcompile,
comparing loss and parameters against a ZeRO-0 eager baseline. 10 steps
so the run crosses the `WARMUP` boundary and the prefetch and
selective-gather passes also see the frozen parameters. `compare_loss()`
takes an optional `model_cls`, so the existing callers are unchanged.
Without the fix the test fails at `engine.compile()` with the same
`KeyError`.
### Notes
ZeRO-1/2 does not hit this, since `init_z1()` iterates
`optimizer.bit16_groups`, which contains only trainable parameters. I
have not checked whether frozen parameters work end to end there, so
this change is scoped to ZeRO-3.
`grad_partitions.get(p.ds_id, torch.Tensor())` would also cover a
trainable parameter that was never passed to the optimizer, but that
case is already broken in eager ZeRO-3, so I kept the narrower guard to
match `set_grad_buffer()`. Happy to switch if you prefer the wider one.
Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>
DeepSpeed
04b4b0cb - [DeepCompile] Fix KeyError on frozen parameters in ZeRO-3 (#8214)
