chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .config import AUTOTP_MODE, get_tensor_parallel_config
|
||||
from .tp_manager import TpTrainingManager
|
||||
@@ -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
|
||||
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import base64
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
import hjson
|
||||
import torch
|
||||
|
||||
from deepspeed.runtime.config_utils import dict_raise_error_on_duplicate_keys
|
||||
|
||||
_TP_MODEL_INIT_ARGS = None
|
||||
|
||||
|
||||
def load_ds_config(config: Union[str, dict]) -> dict:
|
||||
if isinstance(config, dict):
|
||||
return config
|
||||
if isinstance(config, str):
|
||||
if os.path.exists(config):
|
||||
return hjson.load(open(config, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys)
|
||||
try:
|
||||
config_decoded = base64.urlsafe_b64decode(config).decode('utf-8')
|
||||
return hjson.loads(config_decoded)
|
||||
except (UnicodeDecodeError, AttributeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. "
|
||||
f"Received: {config}") from exc
|
||||
raise ValueError(f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. "
|
||||
f"Received: {config}")
|
||||
|
||||
|
||||
def record_tp_model_init_args(tp_size, dtype, tp_group, dist_module):
|
||||
global _TP_MODEL_INIT_ARGS
|
||||
new_args = {
|
||||
"tp_size": tp_size,
|
||||
"dtype": dtype,
|
||||
"tp_group": tp_group,
|
||||
}
|
||||
|
||||
if _TP_MODEL_INIT_ARGS is None:
|
||||
_TP_MODEL_INIT_ARGS = new_args
|
||||
return
|
||||
|
||||
if _TP_MODEL_INIT_ARGS["tp_size"] != tp_size or _TP_MODEL_INIT_ARGS["dtype"] != dtype:
|
||||
raise ValueError("Conflicting tp_model_init arguments detected across multiple calls.")
|
||||
|
||||
existing_group = _TP_MODEL_INIT_ARGS.get("tp_group")
|
||||
if existing_group is None and tp_group is None:
|
||||
return
|
||||
if (existing_group is None) != (tp_group is None):
|
||||
raise ValueError("Conflicting tp_model_init arguments detected across multiple calls.")
|
||||
|
||||
existing_group_size = tp_group_world_size(existing_group, dist_module)
|
||||
new_group_size = tp_group_world_size(tp_group, dist_module)
|
||||
if existing_group_size != new_group_size:
|
||||
raise ValueError("Conflicting tp_model_init arguments detected across multiple calls.")
|
||||
|
||||
|
||||
def tp_group_world_size(tp_group, dist_module):
|
||||
if tp_group is None or dist_module is None:
|
||||
return None
|
||||
return dist_module.get_world_size(group=tp_group)
|
||||
|
||||
|
||||
def infer_config_dtype(config_dict: dict) -> Optional[torch.dtype]:
|
||||
bf16_config = config_dict.get("bf16", {})
|
||||
if isinstance(bf16_config, dict) and bf16_config.get("enabled", False):
|
||||
return torch.bfloat16
|
||||
fp16_config = config_dict.get("fp16", {})
|
||||
if isinstance(fp16_config, dict) and fp16_config.get("enabled", False):
|
||||
return torch.float16
|
||||
return None
|
||||
|
||||
|
||||
def merge_tp_model_init_into_config(config_dict: dict, mpu, mesh_param, dist_module):
|
||||
if _TP_MODEL_INIT_ARGS is None:
|
||||
return
|
||||
|
||||
tp_size = _TP_MODEL_INIT_ARGS["tp_size"]
|
||||
dtype = _TP_MODEL_INIT_ARGS["dtype"]
|
||||
tp_group = _TP_MODEL_INIT_ARGS["tp_group"]
|
||||
|
||||
if tp_group is not None and mpu is not None:
|
||||
raise ValueError("tp_model_init provided tp_group; deepspeed.initialize must not receive mpu.")
|
||||
if tp_group is None and mpu is None and mesh_param is None:
|
||||
# Auto-create TP groups for compatibility with HF Trainer (mpu is not passed).
|
||||
from deepspeed.utils import groups
|
||||
groups._init_tp_mesh_device(tensor_model_parallel_size=tp_size)
|
||||
|
||||
tp_section = config_dict.get("tensor_parallel")
|
||||
if tp_section is None:
|
||||
tp_section = {}
|
||||
config_dict["tensor_parallel"] = tp_section
|
||||
|
||||
config_autotp_size = tp_section.get("autotp_size")
|
||||
if config_autotp_size is not None and config_autotp_size != tp_size:
|
||||
raise ValueError(
|
||||
f"Conflicting tensor_parallel.autotp_size in config ({config_autotp_size}) and tp_model_init ({tp_size}).")
|
||||
|
||||
if config_autotp_size is None:
|
||||
tp_section["autotp_size"] = tp_size
|
||||
|
||||
tp_config = tp_section.get("tp") or {}
|
||||
if not isinstance(tp_config, dict):
|
||||
raise ValueError("tensor_parallel.tp must be a dict when provided.")
|
||||
|
||||
config_tp_size = tp_config.get("tp_size")
|
||||
if config_tp_size is not None and config_tp_size != tp_size:
|
||||
raise ValueError(
|
||||
f"Conflicting tensor_parallel.tp.tp_size in config ({config_tp_size}) and tp_model_init ({tp_size}).")
|
||||
if config_tp_size is None:
|
||||
tp_config["tp_size"] = tp_size
|
||||
|
||||
if tp_group is not None:
|
||||
config_tp_group = tp_config.get("tp_group")
|
||||
if config_tp_group is not None and config_tp_group is not tp_group:
|
||||
raise ValueError("Conflicting tensor_parallel.tp.tp_group in config and tp_model_init.")
|
||||
tp_config["tp_group"] = tp_group
|
||||
|
||||
tp_group_size = tp_group_world_size(tp_group, dist_module)
|
||||
if tp_group_size is not None and tp_group_size != tp_size:
|
||||
raise ValueError(f"tp_model_init tp_size ({tp_size}) does not match tp_group size ({tp_group_size}).")
|
||||
|
||||
tp_section["tp"] = tp_config
|
||||
|
||||
config_dtype = infer_config_dtype(config_dict)
|
||||
if config_dtype is not None and config_dtype != dtype:
|
||||
raise ValueError(f"Conflicting dtype: config uses {config_dtype} but tp_model_init requested {dtype}.")
|
||||
|
||||
tp_dtype = tp_section.get("dtype")
|
||||
if tp_dtype is not None:
|
||||
if isinstance(tp_dtype, str):
|
||||
tp_dtype_map = {
|
||||
"fp16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
"fp32": torch.float32,
|
||||
}
|
||||
tp_dtype_value = tp_dtype_map.get(tp_dtype.lower())
|
||||
else:
|
||||
tp_dtype_value = tp_dtype
|
||||
if tp_dtype_value is not None and tp_dtype_value != dtype:
|
||||
raise ValueError(f"Conflicting tensor_parallel.dtype in config ({tp_dtype}) and tp_model_init ({dtype}).")
|
||||
@@ -0,0 +1,66 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from .config import TPTrainingConfig, TPConfig
|
||||
from deepspeed.utils import groups
|
||||
import deepspeed.comm as dist
|
||||
|
||||
|
||||
class TpTrainingManager():
|
||||
|
||||
def __init__(self, model, tp_size, dtype):
|
||||
self.module = model
|
||||
self.config = self._initialize_config(dtype)
|
||||
|
||||
from deepspeed.module_inject.auto_tp import AutoTP
|
||||
from deepspeed import get_accelerator
|
||||
|
||||
# Parse model configuration
|
||||
parser_dict = AutoTP.tp_parser(model)
|
||||
print("AutoTP: ", parser_dict)
|
||||
|
||||
# Initialize TP configuration and model
|
||||
self._initialize_tp_config(tp_size)
|
||||
self._get_model_config_generate()
|
||||
|
||||
# Synchronize random number generator state across devices
|
||||
_rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name())
|
||||
dist.broadcast(_rng_state, groups.get_tensor_model_parallel_src_rank(), self.tp_config.tp_group)
|
||||
get_accelerator().set_rng_state(_rng_state.cpu())
|
||||
|
||||
# Apply injection policies
|
||||
self._apply_policies(parser_dict)
|
||||
|
||||
def _initialize_config(self, dtype):
|
||||
"""Initialize and return the DeepSpeed TP training configuration."""
|
||||
config = TPTrainingConfig()
|
||||
config.dtype = dtype
|
||||
return config
|
||||
|
||||
def _apply_policies(self, parser_dict):
|
||||
"""Apply injection policies to the parsed modules."""
|
||||
for client_module, injection_policy in parser_dict:
|
||||
self.config.injection_policy_tuple = injection_policy
|
||||
self._apply_injection_policy(self.config, client_module)
|
||||
|
||||
def _apply_injection_policy(self, config, client_module=None):
|
||||
from deepspeed.module_inject import replace_transformer_layer
|
||||
"""Apply the given injection policy to a client module."""
|
||||
if isinstance(self.module, torch.nn.Module):
|
||||
replace_transformer_layer(client_module, self.module, None, self.config, self.model_config)
|
||||
|
||||
def _initialize_tp_config(self, tp_size):
|
||||
"""Perform TP configuration initialization."""
|
||||
self.tp_config = TPConfig()
|
||||
self.tp_config.tp_size = tp_size
|
||||
|
||||
groups._init_tp_mesh_device(tp_size)
|
||||
self.tp_config.tp_group = groups.get_tensor_model_parallel_group()
|
||||
self.config.tensor_parallel = self.tp_config
|
||||
|
||||
def _get_model_config_generate(self):
|
||||
"""Generate and apply HF model configuration."""
|
||||
self.model_config = getattr(self.module, 'config', None)
|
||||
Reference in New Issue
Block a user