Fix ZeRO parameter alignment for grouped_mm (#8277)

## Summary

ZeRO stage 1/2 stores model parameters as views into flattened fp16/bf16
buffers. Individual parameter views can start at non-16-byte offsets
even when the flat buffer itself is aligned, which breaks
alignment-sensitive kernels such as `torch._grouped_mm`.

Pad parameter boundaries inside the ZeRO flat buffer so model parameter
views remain 16-byte aligned without duplicating misaligned parameters.
The padded layout is propagated through partition/gradient offsets,
LP↔HP linkage, DeepCompile gradient buffers, checkpoint restore, and
`zero_to_fp32` reconstruction.

Older checkpoints without parameter-alignment padding remain loadable;
their compact layout is converted when restored.

## Tests

- ZeRO-1/2 BF16 regression with a deliberately misaligned parameter
layout.
- Verifies zero-copy aligned flat-buffer views across optimizer steps.
- Verifies checkpoint reload with `load_module_only=True` and
`load_optimizer_states=False`.
- Synced with current `master`, including the ZeRO-1/2 DeepCompile file
rename.

Fixes #8276

---------

Signed-off-by: Cao Yuhang <caoyuhang@fwerkor.com>
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Co-authored-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

Author: fwerkor

deepspeed/checkpoint/constants.py

@@ -18,6 +18,7 @@
 SINGLE_PARTITION_OF_FP32_GROUPS = "single_partition_of_fp32_groups"
 PARAM_GROUPS = 'param_groups'
 GROUP_PADDINGS = 'group_paddings'
+PARAM_ALIGNMENT_PADDINGS = 'param_alignment_paddings'
 PARTITION_COUNT = 'partition_count'
 ZERO_STAGE = 'zero_stage'
 CLIP_GRAD = 'clip_grad'

deepspeed/compile/init_z1_and_2.py

@@ -8,6 +8,7 @@
 
 import torch
 
+import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
 from .passes import zero_1_and_2_compile, zero3_compile
 from .backend import make_backend, launch_compile_passes, init_schedule
@@ -51,39 +52,47 @@ def _build_partition_grad_views(optimizer, group_idx):
             optimizer.all_grad_tensors[group_idx] = original_all_grad_tensors
 
 
+def _clone_partition_grad_views(partition_grad_views):
+    cloned_views = []
+    for view in partition_grad_views:
+        cloned = view.clone().detach()
+        if getattr(view, "_zero_padding", False):
+            cloned._zero_padding = True
+        cloned_views.append(cloned)
+    return cloned_views
+
+
 def _build_flat_partition_grad_views(optimizer, group_idx):
     partition_size = int(optimizer.partition_size[group_idx])
     dtype = optimizer.gradient_accumulation_dtype
     device = get_accelerator().current_device_name()
     flat_buffer = torch.zeros(partition_size, dtype=dtype, device=device)
+    partition_id = dist.get_rank(group=optimizer.real_dp_process_group[group_idx])
 
     views = []
     current_size = 0
-    for i, tensor in enumerate(optimizer.params_in_partition[group_idx]):
-        num_elements = tensor.numel()
-        tensor_offset = 0
-
-        if i == 0 and optimizer.first_offset[group_idx] > 0:
-            tensor_offset = int(optimizer.first_offset[group_idx])
-            num_elements -= tensor_offset
-
-        if num_elements > partition_size - current_size:
-            num_elements = partition_size - current_size
-
-        if num_elements <= 0:
-            continue
-
-        view = flat_buffer.narrow(0, current_size, int(num_elements))
-        if tensor_offset == 0 and num_elements == tensor.numel():
-            view = view.view(tensor.shape)
-        views.append(view)
-        current_size += int(num_elements)
-
-        if current_size >= partition_size:
-            break
+    for tensor in optimizer.params_in_partition[group_idx]:
+        param_id = optimizer.get_param_id(tensor)
+        source_offset = optimizer.grad_start_offset[group_idx][partition_id][param_id]
+        dest_offset = optimizer.grad_partition_insertion_offset[group_idx][partition_id][param_id]
+        num_elements = min(tensor.numel() - source_offset, partition_size - dest_offset)
+
+        if dest_offset > current_size:
+            padding = flat_buffer.narrow(0, current_size, dest_offset - current_size)
+            padding._zero_padding = True
+            views.append(padding)
+
+        if num_elements > 0:
+            view = flat_buffer.narrow(0, dest_offset, int(num_elements))
+            if source_offset == 0 and num_elements == tensor.numel():
+                view = view.view(tensor.shape)
+            views.append(view)
+            current_size = dest_offset + int(num_elements)
 
     if current_size < partition_size:
-        views.append(flat_buffer.narrow(0, current_size, partition_size - current_size))
+        padding = flat_buffer.narrow(0, current_size, partition_size - current_size)
+        padding._zero_padding = True
+        views.append(padding)
 
     return flat_buffer, views
 
@@ -101,7 +110,12 @@ def init_z1_and_2(engine, backend, compile_config, compile_kwargs, schedule=None
     if use_z2:
         grad_buffer = {}
         for i, group in enumerate(optimizer.bit16_groups):
-            grad_buffer[i] = [p.clone().detach() for p in _build_partition_grad_views(optimizer, i)]
+            partition_grad_views = _build_partition_grad_views(optimizer, i)
+            grad_buffer[i] = _clone_partition_grad_views(partition_grad_views)
+            param_grad_buffers = [
+                cloned for original, cloned in zip(partition_grad_views, grad_buffer[i])
+                if not getattr(original, "_zero_padding", False)
+            ]
 
             index_in_partition = 0
             first_in_partition = True
@@ -111,7 +125,7 @@ def init_z1_and_2(engine, backend, compile_config, compile_kwargs, schedule=None
                 in_partition = optimizer.is_param_in_current_partition[param_id]
 
                 if in_partition:
-                    buf = grad_buffer[i][index_in_partition]
+                    buf = param_grad_buffers[index_in_partition]
                     offset = optimizer.first_offset[i] if first_in_partition else 0
                     dc.register_param(p.param_id, p.shape, p, buf, int(offset))
                     index_in_partition += 1
@@ -157,7 +171,8 @@ def set_z1_grad_buffer(is_gradient_accumulation_boundary):
                 flat_grad_buffer, group_grad_buffers = _build_flat_partition_grad_views(optimizer, group_idx)
                 current_grad_buffers[group_idx] = _FlatPartitionGradBufferGroup(
                     group_grad_buffers, flat_grad_buffer, lambda group_idx=group_idx: release_grad_buffer(group_idx))
-                for (param_id, _, offset), grad_buffer in zip(grad_buffer_metadata[group_idx], group_grad_buffers):
+                param_grad_buffers = [g for g in group_grad_buffers if not getattr(g, "_zero_padding", False)]
+                for (param_id, _, offset), grad_buffer in zip(grad_buffer_metadata[group_idx], param_grad_buffers):
                     dc.update_param_grad_buffer(param_id, grad_buffer, offset)
             optimizer.averaged_gradients = current_grad_buffers
 
@@ -179,7 +194,6 @@ def release_grad_buffer(group_idx=None):
             schedule.append((0, [zero_1_and_2_compile.add_z1_reduce]))
     else:
         for opt in schedule:
-            # avoid typical misconfiguration
             if zero3_compile.add_z3_gather_release in opt[1]:
                 raise ValueError("The schedule contains the ZeRO-3 pass add_z3_gather_release, but ZeRO stage 1 "
                                  "or 2 is enabled. Use zero_1_and_2_compile.add_z1_reduce (stage 1) or add_z2_reduce "

deepspeed/runtime/engine.py

@@ -1723,6 +1723,9 @@ def zero_allgather_partitions(self):
     def zero_round_robin_gradients(self):
         return self._config.zero_config.round_robin_gradients
 
+    def zero_parameter_alignment(self):
+        return self._config.zero_config.parameter_alignment
+
     def zero_hpz_partition_size(self):
         return self._config.zero_config.zero_hpz_partition_size
 
@@ -2741,6 +2744,7 @@ def _configure_zero_optimizer(self, optimizer):
                 ignore_unused_parameters=self.zero_ignore_unused_parameters(),
                 partition_grads=zero_stage == ZeroStageEnum.gradients,
                 round_robin_gradients=round_robin_gradients,
+                parameter_alignment=self.zero_parameter_alignment(),
                 has_moe_layers=self.has_moe_layers,
                 fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(),
                 bf16_master_weights_and_gradients=self.bf16_master_weights_and_gradients(),

deepspeed/runtime/zero/config.py

@@ -41,6 +41,7 @@
     "offload_optimizer": {...},
     "ignore_unused_parameters": [true|false],
     "round_robin_gradients": [true|false],
+    "parameter_alignment": [true|false],
     "zero_hpz_partition_size": 1,
     "zero_quantized_weights": [true|false],
     "zero_quantized_nontrainable_weights": [true|false],
@@ -313,6 +314,14 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     Performance benefit grows with gradient accumulation steps (more copying
     between optimizer steps) or GPU count (increased parallelism).
     """
+
+    parameter_alignment: bool = False
+    """
+    Pad ZeRO Stage 1 and 2 flat buffers between parameters so each parameter
+    starts at a 16-byte-aligned address. This is disabled by default because
+    the padding increases flat-buffer and optimizer-state memory usage.
+    """
+
     zero_hpz_partition_size: int = Field(1, ge=0)
     """
     Number of ranks in zero parameters partitioning secondary group