Drop the unused kv-head override from the shard-size helpers

The override existed because get_shard_size read a process-wide num_kv_heads
that was only set during AutoTP replacement. The inference engine patches the
alibi helpers before that happens, so it probed the head count itself and passed
it in to bypass the uninitialized global. That was the parameter's only caller.

With the count carried by a per-model AutoTPMeta, the alibi path receives its
own model's value like everyone else, and the previous commit removed the
round trip that read meta.num_kv_heads only to pass it straight back. Nothing
outside tp_shard supplies the argument now, so remove it and the two tests
written against it; get_shard_size has a single source for the head count again.

Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Author: delock

deepspeed/module_inject/tp_shard.py

@@ -82,7 +82,7 @@ def from_model_config(cls, model_config, tp_grain_size: int = 1) -> "AutoTPMeta"
                    tp_grain_size=tp_grain_size)
 
 
-def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None, eff_num_kv_heads=None):
+def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None, mp_group=None):
     """Size of one shard of ``total_size`` split across a tensor-parallel group of ``mp_size``.
 
     ``meta`` carries this model's ``num_kv_heads`` / ``tp_grain_size`` so the split is stable
@@ -92,14 +92,8 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
     ``rank`` is the rank *within the tensor-parallel group*, i.e. in ``[0, mp_size)``, matching
     ``dist.get_rank(group=mp_group)`` and the index used by ``get_shard_size_list``. It is not a
     global rank.
-
-    ``eff_num_kv_heads`` is the head count this split is actually aligned to; it defaults to
-    ``meta.num_kv_heads``. Passing it explicitly lets callers split fused sub-parameters (Q/K/V)
-    against their respective head counts without reimplementing the KV-head-aligned partition
-    logic at the call site.
     """
-    if eff_num_kv_heads is None:
-        eff_num_kv_heads = meta.num_kv_heads
+    num_kv_heads = meta.num_kv_heads
     tp_grain_size = meta.tp_grain_size
     last_linear = ["lm_head", "embed_out"]
     # MoE MLP layer use near even division will get better perf.
@@ -117,10 +111,10 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
                 raise ValueError("get_shard_size requires a group-local rank or process group when mp_size "
                                  f"({mp_size}) differs from the distributed world size ({world_size}).")
             rank = dist.get_rank()
-    if eff_num_kv_heads is not None and total_size % eff_num_kv_heads == 0 and "mlp" not in str(name) and \
+    if num_kv_heads is not None and total_size % num_kv_heads == 0 and "mlp" not in str(name) and \
         str(name) not in last_linear and not_moe_mlp_layer:
-        my_slices = (eff_num_kv_heads // mp_size) + (1 if rank < (eff_num_kv_heads % mp_size) else 0)
-        return total_size * my_slices // eff_num_kv_heads
+        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 >= tp_grain_size:
             grain_size, remainder = divmod(total_size, tp_grain_size)
@@ -135,16 +129,14 @@ def get_shard_size(total_size, mp_size, meta: AutoTPMeta, name=None, rank=None,
             return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
 
 
-def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None, eff_num_kv_heads=None):
+def get_shard_size_list(total_size, mp_size, meta: AutoTPMeta, name=None):
     shard_sizes = []
-    if eff_num_kv_heads is None:
-        eff_num_kv_heads = meta.num_kv_heads
     for i in range(mp_size):
-        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i, eff_num_kv_heads=eff_num_kv_heads))
+        shard_sizes.append(get_shard_size(total_size, mp_size, meta, name, i))
     # Shards must tile the dimension exactly, otherwise the partitioned weights no longer
     # reconstruct the original tensor.
     assert sum(shard_sizes) == total_size, (
         f"AutoTP shard sizes {shard_sizes} for layer '{name}' do not sum to the dimension size "
         f"{total_size} with tp_size={mp_size}, tp_grain_size={meta.tp_grain_size} and "
-        f"num_kv_heads={eff_num_kv_heads}.")
+        f"num_kv_heads={meta.num_kv_heads}.")
     return shard_sizes

tests/unit/module_inject/test_tp_shard.py

@@ -64,19 +64,3 @@ def test_shard_size_refuses_to_guess_subgroup_rank(monkeypatch):
 
     with pytest.raises(ValueError, match="group-local rank or process group"):
         get_shard_size(12, 2, AutoTPMeta())
-
-
-def test_explicit_num_kv_heads_overrides_meta():
-    # Fused Q/K/V sub-parameters are split against their own head counts, which differ from
-    # the model-wide kv-head count carried by the meta.
-    meta = AutoTPMeta(num_kv_heads=2)
-
-    assert get_shard_size_list(384, 4, meta, "self_attn.q_proj", eff_num_kv_heads=6) == [128, 128, 64, 64]
-
-
-def test_explicit_num_kv_heads_matches_meta_value():
-    expected = get_shard_size_list(384, 4, AutoTPMeta(num_kv_heads=6), "self_attn.q_proj")
-
-    actual = get_shard_size_list(384, 4, AutoTPMeta(), "self_attn.q_proj", eff_num_kv_heads=6)
-
-    assert actual == expected