Missing strided copy for gated MLP (#3788)

Co-authored-by: Ammar Ahmad Awan <[email protected]> Co-authored-by: Jeff Rasley <[email protected]> Co-authored-by: Logan Adams <[email protected]>
deepspeedai · tjruwase · Jul 3, 2023 · Jun 21, 2023 · Jun 21, 2023 · Jun 22, 2023
commit b7e1010b391304617e9f6e45df5d9c636a5d591f
diff --git a/deepspeed/module_inject/containers/features/gated_mlp.py b/deepspeed/module_inject/containers/features/gated_mlp.py
@@ -48,7 +48,14 @@ def mlp_inter_mp(self, mp_replace, reversed_dim=False):
                                       int8=reversed_dim,
                                       allocate_tensor=reversed_dim) if src is not None else None
         else:
-            super().mlp_inter_mp(mp_replace)
+            self.module.mlp.inter_w = mp_replace.strided_copy(self.module.mlp.inter_w,
+                                                              self._h4h_w,
+                                                              num_splits=2,
+                                                              int8=reversed_dim)
+            self.module.mlp.inter_b = mp_replace.strided_copy(self.module.mlp.inter_b,
+                                                              self._h4h_b,
+                                                              num_splits=2,
+                                                              int8=reversed_dim)
 
     def release_mlp(self):
         super().release_mlp()