Fix ZeRO parameter alignment for grouped_mm

Author: fwerkor

deepspeed/runtime/zero/stage_1_and_2.py

@@ -795,7 +795,14 @@ def _configure_moe_settings(self):
     def _update_model_bit16_weights(self, group_index):
         updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index])
         for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params):
-            p.data = q.data
+            # Some kernels, including torch grouped_mm, require 16-byte aligned parameter pointers.
+            if q.data_ptr() % 16 == 0:
+                p.data = q.data
+            elif (p.data.shape == q.shape and p.data.dtype == q.dtype and p.data.device == q.device
+                  and p.data.data_ptr() % 16 == 0):
+                p.data.copy_(q.data)
+            else:
+                p.data = q.data.clone()
 
         # set model fp16 weight to slices of reordered flattened buffer
         for param_index, param in enumerate(self.bit16_groups[group_index]):

tests/unit/v1/zero/test_stage2_flatten_on_gpu.py

@@ -18,6 +18,17 @@
 _DTYPE_MAP = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
 
 
+class _MisalignedParamModel(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.offset = torch.nn.Parameter(torch.ones(1))
+        self.weight = torch.nn.Parameter(torch.ones(8, 8))
+
+    def forward(self, x):
+        return (x @ self.weight).sum() + self.offset.sum()
+
+
 def _apply_dtype_to_config(config_dict, dtype):
     """Set bf16/fp16 in config_dict based on dtype; skip if not supported."""
     if dtype == "bf16":
@@ -31,6 +42,46 @@ def _apply_dtype_to_config(config_dict, dtype):
     # fp32: no half-precision block
 
 
+class TestStage1ParamAlignment(DistributedTest):
+    world_size = 1
+
+    def test_model_params_remain_16_byte_aligned(self):
+        if not get_accelerator().is_bf16_supported():
+            pytest.skip("bf16 is not supported on this accelerator")
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "bf16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": 1
+            },
+        }
+        model = _MisalignedParamModel()
+        optimizer = torch.optim.AdamW(model.parameters(), lr=0.1)
+        engine, _, _, _ = deepspeed.initialize(config=config_dict,
+                                               model=model,
+                                               optimizer=optimizer,
+                                               model_parameters=model.parameters())
+
+        opt = engine.optimizer
+        flat_views = opt.unflatten(opt.bit16_groups_flat[0], opt.round_robin_bit16_meta[0])
+        assert flat_views[1].data_ptr() % 16 != 0
+        assert engine.module.weight.data_ptr() % 16 == 0
+        weight_before_step = engine.module.weight.detach().clone()
+
+        data = torch.ones(1, 8, device=engine.device, dtype=torch.bfloat16)
+        loss = engine(data)
+        engine.backward(loss)
+        engine.step()
+
+        assert engine.module.weight.data_ptr() % 16 == 0
+        assert not torch.equal(engine.module.weight, weight_before_step)
+        flat_views = opt.unflatten(opt.bit16_groups_flat[0], opt.round_robin_bit16_meta[0])
+        assert torch.equal(engine.module.weight, flat_views[1])
+
+
 @pytest.mark.parametrize("zero_stage", [1, 2])
 @pytest.mark.parametrize("dtype", ["fp32", "fp16", "bf16"], ids=["fp32", "fp16", "bf16"])
 class TestStage2FlattenOnGPU(DistributedTest):