Refactor consolidate transpose (#7934)

refactor(module_inject): consolidate duplicate transpose functions

- Extract the duplicated `transpose` function into
`deepspeed/module_inject/utils.py`.
- Remove redundant `transpose` definitions from `policy.py` and
`load_checkpoint.py`.
- This resolves an existing `TODO (lekurile)` to consolidate the
function across containers.

---------

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: nathon-lee <248585198+nathon-lee@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>

Author: 3 people

deepspeed/module_inject/load_checkpoint.py

@@ -17,6 +17,7 @@
 import gc
 from deepspeed.accelerator import get_accelerator
 import re
+from .utils import transpose
 
 
 def load_model_with_checkpoint(r_module,
@@ -42,14 +43,6 @@ def prefix_check():
 
     skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix
 
-    def transpose(data):
-        with torch.no_grad():
-            data = data.contiguous()
-            data1 = data.transpose(-1, -2).reshape(-1)
-            data.reshape(-1).copy_(data1)
-            data1 = None
-        return data.reshape(data.shape[-1], data.shape[-2])
-
     def load(module, prefix):
         args = (sd[0], prefix, {}, True, [], [], error_msgs)
 

deepspeed/module_inject/policy.py

@@ -7,6 +7,7 @@
 from deepspeed.utils.types import ActivationFuncType, NormType
 import torch
 from deepspeed.accelerator import get_accelerator
+from .utils import transpose
 
 transformer_param_names = (
         'attn_qkvw', \
@@ -109,16 +110,6 @@ def layernorm(self):
         raise NotImplementedError
 
 
-# TODO (lekurile): This function exists in base container as well, consolidate as some point
-def transpose(data):
-    with torch.no_grad():
-        data = data.contiguous()
-        data1 = data.transpose(-1, -2).reshape(-1)
-        data.reshape(-1).copy_(data1)
-        data1 = None
-    return data.reshape(data.shape[-1], data.shape[-2])
-
-
 # TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point
 def _transpose(x, heads=1, mp_replace=None):
     heads = heads // mp_replace.mp_size  # type: ignore

deepspeed/module_inject/utils.py

@@ -3,9 +3,19 @@
 
 # DeepSpeed Team
 
+import torch
 from deepspeed.utils import log_dist
 
 
+def transpose(data):
+    with torch.no_grad():
+        data = data.contiguous()
+        data1 = data.transpose(-1, -2).reshape(-1)
+        data.reshape(-1).copy_(data1)
+        data1 = None
+    return data.reshape(data.shape[-1], data.shape[-2])
+
+
 # helper function to map between DS policies and DS containers
 def policy_to_ds_container(**kwargs):
     from .containers import HFGPT2LayerPolicy, DS_GPT2Container