fix uneven heads issue (microsoft#25)

deepspeed/module_inject/tp_shard.py

@@ -19,17 +19,14 @@ def get_num_kv_heads():
 
 def get_shard_size(total_size, mp_size, rank=None):
     global num_kv_heads
-    # When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
-    if num_kv_heads != None:
-        if (rank == None):
-            rank = dist.get_rank()
+    # When we have num_kv_heads defined, uneven division is possible, otherwise enforce near even division
+    if rank == None:
+        rank = dist.get_rank()
+    if num_kv_heads != None and total_size % num_kv_heads == 0:
         my_slices = (num_kv_heads // mp_size) + (1 if rank < (num_kv_heads % mp_size) else 0)
         return total_size * my_slices // num_kv_heads
     else:
-        if total_size % mp_size == 0:
-            return total_size // mp_size
-        else:
-            assert False, f"Number of attention heads ({total_size}) must be divisible by mp_size ({mp_size})"
+        return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
 def get_shard_size_list(total_size, mp_size):