[DeepCompile] Fix KeyError on frozen parameters in ZeRO-3

[roycho96]

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

# 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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c38d6d9d14

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a non-merge commit but its message has no Signed-off-by: trailer, which violates the repository's Commit & CI requirement and can cause DCO/signoff validation to reject the change. Please recommit/amend with --signoff using the configured Git identity before merging.

AGENTS.md reference: AGENTS.md:L6-L8

Useful? React with 👍 / 👎.