goodjava
diff --git a/‎.github/workflows/amd.yml‎
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/amd.yml‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎deepspeed/module_inject/replace_module.py‎
Lines changed: 8 additions & 0 deletions b/‎deepspeed/module_inject/replace_module.py‎
Lines changed: 8 additions & 0 deletions
@@ -15,7 +15,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  unit-tests:
+  amd-tests:
     # The type of runner that the job will run on
     runs-on: [self-hosted, amd]
 
 
@@ -464,20 +464,28 @@ def _slice_embedding(child, name, conv_linear_layer):
 
         def update_mp_params(child):
             if hasattr(child, 'n_heads'):
+                assert child.n_heads%mp_size == 0, "n_heads ({}) must be divisible by mp_size ({})".format(child.n_heads, mp_size)
                 child.n_heads = child.n_heads // mp_size
             if hasattr(child, 'inner_dim'):
+                assert child.inner_dim%mp_size == 0, "inner_dim ({}) must be divisible by mp_size ({})".format(child.inner_dim, mp_size)
                 child.inner_dim = child.inner_dim // mp_size
             if hasattr(child, 'num_heads'):
+                assert child.num_heads%mp_size == 0, "num_heads ({}) must be divisible by mp_size ({})".format(child.num_heads, mp_size)
                 child.num_heads = child.num_heads // mp_size
             if hasattr(child, 'num_attention_heads'):
+                assert child.num_attention_heads%mp_size == 0, "num_attention_heads ({}) must be divisible by mp_size ({})".format(child.num_attention_heads, mp_size)
                 child.num_attention_heads = child.num_attention_heads // mp_size
             if hasattr(child, 'num_attn_heads'):
+                assert child.num_attn_heads%mp_size == 0, "num_attn_heads ({}) must be divisible by mp_size ({})".format(child.num_attn_heads, mp_size)
                 child.num_attn_heads = child.num_attn_heads // mp_size
             if hasattr(child, 'all_head_size'):
+                assert child.all_head_size%mp_size == 0, "all_head_size ({}) must be divisible by mp_size ({})".format(child.all_head_size, mp_size)
                 child.all_head_size = child.all_head_size // mp_size
             if hasattr(child, 'embed_dim'):
+                assert child.embed_dim%mp_size == 0, "embed_dim must ({}) be divisible by mp_size ({})".format(child.embed_dim, mp_size)
                 child.embed_dim = child.embed_dim // mp_size
             if hasattr(child, 'hidden_size'):
+                assert child.hidden_size%mp_size == 0, "hidden_size ({}) must be divisible by mp_size ({})".format(child.hidden_size, mp_size)
                 child.hidden_size = child.hidden_size // mp_size
 
         conv_linear_layer = False