From c38d6d9d1404dcb9337db7ef6fff37bf4fd652e9 Mon Sep 17 00:00:00 2001 From: Sung Hyun Cho Date: Wed, 5 Aug 2026 22:15:32 +0900 Subject: [PATCH] Fix DeepCompile ZeRO-3 KeyError for models with frozen parameters init_z3 looks up every module parameter in the stage-3 optimizer's __param_id_to_grad_partition map, but that map is built only from the trainable parameter groups: _get_trainable_parameter_groups filters out parameters with requires_grad=False. With ZeRO-3 and "compile": {"deepcompile": true}, engine.compile() therefore raised KeyError on the first frozen parameter, which breaks any LoRA/PEFT or partial-freeze run. Register frozen parameters with an empty grad buffer instead, mirroring the requires_grad guard already used by set_grad_buffer in the same file. The empty buffer is never consumed: add_gather_and_reduce schedules no reduce op for a parameter without a grad node, and the buffer is only read from flushReduceBucket via those reduce ops. Frozen parameters are still registered with the native handle because they are partitioned and need gather/release ops in the forward graph. Add TestDeepCompile::test_frozen_params using SimpleFrozenModel. It compares loss and parameters against a ZeRO-0 eager baseline across the warmup boundary, so the prefetch and selective-gather passes also run with frozen parameters present. compare_loss takes an optional model_cls argument, leaving the existing callers unchanged. Signed-off-by: Sung Hyun Cho --- deepspeed/compile/init_z3.py | 5 +++- tests/unit/v1/compile/test_compile_zero.py | 32 +++++++++++++++++++++- tests/unit/v1/compile/util.py | 7 +++-- 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/deepspeed/compile/init_z3.py b/deepspeed/compile/init_z3.py index cf114879a90b..afe1c71c9c0c 100644 --- a/deepspeed/compile/init_z3.py +++ b/deepspeed/compile/init_z3.py @@ -127,7 +127,10 @@ def init_z3(engine, backend, compile_config, compile_kwargs, schedule=None): for p in engine.module.parameters(): grad_buffer = torch.Tensor() - if use_opt: + # Frozen params (e.g. the base weights of a LoRA setup) are absent from the optimizer's + # grad partition map, which is built from the trainable groups only. They keep the empty + # buffer: no reduce op is scheduled for a param without a grad node, so it is never read. + if use_opt and p.requires_grad: grad_buffer = optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[p.ds_id] # Disable persistent param diff --git a/tests/unit/v1/compile/test_compile_zero.py b/tests/unit/v1/compile/test_compile_zero.py index bd527833775b..bf8e425cc222 100644 --- a/tests/unit/v1/compile/test_compile_zero.py +++ b/tests/unit/v1/compile/test_compile_zero.py @@ -12,7 +12,7 @@ from unit.v1.compile.util import compare_loss from unit.common import DistributedTest -from unit.simple_model import SimpleModel +from unit.simple_model import SimpleModel, SimpleFrozenModel from unit.util import bf16_required_version_check, skip_on_arch import deepspeed from deepspeed.ops.aio import AsyncIOBuilder @@ -265,3 +265,33 @@ def test_fusing_allgather_and_autocast(self, zero_stage, dtype): } compare_loss(self, config_dict, torch.float32) + + @pytest.mark.parametrize('dtype', [torch.float32]) + @pytest.mark.parametrize('zero_stage', [3]) + def test_frozen_params(self, zero_stage, dtype): + """Test that models with frozen params (e.g. LoRA/PEFT) work correctly with DeepCompile""" + if not required_torch_version(min_version=2.6): + pytest.skip("DeepCompile requires PyTorch >= v2.6") + + if get_accelerator().device_name() == "cpu": + pytest.skip("CPU does not support this test yet") + + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "steps_per_print": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 0.00015 + } + }, + "zero_optimization": { + "stage": zero_stage, + }, + "compile": { + "deepcompile": True + } + } + + # Need warmup steps so that the profiling passes also run with frozen params present + compare_loss(self, config_dict, dtype, iteration=10, model_cls=SimpleFrozenModel) diff --git a/tests/unit/v1/compile/util.py b/tests/unit/v1/compile/util.py index ce5d2518049a..c61554091a1c 100644 --- a/tests/unit/v1/compile/util.py +++ b/tests/unit/v1/compile/util.py @@ -19,8 +19,9 @@ from unit.common import allclose_on_all_ranks -def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None, rtol=None, atol=None): +def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None, rtol=None, atol=None, model_cls=None): hidden_dim = hidden_dim_override if hidden_dim_override is not None else 10 + model_cls = SimpleModel if model_cls is None else model_cls # the default tolerances are too small for the ZeRO-0 eager vs ZeRO-3 compiled comparison RTOL = 5e-1 if rtol is None else rtol @@ -40,7 +41,7 @@ def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None, rto get_accelerator().manual_seed_all(seed) device = torch.device(get_accelerator().current_device_name()) - model = SimpleModel(hidden_dim) + model = model_cls(hidden_dim) i = get_accelerator().current_device() baseline_model = deepcopy(model) @@ -53,7 +54,7 @@ def compare_loss(self, config, dtype, iteration=5, hidden_dim_override=None, rto if config["zero_optimization"]["stage"] == 3: with deepspeed.zero.Init(config_dict_or_path=config): - target_model = SimpleModel(hidden_dim) + target_model = model_cls(hidden_dim) with GatheredParameters(target_model.parameters(), modifier_rank=0): for p1, p2 in zip(target_model.parameters(), model.parameters()): p1.data.copy_(p2.data)