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)