chore: import upstream snapshot with attribution

deepspeed/runtime/tensor_parallel/config.py

@@ -0,0 +1,163 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from enum import Enum
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+import torch
+from pydantic import Field
+from typing import Optional, Dict, Any
+
+
+class AUTOTP_MODE(Enum):
+    TRAINING = "TRAINING"
+    INFERENCE = "INFERENCE"
+
+
+class TPConfig(DeepSpeedConfigModel):
+    """ Configure tensor parallelism settings """
+
+    tp_size: int = 1
+    """ Number of devices to split the model across using tensor parallelism. """
+
+    tp_grain_size: int = 1
+    "The variable required by the autoTP parser has not been activated in training yet"
+    "as it depends on the gather logic that supports uneven partitioning. "
+    "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size."
+
+    mpu: object = None
+    """
+    A model parallelism unit object that implements
+    ``get_{model,data}_parallel_{rank,group,world_size}()``.
+    """
+
+    tp_group: object = None
+
+
+class TPTrainingConfig(DeepSpeedConfigModel):
+
+    dtype: torch.dtype = torch.float16
+    """
+    Desired model data type, will convert model to this type.
+    """
+
+    autotp_size: int = 0
+    """
+    In automatic tensor-parallelism training, 'tensor_parallel_size'
+    When set to 0, indicates that it is disabled.
+    """
+    tp_overlap_comm: bool = False
+    """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """
+
+    tensor_parallel: TPConfig = Field({}, alias="tp")
+    """
+    Configuration for tensor parallelism used to split the model across several
+    GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`.
+    """
+
+    injection_policy_tuple: Optional[tuple] = None
+
+    # New configurable AutoTP settings
+    partition_config: Optional[Dict[str, Any]] = None
+    """
+    Configuration for the new configurable AutoTP API.
+    Allows users to specify custom layer partitioning rules via TPLayerSpec.
+
+    Example:
+        "partition_config": {
+            "use_default_specs": false,
+            "layer_specs": [
+                {
+                    "patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"],
+                    "partition_type": "row"
+                },
+                {
+                    "patterns": [".*\\.[qkv]_proj\\.weight$"],
+                    "partition_type": "column"
+                },
+                {
+                    "patterns": [".*\\.gate_up_proj\\.weight$"],
+                    "partition_type": "column",
+                    "shape": [2, -1],
+                    "partition_dim": 0
+                }
+            ]
+        }
+    """
+
+    preset_model: Optional[str] = None
+    """
+    Use a built-in preset for common model architectures.
+    Available presets: "llama", "bloom", "chatglm", "mixtral", "deepseek_v2", "qwen2", "phi3"
+    """
+
+    #The following parameters are required by autoTP parser.
+    ########################################
+    keep_module_on_host: bool = False
+    """
+    When loading checkpoints to model parameters, they are moved to the device. In very large models
+    this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on
+    host and not move them directly to the device (giving an option to quantize checkpoint data before
+    moving it to the device for example).
+    """
+
+    replace_with_kernel_inject: bool = Field(False, alias="kernel_inject")
+    """
+    Set to true to inject inference kernels for models such as, Bert, GPT2,
+    GPT-Neo and GPT-J.  Otherwise, the injection_dict provides the names of two
+    linear layers as a tuple:
+    `(attention_output projection, transformer output projection)`
+    """
+
+    ########################################
+
+    def get_partition_config_object(self):
+        """
+        Get the AutoTPConfig object from the configuration.
+        Returns None if no custom config is specified.
+        """
+        from deepspeed.module_inject.autotp_config import AutoTPConfig, AutoTPPresets, merge_autotp_configs
+
+        config = None
+
+        # First check for preset
+        if self.preset_model:
+            config = AutoTPPresets.get_preset(self.preset_model)
+
+        # Then check for custom config
+        if self.partition_config:
+            custom_config = AutoTPConfig.from_dict(self.partition_config)
+            if config and custom_config.use_default_specs:
+                config = merge_autotp_configs(config, custom_config)
+            else:
+                config = custom_config
+
+        if config:
+            config.tp_size = self.autotp_size
+
+        return config
+
+
+def get_tensor_parallel_config(ds_config):
+
+    if 'tensor_parallel' in ds_config:
+        return TPTrainingConfig(**ds_config['tensor_parallel'])
+    return TPTrainingConfig()
+
+
+def _get_hf_tp_plan(model):
+    """Extract tp_plan from HuggingFace model.
+
+    Prefer base_model_tp_plan (from model config) over _tp_plan (runtime attribute)
+    because _tp_plan often contains duplicate entries with a 'model.' prefix added
+    by HuggingFace, which causes spurious duplicate-match warnings during conversion.
+    """
+    config = getattr(model, 'config', None)
+    if config and getattr(config, 'base_model_tp_plan', None):
+        return model.config.base_model_tp_plan
+
+    if getattr(model, '_tp_plan', None):
+        return model._tp_plan
+
+    return None