[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta

deepspeed/inference/engine.py

@@ -25,6 +25,7 @@
 from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
 from deepspeed.accelerator import get_accelerator
 from ..module_inject.policy import TransformerPolicy
+from deepspeed.module_inject.tp_shard import AutoTPMeta
 from ..module_inject.auto_tp import AutoTP
 
 from ..module_inject.replace_policy import generic_policies
@@ -218,52 +219,44 @@ def build_alibi_tensor(self):
             if hasattr(self.module.transformer, 'build_alibi_tensor'):
                 # The heads must be sliced with the same tensor-parallel group that partitioned
                 # the attention weights, so bind it rather than letting the helper guess.
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                shard_sizes = get_head_shard_sizes(num_heads, self.mp_group, num_kv_heads)
-                self.module.transformer.build_alibi_tensor = functools.partial(build_bloom_alibi_tensor,
-                                                                               mp_group=self.mp_group,
-                                                                               head_shard_sizes=shard_sizes,
-                                                                               total_num_heads=num_heads)
+                meta = self._autotp_meta(self.module.transformer)
+                shard_sizes = get_head_shard_sizes(meta, self.mp_group)
+                self.module.transformer.build_alibi_tensor = functools.partial(
+                    build_bloom_alibi_tensor,
+                    mp_group=self.mp_group,
+                    head_shard_sizes=shard_sizes,
+                    total_num_heads=meta.num_attention_heads)
             if hasattr(self.module.transformer, 'build_mpt_alibi_tensor'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
+                meta = self._autotp_meta(self.module.transformer)
                 install_head_sharded_helper(self.module.transformer, 'build_mpt_alibi_tensor', build_mpt_alibi_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
+                                            meta, self.mp_group)
         if hasattr(self.module, 'model'):
             if hasattr(self.module.model, 'get_alibi_mask'):
-                num_heads = self._get_model_head_count(self.module.model)
-                num_kv_heads = self._get_model_kv_head_count(self.module.model, num_heads)
-                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, self.mp_group,
-                                            num_heads, num_kv_heads)
+                meta = self._autotp_meta(self.module.model)
+                install_head_sharded_helper(self.module.model, 'get_alibi_mask', get_alibi_mask, meta, self.mp_group)
 
     def build_attn_bias(self):
         if hasattr(self.module, 'transformer'):
             if hasattr(self.module.transformer, '_attn_bias'):
-                num_heads = self._get_model_head_count(self.module.transformer)
-                num_kv_heads = self._get_model_kv_head_count(self.module.transformer, num_heads)
-                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor,
-                                            self.mp_group, num_heads, num_kv_heads)
-
-    def _get_model_head_count(self, module):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            for name in ("num_heads", "n_heads", "n_head", "num_attention_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
-        raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
-
-    def _get_model_kv_head_count(self, module, num_heads):
-        for source in (module, getattr(module, "config", None), getattr(self.module, "config", None)):
-            if source is None:
-                continue
-            for name in ("num_key_value_heads", "num_kv_heads", "n_head_kv", "kv_n_heads"):
-                value = getattr(source, name, None)
-                if value is not None:
-                    return value
-        return num_heads
+                meta = self._autotp_meta(self.module.transformer)
+                install_head_sharded_helper(self.module.transformer, '_attn_bias', build_mpt_atten_bias_tensor, meta,
+                                            self.mp_group)
+
+    def _autotp_meta(self, module):
+        # from_model_config extracts from a single config object, but the head counts may live on
+        # the module, its config, or the top-level model config. Probe each source via the shared
+        # from_model_config and merge per field -- num_attention_heads is mandatory (alibi needs
+        # it), num_kv_heads is None for non-GQA. Subsumes the former
+        # _get_model_head_count / _get_model_kv_head_count pair.
+        metas = [
+            AutoTPMeta.from_model_config(s)
+            for s in (module, getattr(module, "config", None), getattr(self.module, "config", None)) if s is not None
+        ]
+        num_heads = next((m.num_attention_heads for m in metas if m.num_attention_heads is not None), None)
+        if num_heads is None:
+            raise ValueError(f"Cannot determine the attention head count for {module.__class__.__name__}.")
+        num_kv_heads = next((m.num_kv_heads for m in metas if m.num_kv_heads is not None), None)
+        return AutoTPMeta(num_attention_heads=num_heads, num_kv_heads=num_kv_heads)
 
     def _pre_forward_hook(self, module, *inputs, **kwargs):
         if self.use_cuda_events:

deepspeed/module_inject/auto_tp.py

@@ -15,7 +15,7 @@
 from .layers import *
 from deepspeed.accelerator import get_accelerator
 from .fusedqkv_utils import require_tp_fused_qkvw
-from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size, get_shard_size_list
 from deepspeed.utils import groups
 from deepspeed.utils.logging import log_dist
 from deepspeed.module_inject.layers import is_autotp_training_mode
@@ -205,14 +205,19 @@ def __init__(self,
                  linear_layer_setting,
                  orig_layer_impl,
                  keep_module_on_host=False,
-                 partition_config: Optional[AutoTPConfig] = None):
+                 partition_config: Optional[AutoTPConfig] = None,
+                 model_config=None,
+                 tp_grain_size: int = 1):
         self.module = module
         self.all_reduce_linears = all_reduce_linears
         self.prefix = prefix
         self.state_dict = state_dict
 
         self.mp_size = None
         self.mp_group = None
+        # Per-model TP metadata threaded through every layer / helper so each AutoTP instance
+        # shards by its own model's kv-head / grain values.
+        self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size)
         self.linear_layer_setting = linear_layer_setting
         self.orig_layer_impl = orig_layer_impl
         self.linear_policies = None
@@ -378,14 +383,14 @@ def _replace(self, child, name, conv_linear_layer):
         # For Yuan model
         if 'Yuan' in str(self.module):
             if 'v_proj' in name:
-                return Yuan_LinearLayer(child, self.mp_group)
+                return Yuan_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
 
             elif 'o_proj' in name:
-                return Yuan_LinearAllreduce(child, self.mp_group)
+                return Yuan_LinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
         # For MLP including chunk layer.
         if 'gate_up_proj' in name or ('dense_h_to_4h' in name and 'GLM' in str(self.module)):
-            return GateUpPack_LinearLayer(child, self.mp_group)
+            return GateUpPack_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             # For Arctic model, bypass to all_reduce replacement for w2 weights
         arctic_w2_all_reduce_linear = False
         if 'Arctic' in str(self.module) and 'w2' in name:
@@ -398,21 +403,21 @@ def _replace(self, child, name, conv_linear_layer):
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                return Conv_LinearALlreduce(child, self.mp_group, name=name)
+                return Conv_LinearALlreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
             elif name == "lm_head" or name == 'embed_out':
-                return LmHeadLinearAllreduce(child, self.mp_group)
+                return LmHeadLinearAllreduce(child, self.mp_group, tp_meta=self.tp_meta)
 
-            return LinearAllreduce(child, self.mp_group, name=name)
+            return LinearAllreduce(child, self.mp_group, name=name, tp_meta=self.tp_meta)
         else:
 
             setattr(child, "replaced", True)
             if self.conv_linear_layer:
-                conv_LinearLayer(child, self.mp_group)
+                conv_LinearLayer(child, self.mp_group, tp_meta=self.tp_meta)
             elif require_tp_fused_qkvw(name, self.mp_size):
                 #Check and handle fused qkv for TP
-                return fused_LinearLayer(child, self.mp_group, fused_module=self.module)
+                return fused_LinearLayer(child, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
 
-            return LinearLayer(child, self.mp_group, name=name)
+            return LinearLayer(child, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _replace_with_config(self, child, name):
         """
@@ -450,10 +455,10 @@ def _replace_with_config(self, child, name):
     def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create row-parallel layer (AllReduce after forward)."""
         if self.conv_linear_layer:
-            return Conv_LinearALlreduce(module, self.mp_group, name=name)
+            return Conv_LinearALlreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
         # Check for lm_head / embed_out
         if name == "lm_head" or name == 'embed_out':
-            return LmHeadLinearAllreduce(module, self.mp_group)
+            return LmHeadLinearAllreduce(module, self.mp_group, tp_meta=self.tp_meta)
 
         if spec.shape is not None:
             return SubParamLinearAllreduce(
@@ -462,17 +467,22 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearAllreduce(module, self.mp_group, name=name)
+        return LinearAllreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
 
     def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
         """Create column-parallel layer (AllReduce in backward)."""
         if self.conv_linear_layer:
-            return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+            return conv_LinearLayer(module,
+                                    self.mp_group,
+                                    name=name,
+                                    gather_output=spec.gather_output,
+                                    tp_meta=self.tp_meta)
         # Only use fused-QKV heuristics when no partition_config is provided.
         elif self.partition_config is None and require_tp_fused_qkvw(name, self.mp_size):
             # Check and handle fused qkv for TP
-            return fused_LinearLayer(module, self.mp_group, fused_module=self.module)
+            return fused_LinearLayer(module, self.mp_group, fused_module=self.module, tp_meta=self.tp_meta)
         if spec.shape is not None:
             if spec.gather_output:
                 raise NotImplementedError("AutoTP gather_output does not yet support shaped sub-parameter layers.")
@@ -482,8 +492,9 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
                 shape=spec.shape,
                 partition_dim=spec.get_partition_dim(),
                 name=name,
+                tp_meta=self.tp_meta,
             )
-        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
+        return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output, tp_meta=self.tp_meta)
 
     def _configure_gathered_column_tie_fallbacks(self):
         """Configure a replicated fallback for gathered output layers tied to embeddings."""
@@ -561,7 +572,7 @@ def _slice_embedding(self, child, name, conv_linear_layer):
         mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
 
         original_shape = tuple(child.weight.shape)
-        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, name)
+        partition_sizes = get_shard_size_list(original_shape[1], self.mp_size, self.tp_meta, name)
         if hasattr(child.weight, 'ds_tensor'):
             data = child.weight.ds_tensor.data.split(partition_sizes, dim=1)
         else:
@@ -645,7 +656,7 @@ def update_mp_params(self, child, name=None):
                 param_val = getattr(child, param)
                 # get_shard_size selects its partitioning strategy from the module name, so the
                 # attributes must be sharded under the same name as the weights they describe.
-                setattr(child, param, get_shard_size(param_val, self.mp_size, name, rank=tp_index))
+                setattr(child, param, get_shard_size(param_val, self.mp_size, self.tp_meta, name, rank=tp_index))
         setattr(child, "replaced", True)
 
     def update_linear_policies(self):
@@ -754,21 +765,6 @@ def _replace_module(self, r_module, prev_name='', prev_class_name=''):
                 self._replace_module(child, name, class_name)
         return r_module
 
-    @staticmethod
-    def get_model_num_kv_heads(config):
-        num_kv_heads = None
-        # multi_query_group_num is for chatglm2 & chatglm3
-        kv_head_names = [
-            'multi_query_group_num', 'num_kv_heads', 'num_key_value_heads', 'num_attention_heads', 'n_heads',
-            'attention_heads'
-        ]
-        for name in kv_head_names:
-            if hasattr(config, name):
-                num_kv_heads = getattr(config, name)
-                if num_kv_heads is not None:
-                    break
-        return num_kv_heads
-
     def _replace_last_linear_module(self, r_module):
         if hasattr(r_module, "lm_head"):
             name = "lm_head"

deepspeed/module_inject/auto_tp_model_utils.py

@@ -8,7 +8,7 @@
 from deepspeed import comm as dist
 import torch
 from typing import Optional
-from deepspeed.module_inject.tp_shard import get_shard_size_list
+from deepspeed.module_inject.tp_shard import AutoTPMeta, get_shard_size_list
 
 
 class _HeadCountProxy:
@@ -27,16 +27,12 @@ def __setattr__(self, name, value):
             setattr(self._module, name, value)
 
 
-def get_head_shard_sizes(num_heads, mp_group=None, num_kv_heads=None):
+def get_head_shard_sizes(meta: AutoTPMeta, mp_group=None):
     tp_world_size = dist.get_world_size(group=mp_group)
-    return get_shard_size_list(
-        num_heads,
-        tp_world_size,
-        num_kv_heads=num_kv_heads,
-    )
+    return get_shard_size_list(meta.num_attention_heads, tp_world_size, meta)
 
 
-def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=None, num_kv_heads=None):
+def install_head_sharded_helper(module, name, wrapper, meta, mp_group=None):
     """Give ``module`` a head-slicing wrapper around one of its own methods.
 
     The wrapper is bound to this instance instead of installed on its class. A class-wide patch
@@ -48,10 +44,15 @@ def install_head_sharded_helper(module, name, wrapper, mp_group=None, num_heads=
     if original_name not in module.__dict__:
         # Wrapping an already wrapped instance would make it delegate to itself.
         setattr(module, original_name, getattr(module, name))
-    shard_sizes = get_head_shard_sizes(num_heads, mp_group, num_kv_heads) if num_heads is not None else None
+    total_num_heads = meta.num_attention_heads
+    shard_sizes = get_head_shard_sizes(meta, mp_group) if total_num_heads is not None else None
     setattr(
         module, name,
-        functools.partial(wrapper, module, mp_group=mp_group, head_shard_sizes=shard_sizes, total_num_heads=num_heads))
+        functools.partial(wrapper,
+                          module,
+                          mp_group=mp_group,
+                          head_shard_sizes=shard_sizes,
+                          total_num_heads=total_num_heads))
 
 
 def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads=None):
@@ -64,7 +65,9 @@ def _head_shard(num_heads, mp_group=None, head_shard_sizes=None, total_num_heads
     tp_world_size = dist.get_world_size(group=mp_group)
     tp_index = dist.get_rank(group=mp_group)
     full_num_heads = total_num_heads if total_num_heads is not None else num_heads
-    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size)
+    # The fallback split is only reached when no per-head shard sizes were recorded, i.e. when
+    # the model is not GQA-aware, so an even split (default meta, no kv-heads) is correct here.
+    shard_sizes = head_shard_sizes or get_shard_size_list(full_num_heads, tp_world_size, AutoTPMeta())
     if len(shard_sizes) != tp_world_size or sum(shard_sizes) != full_num_heads:
         raise ValueError(f"Head shard sizes {shard_sizes} do not partition {full_num_heads} heads across "
                          f"{tp_world_size} tensor-parallel ranks.")