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

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

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)

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)