Ensure all parameters are aligned

VectorInstitute · Nov 15, 2023 · 25e9cd5 · 25e9cd5
1 parent 901d807
commit 25e9cd5
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/deepspeed/inference/v2/model_implementations/flat_model_helpers.py b/deepspeed/inference/v2/model_implementations/flat_model_helpers.py
@@ -16,6 +16,13 @@
 from ..inference_utils import elem_size
 
 
+def pad_to_aligned_offset(offset: int, alignment: int = 256) -> int:
+    """
+    Pad the provided offset to a well-aligned value.
+    """
+    return ((offset + alignment - 1) // alignment) * alignment
+
+
 class TensorMetadata(DeepSpeedConfigModel):
     """
     A class to represent a tensor specification.
@@ -149,15 +156,15 @@ def process_layer(layer_container: LayerContainer, l_name: str, cur_offset: int)
                                                        strides=param.stride(),
                                                        offset=cur_offset)
 
-            cur_offset += elem_size(param.dtype) * param.numel()
+            cur_offset += pad_to_aligned_offset(elem_size(param.dtype) * param.numel())
 
             for t_name, tensor in param.aux_attrs.items():
                 param_metadata.aux_params[t_name] = TensorMetadata(dtype=str(tensor.dtype),
                                                                    shape=tensor.shape,
                                                                    strides=tensor.stride(),
                                                                    offset=cur_offset)
 
-                cur_offset += elem_size(param.dtype) * param.numel()
+                cur_offset += pad_to_aligned_offset(elem_size(param.dtype) * param.numel())
 
             layer_metadata.params[p_name] = param_metadata