chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
Executable
+11
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .replace_module import replace_transformer_layer, revert_transformer_layer, ReplaceWithTensorSlicing, GroupQuantizer, generic_injection
|
||||
from .module_quantize import quantize_transformer_layer
|
||||
from .replace_policy import HFBertLayerPolicy
|
||||
from .layers import LinearAllreduce, LinearLayer, EmbeddingLayer, Normalize, set_autotp_mode, SubParamLinearLayer, SubParamLinearAllreduce
|
||||
from .policy import DSPolicy
|
||||
from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType, AutoTPPresets, merge_autotp_configs
|
||||
@@ -0,0 +1,599 @@
|
||||
# Copyright (c) DeepSpeed Team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""AutoEP: Automatic Expert Parallelism for MoE models.
|
||||
|
||||
Phase 3: MoE layer detection and structural validation.
|
||||
Phase 5: Layer replacement (replace_moe_layer filled in).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from deepspeed.module_inject.auto_ep_config import (
|
||||
fill_autoep_config_from_hf,
|
||||
AutoEPConfig,
|
||||
MoELayerSpec,
|
||||
MoEModelPreset,
|
||||
)
|
||||
from deepspeed.module_inject.auto_ep_presets.base import ForwardContract
|
||||
from deepspeed.module_inject.auto_ep_presets.registry import (
|
||||
apply_config_overrides,
|
||||
get_preset_adapter,
|
||||
preset_name_for_hf_model_type,
|
||||
resolve_preset_candidates,
|
||||
unsupported_preset_for_hf_model_type,
|
||||
)
|
||||
from deepspeed.moe.fused_expert_layout import classify_fused_gate_up_layout
|
||||
from deepspeed.runtime.zero.utils import is_zero_param
|
||||
from deepspeed.utils import logger
|
||||
|
||||
|
||||
def _remove_transformers_output_capture_hooks(model: nn.Module) -> int:
|
||||
"""Remove HF output-capturing hooks so they can be reinstalled after AutoEP conversion."""
|
||||
removed = 0
|
||||
for module in model.modules():
|
||||
hooks = getattr(module, "_forward_hooks", None)
|
||||
if not hooks:
|
||||
continue
|
||||
|
||||
for hook_id, hook in list(hooks.items()):
|
||||
if getattr(hook, "__name__", "") != "output_capturing_hook":
|
||||
continue
|
||||
del hooks[hook_id]
|
||||
removed += 1
|
||||
hooks_with_kwargs = getattr(module, "_forward_hooks_with_kwargs", None)
|
||||
if hooks_with_kwargs is not None:
|
||||
hooks_with_kwargs.pop(hook_id, None)
|
||||
hooks_always_called = getattr(module, "_forward_hooks_always_called", None)
|
||||
if hooks_always_called is not None:
|
||||
hooks_always_called.pop(hook_id, None)
|
||||
return removed
|
||||
|
||||
|
||||
def _is_known_hf_model_type(model_type: str | None) -> bool:
|
||||
if model_type is None:
|
||||
return False
|
||||
return (preset_name_for_hf_model_type(model_type) is not None
|
||||
or unsupported_preset_for_hf_model_type(model_type) is not None)
|
||||
|
||||
|
||||
def _raise_if_duplicate_moe_specs(specs: list[MoELayerSpec]) -> None:
|
||||
by_module: dict[str, list[MoELayerSpec]] = {}
|
||||
for spec in specs:
|
||||
by_module.setdefault(spec.moe_module_name, []).append(spec)
|
||||
|
||||
duplicates = {name: matches for name, matches in by_module.items() if len(matches) > 1}
|
||||
if not duplicates:
|
||||
return
|
||||
|
||||
details = "; ".join(f"{name}: {', '.join(spec.model_family for spec in matches)}"
|
||||
for name, matches in sorted(duplicates.items()))
|
||||
raise ValueError("AutoEP detection is ambiguous and produced multiple replacement specs for the same "
|
||||
f"MoE module(s): {details}. Set expert_parallel.preset_model or provide custom "
|
||||
"AutoEP patterns so each MoE module matches exactly one preset.")
|
||||
|
||||
|
||||
def _source_param_shape(param: torch.Tensor | nn.Parameter) -> torch.Size:
|
||||
if is_zero_param(param):
|
||||
return torch.Size(param.ds_shape)
|
||||
return torch.Size(param.shape)
|
||||
|
||||
|
||||
def _source_param_ndim(param: torch.Tensor | nn.Parameter) -> int:
|
||||
return len(_source_param_shape(param))
|
||||
|
||||
|
||||
def _has_3d_expert_params(module: nn.Module, preset: MoEModelPreset) -> bool:
|
||||
"""Check if module stores expert weights as 3D parameter tensors (transformers 5.0.0+).
|
||||
|
||||
Returns True if the module has a parameter named preset.expert_w1 (e.g., "gate_up_proj")
|
||||
with 3 dimensions (num_experts, ..., ...).
|
||||
"""
|
||||
w1_name = preset.expert_w1
|
||||
param = getattr(module, w1_name, None)
|
||||
if param is None:
|
||||
return False
|
||||
if isinstance(param, nn.Parameter) or isinstance(param, torch.Tensor):
|
||||
return _source_param_ndim(param) == 3
|
||||
return False
|
||||
|
||||
|
||||
def _get_num_experts_from_config(model_config, preset: MoEModelPreset) -> int | None:
|
||||
"""Extract num_experts from model.config using the preset's attribute name."""
|
||||
return getattr(model_config, preset.num_experts_attr, None)
|
||||
|
||||
|
||||
def _get_top_k_from_config(model_config, preset: MoEModelPreset) -> int | None:
|
||||
"""Extract top_k from model.config using the preset's attribute name."""
|
||||
return getattr(model_config, preset.top_k_attr, None)
|
||||
|
||||
|
||||
def _as_finite_float(value, field_name: str) -> float:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"{field_name} must be a finite number")
|
||||
|
||||
value = float(value)
|
||||
if not math.isfinite(value):
|
||||
raise ValueError(f"{field_name} must be a finite number")
|
||||
return value
|
||||
|
||||
|
||||
def _resolve_route_scale(config: AutoEPConfig, model_config) -> float:
|
||||
"""Resolve the single scale applied by TokenChoiceTopKRouter."""
|
||||
routed_scaling_factor = config.routed_scaling_factor
|
||||
|
||||
if routed_scaling_factor != "auto":
|
||||
route_scale = _as_finite_float(routed_scaling_factor, "routed_scaling_factor")
|
||||
if config.route_scale != 1.0:
|
||||
logger.warning("AutoEP: routed_scaling_factor=%s overrides route_scale=%s.", routed_scaling_factor,
|
||||
config.route_scale)
|
||||
return route_scale
|
||||
|
||||
cfg_routed_scaling_factor = getattr(model_config, 'routed_scaling_factor', None)
|
||||
if cfg_routed_scaling_factor is not None:
|
||||
route_scale = _as_finite_float(cfg_routed_scaling_factor, "model.config.routed_scaling_factor")
|
||||
if config.route_scale != 1.0:
|
||||
logger.warning("AutoEP: model.config.routed_scaling_factor=%s overrides route_scale=%s.",
|
||||
cfg_routed_scaling_factor, config.route_scale)
|
||||
return route_scale
|
||||
|
||||
return _as_finite_float(config.route_scale, "route_scale")
|
||||
|
||||
|
||||
def _detect_expert_storage(experts_module: nn.Module, preset: MoEModelPreset) -> Literal["fused_3d", "module_list"]:
|
||||
"""Determine whether experts are stored as fused 3D tensors or nn.ModuleList."""
|
||||
if _has_3d_expert_params(experts_module, preset):
|
||||
return "fused_3d"
|
||||
if isinstance(experts_module, nn.ModuleList):
|
||||
return "module_list"
|
||||
# Check children for 3D params as fallback
|
||||
for name, param in experts_module.named_parameters(recurse=False):
|
||||
if _source_param_ndim(param) == 3:
|
||||
return "fused_3d"
|
||||
return "module_list"
|
||||
|
||||
|
||||
def _infer_hidden_and_ffn_size(
|
||||
experts_module: nn.Module,
|
||||
preset: MoEModelPreset,
|
||||
storage: Literal["fused_3d", "module_list"],
|
||||
num_experts: int,
|
||||
) -> tuple[int, int]:
|
||||
"""Infer hidden_size and ffn_hidden_size from expert weight shapes."""
|
||||
if storage == "fused_3d":
|
||||
w1_param = getattr(experts_module, preset.expert_w1, None)
|
||||
w2_param = getattr(experts_module, preset.expert_w2, None)
|
||||
if w1_param is not None and w2_param is not None:
|
||||
w1_shape = _source_param_shape(w1_param)
|
||||
w2_shape = _source_param_shape(w2_param)
|
||||
if preset.expert_w3 is None:
|
||||
layout = classify_fused_gate_up_layout(tuple(w1_shape), tuple(w2_shape))
|
||||
if layout is None:
|
||||
raise ValueError("expert_w3=None expects fused gate+up weights with either "
|
||||
f"[E, 2*ffn, hidden]/[E, hidden, ffn] or [E, hidden, 2*ffn]/[E, ffn, hidden], "
|
||||
f"but got {preset.expert_w1}={tuple(w1_shape)} and "
|
||||
f"{preset.expert_w2}={tuple(w2_shape)}.")
|
||||
hidden_size = layout.hidden_size
|
||||
ffn_hidden_size = layout.ffn_hidden_size
|
||||
else:
|
||||
# Separate gate and up: w1 shape is [E, ffn, hidden]
|
||||
w3_param = getattr(experts_module, preset.expert_w3, None)
|
||||
if w3_param is None:
|
||||
raise ValueError(f"expert_w3='{preset.expert_w3}' is set but no such weight "
|
||||
f"exists on experts module.")
|
||||
hidden_size = w1_shape[2]
|
||||
ffn_hidden_size = w1_shape[1]
|
||||
return hidden_size, ffn_hidden_size
|
||||
elif storage == "module_list":
|
||||
# Legacy: individual expert modules
|
||||
if isinstance(experts_module, nn.ModuleList) and len(experts_module) > 0:
|
||||
expert0 = experts_module[0]
|
||||
w1 = getattr(expert0, preset.expert_w1, None)
|
||||
if w1 is None:
|
||||
# Try weight attribute for nn.Linear
|
||||
for name, child in expert0.named_children():
|
||||
if preset.expert_w1 in name:
|
||||
w1 = child.weight if hasattr(child, 'weight') else None
|
||||
break
|
||||
if w1 is not None:
|
||||
if isinstance(w1, nn.Linear):
|
||||
return w1.in_features, w1.out_features
|
||||
elif isinstance(w1, (nn.Parameter, torch.Tensor)):
|
||||
w1_shape = _source_param_shape(w1)
|
||||
if len(w1_shape) == 2:
|
||||
return w1_shape[1], w1_shape[0]
|
||||
|
||||
raise ValueError(f"Could not infer hidden_size/ffn_hidden_size from experts module "
|
||||
f"with storage={storage}, preset.expert_w1={preset.expert_w1}")
|
||||
|
||||
|
||||
def _detect_forward_contract(
|
||||
moe_module: nn.Module,
|
||||
router_module: nn.Module,
|
||||
) -> ForwardContract:
|
||||
"""Detect the forward contract for router logits capture.
|
||||
|
||||
Returns:
|
||||
ForwardContract with router-logit return and capture metadata.
|
||||
"""
|
||||
# Check for OutputRecorder on the model (transformers 5.0.0 pattern)
|
||||
# Look for _can_record_outputs attribute on parent modules
|
||||
capture_target: Literal["moe_block", "router", "none"] = "none"
|
||||
capture_index: int | None = None
|
||||
capture_layer_name: str | None = None
|
||||
return_router_logits = False
|
||||
|
||||
# Check for OutputRecorder pattern on router class
|
||||
router_class = type(router_module)
|
||||
if hasattr(router_class, '_can_record_outputs'):
|
||||
capture_target = "router"
|
||||
record_config = router_class._can_record_outputs
|
||||
if isinstance(record_config, dict):
|
||||
for key, val in record_config.items():
|
||||
if isinstance(val, dict):
|
||||
capture_index = val.get('index', 0)
|
||||
capture_layer_name = val.get('layer_name', None)
|
||||
else:
|
||||
capture_index = 0
|
||||
elif isinstance(record_config, (list, tuple)):
|
||||
capture_index = 0
|
||||
logger.debug(f"Detected OutputRecorder on router class {router_class.__name__}: "
|
||||
f"index={capture_index}, layer_name={capture_layer_name}")
|
||||
|
||||
# Check if MoE block has tuple return contract (legacy transformers)
|
||||
if hasattr(moe_module, '_can_record_outputs'):
|
||||
record_config = moe_module._can_record_outputs
|
||||
if record_config:
|
||||
capture_target = "moe_block"
|
||||
return_router_logits = True
|
||||
if isinstance(record_config, dict):
|
||||
for key, val in record_config.items():
|
||||
if isinstance(val, dict):
|
||||
capture_index = val.get('index', None)
|
||||
elif isinstance(val, int):
|
||||
capture_index = val
|
||||
|
||||
return ForwardContract(
|
||||
return_router_logits=return_router_logits,
|
||||
capture_target=capture_target,
|
||||
capture_index=capture_index,
|
||||
capture_layer_name=capture_layer_name,
|
||||
)
|
||||
|
||||
|
||||
class AutoEP:
|
||||
"""Automatic Expert Parallelism: detect and replace MoE layers."""
|
||||
|
||||
def __init__(self, model: nn.Module, config: AutoEPConfig) -> None:
|
||||
self.model = model
|
||||
self.config = config
|
||||
self.model_config = getattr(model, 'config', None)
|
||||
self._retargeted_transformers_output_recorders: set[str] = set()
|
||||
fill_autoep_config_from_hf(self.config, self.model_config)
|
||||
|
||||
def ep_parser(self) -> list[MoELayerSpec]:
|
||||
"""Traverse model and detect MoE layers. Returns list of MoELayerSpec."""
|
||||
specs = []
|
||||
|
||||
# Determine which preset(s) to use
|
||||
presets_to_try = self._resolve_presets()
|
||||
|
||||
for preset_name, preset in presets_to_try:
|
||||
adapter = get_preset_adapter(preset.preset_adapter)
|
||||
pattern = re.compile(preset.moe_layer_pattern)
|
||||
|
||||
for module_name, module in self.model.named_modules():
|
||||
if not pattern.fullmatch(module_name):
|
||||
continue
|
||||
|
||||
# Structural validation: check for experts child
|
||||
experts_child = getattr(module, preset.experts_pattern, None)
|
||||
if experts_child is None:
|
||||
logger.debug(
|
||||
"Skipping %s: pattern matched but no '%s' child (likely dense FFN)",
|
||||
module_name,
|
||||
preset.experts_pattern,
|
||||
)
|
||||
continue
|
||||
|
||||
expert_layout = adapter.resolve_expert_layout(experts_child, preset)
|
||||
|
||||
# Accept both: nn.ModuleList (legacy) and Experts class (transformers 5.0.0+)
|
||||
has_expert_params = (isinstance(experts_child, nn.ModuleList)
|
||||
or _has_3d_expert_params(experts_child, expert_layout))
|
||||
if not has_expert_params:
|
||||
logger.debug(
|
||||
"Skipping %s: '%s' child exists but has no expert parameters",
|
||||
module_name,
|
||||
preset.experts_pattern,
|
||||
)
|
||||
continue
|
||||
|
||||
# Check for router
|
||||
router_child = getattr(module, preset.router_pattern, None)
|
||||
if router_child is None:
|
||||
logger.debug(
|
||||
"Skipping %s: no router child '%s'",
|
||||
module_name,
|
||||
preset.router_pattern,
|
||||
)
|
||||
continue
|
||||
|
||||
# Detect storage format
|
||||
storage = _detect_expert_storage(experts_child, expert_layout)
|
||||
|
||||
# Get num_experts and top_k from config or weights
|
||||
num_experts = None
|
||||
top_k = None
|
||||
|
||||
if self.model_config is not None:
|
||||
num_experts = _get_num_experts_from_config(self.model_config, preset)
|
||||
top_k = _get_top_k_from_config(self.model_config, preset)
|
||||
|
||||
# Validate/derive from router weight shape
|
||||
router_weight = getattr(router_child, 'weight', None)
|
||||
router_weight_shape = _source_param_shape(router_weight) if router_weight is not None else None
|
||||
if router_weight_shape is not None and len(router_weight_shape) == 2:
|
||||
num_experts_from_weight = router_weight_shape[0]
|
||||
hidden_from_weight = router_weight_shape[1]
|
||||
if num_experts is not None and num_experts != num_experts_from_weight:
|
||||
raise ValueError(f"Config num_experts={num_experts} mismatches router weight "
|
||||
f"shape {router_weight_shape} (expected {num_experts_from_weight}) "
|
||||
f"in layer '{module_name}'")
|
||||
num_experts = num_experts_from_weight
|
||||
|
||||
if num_experts is None:
|
||||
raise ValueError(f"Could not determine num_experts for layer '{module_name}'. "
|
||||
f"Set model.config.{preset.num_experts_attr} or use a preset.")
|
||||
|
||||
# Override top_k from config if user specified
|
||||
if isinstance(self.config.top_k, int):
|
||||
top_k = self.config.top_k
|
||||
elif top_k is None:
|
||||
raise ValueError(f"Could not determine top_k for layer '{module_name}'. "
|
||||
f"Set model.config.{preset.top_k_attr} or config top_k.")
|
||||
|
||||
# Infer hidden sizes
|
||||
try:
|
||||
hidden_size, ffn_hidden_size = _infer_hidden_and_ffn_size(experts_child, expert_layout, storage,
|
||||
num_experts)
|
||||
except ValueError as e:
|
||||
if self._requires_selected_preset_detection():
|
||||
raise ValueError(f"AutoEP: preset '{preset_name}' matched layer '{module_name}' "
|
||||
f"with router and experts, but shape inference failed: {e}") from e
|
||||
logger.warning(f"Skipping {module_name}: {e}")
|
||||
continue
|
||||
|
||||
# Cross-validate hidden_size with router
|
||||
if router_weight_shape is not None and len(router_weight_shape) == 2:
|
||||
if hidden_size != router_weight_shape[1]:
|
||||
raise ValueError(f"hidden_size={hidden_size} from expert weights mismatches "
|
||||
f"router weight dim={router_weight_shape[1]} in '{module_name}'")
|
||||
|
||||
# Validate top_k <= num_experts
|
||||
if top_k > num_experts:
|
||||
raise ValueError(f"top_k={top_k} exceeds num_experts={num_experts} "
|
||||
f"in layer '{module_name}'")
|
||||
|
||||
# Resolve score_func
|
||||
if self.config.score_func != "auto":
|
||||
score_func = self.config.score_func
|
||||
else:
|
||||
# Check model config for scoring_func attribute
|
||||
cfg_score = getattr(self.model_config, 'scoring_func', None)
|
||||
if cfg_score in ("softmax", "sigmoid"):
|
||||
score_func = cfg_score
|
||||
else:
|
||||
score_func = preset.score_func
|
||||
|
||||
# Resolve score_apply
|
||||
if self.config.score_apply != "auto":
|
||||
score_apply = self.config.score_apply
|
||||
else:
|
||||
score_apply = preset.score_apply
|
||||
|
||||
route_norm = adapter.resolve_route_norm(self.config, preset, self.model_config)
|
||||
|
||||
route_scale = _resolve_route_scale(self.config, self.model_config)
|
||||
|
||||
group_routing = adapter.resolve_group_routing(self.config, self.model_config)
|
||||
|
||||
# Check gate bias
|
||||
gate_bias = preset.gate_bias
|
||||
if router_weight is not None:
|
||||
gate_bias = getattr(router_child, 'bias', None) is not None
|
||||
|
||||
forward_contract = adapter.adjust_forward_contract(_detect_forward_contract(module, router_child))
|
||||
|
||||
# Check shared experts
|
||||
has_shared = False
|
||||
shared_name = ""
|
||||
shared_gate_name = ""
|
||||
if preset.has_shared_experts and preset.shared_experts_pattern:
|
||||
shared = getattr(module, preset.shared_experts_pattern, None)
|
||||
if shared is not None:
|
||||
has_shared = True
|
||||
shared_name = preset.shared_experts_pattern
|
||||
if preset.shared_experts_gate_pattern:
|
||||
shared_gate = getattr(module, preset.shared_experts_gate_pattern, None)
|
||||
if shared_gate is not None:
|
||||
shared_gate_name = preset.shared_experts_gate_pattern
|
||||
|
||||
# Warn about router stochasticity/precision settings
|
||||
if self.model_config is not None:
|
||||
jitter = getattr(self.model_config, 'router_jitter_noise', 0.0)
|
||||
if jitter and jitter > 0:
|
||||
logger.warning(f"Layer {module_name}: model has router_jitter_noise={jitter}, "
|
||||
f"AutoEP router does not implement jitter.")
|
||||
z_loss = getattr(self.model_config, 'router_z_loss_coef', 0.0)
|
||||
if z_loss and z_loss > 0:
|
||||
logger.warning(f"Layer {module_name}: model has router_z_loss_coef={z_loss}, "
|
||||
f"AutoEP router does not implement z-loss.")
|
||||
|
||||
spec = MoELayerSpec(
|
||||
moe_module_name=module_name,
|
||||
model_family=preset_name,
|
||||
router_name=preset.router_pattern,
|
||||
experts_name=preset.experts_pattern,
|
||||
expert_storage=storage,
|
||||
expert_w1_name=expert_layout.expert_w1,
|
||||
expert_w2_name=expert_layout.expert_w2,
|
||||
expert_w3_name=expert_layout.expert_w3,
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
ffn_hidden_size=ffn_hidden_size,
|
||||
score_func=score_func,
|
||||
score_apply=score_apply,
|
||||
route_norm=route_norm,
|
||||
gate_bias=gate_bias,
|
||||
return_router_logits=forward_contract.return_router_logits,
|
||||
router_logits_capture_target=forward_contract.capture_target,
|
||||
router_logits_capture_index=forward_contract.capture_index,
|
||||
router_logits_capture_layer_name=forward_contract.capture_layer_name,
|
||||
has_shared_experts=has_shared,
|
||||
shared_experts_name=shared_name,
|
||||
shared_experts_gate_name=shared_gate_name,
|
||||
route_scale=route_scale,
|
||||
num_expert_groups=group_routing.num_expert_groups,
|
||||
num_limited_groups=group_routing.num_limited_groups,
|
||||
group_score_func=group_routing.group_score_func,
|
||||
supports_expert_bias=preset.supports_expert_bias,
|
||||
unsupported_router_bias_names=preset.unsupported_router_bias_names,
|
||||
preset_adapter=preset.preset_adapter,
|
||||
router_logits_capture_mode=forward_contract.router_logits_capture_mode,
|
||||
moe_output_shape=forward_contract.moe_output_shape,
|
||||
)
|
||||
specs.append(spec)
|
||||
logger.debug(f"Detected MoE layer: {module_name} (family={preset_name}, "
|
||||
f"experts={num_experts}, top_k={top_k}, storage={storage})")
|
||||
|
||||
if not specs:
|
||||
if self._requires_selected_preset_detection():
|
||||
self._raise_no_moe_layers_detected(presets_to_try)
|
||||
logger.warning("AutoEP: no MoE layers detected in model.")
|
||||
else:
|
||||
_raise_if_duplicate_moe_specs(specs)
|
||||
|
||||
return specs
|
||||
|
||||
def _replace_moe_layer_without_retarget(
|
||||
self,
|
||||
spec: MoELayerSpec,
|
||||
ep_size: int,
|
||||
ep_rank: int,
|
||||
) -> nn.Module:
|
||||
from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
|
||||
|
||||
# Navigate to the parent module and get the child name
|
||||
parts = spec.moe_module_name.split(".")
|
||||
parent = self.model
|
||||
for part in parts[:-1]:
|
||||
parent = getattr(parent, part)
|
||||
child_name = parts[-1]
|
||||
source_module = getattr(parent, child_name)
|
||||
|
||||
# Create replacement layer
|
||||
replacement = AutoEPMoELayer(
|
||||
spec=spec,
|
||||
source_module=source_module,
|
||||
ep_size=ep_size,
|
||||
ep_rank=ep_rank,
|
||||
config=self.config,
|
||||
)
|
||||
|
||||
# Replace in-place on parent
|
||||
setattr(parent, child_name, replacement)
|
||||
return replacement
|
||||
|
||||
def _retarget_transformers_output_recorders(self, spec: MoELayerSpec, replacement: nn.Module) -> None:
|
||||
adapter = get_preset_adapter(spec.preset_adapter)
|
||||
adapter.retarget_transformers_output_recorders(
|
||||
self.model,
|
||||
spec,
|
||||
replacement,
|
||||
self._retargeted_transformers_output_recorders,
|
||||
_remove_transformers_output_capture_hooks,
|
||||
)
|
||||
|
||||
def replace_moe_layer(
|
||||
self,
|
||||
spec: MoELayerSpec,
|
||||
ep_size: int,
|
||||
ep_rank: int,
|
||||
) -> None:
|
||||
"""Replace a single MoE module with AutoEPMoELayer in-place on the model."""
|
||||
replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank)
|
||||
self._retarget_transformers_output_recorders(spec, replacement)
|
||||
|
||||
logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer "
|
||||
f"(ep_size={ep_size}, ep_rank={ep_rank}, "
|
||||
f"local_experts={replacement.num_local_experts})")
|
||||
|
||||
def replace_moe_layers(
|
||||
self,
|
||||
specs: list[MoELayerSpec],
|
||||
ep_size: int,
|
||||
ep_rank: int,
|
||||
) -> None:
|
||||
"""Replace multiple MoE modules and batch post-replacement recorder retargeting."""
|
||||
replacements: list[tuple[MoELayerSpec, nn.Module]] = []
|
||||
for spec in specs:
|
||||
replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank)
|
||||
replacements.append((spec, replacement))
|
||||
logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer "
|
||||
f"(ep_size={ep_size}, ep_rank={ep_rank}, "
|
||||
f"local_experts={replacement.num_local_experts})")
|
||||
|
||||
retarget_groups: OrderedDict[tuple[str, str, type], tuple[MoELayerSpec, nn.Module]] = OrderedDict()
|
||||
for spec, replacement in replacements:
|
||||
retarget_key = (spec.preset_adapter, spec.model_family, replacement.__class__)
|
||||
retarget_groups.setdefault(retarget_key, (spec, replacement))
|
||||
|
||||
for spec, replacement in retarget_groups.values():
|
||||
self._retarget_transformers_output_recorders(spec, replacement)
|
||||
|
||||
def _apply_config_overrides(self, preset: MoEModelPreset) -> MoEModelPreset:
|
||||
return apply_config_overrides(self.config, preset)
|
||||
|
||||
def _requires_selected_preset_detection(self) -> bool:
|
||||
"""Return whether empty detection should fail for the selected preset."""
|
||||
if self.config.preset_model is not None:
|
||||
return True
|
||||
if self.config.moe_layer_pattern is not None:
|
||||
return True
|
||||
if self.model_config is None:
|
||||
return False
|
||||
model_type = getattr(self.model_config, 'model_type', None)
|
||||
return _is_known_hf_model_type(model_type)
|
||||
|
||||
def _raise_no_moe_layers_detected(self, presets_to_try: list[tuple[str, MoEModelPreset]]) -> None:
|
||||
model_type = getattr(self.model_config, 'model_type', None)
|
||||
if self.config.preset_model is not None:
|
||||
source = f"preset_model='{self.config.preset_model}'"
|
||||
elif self.config.moe_layer_pattern is not None:
|
||||
source = f"moe_layer_pattern='{self.config.moe_layer_pattern}'"
|
||||
else:
|
||||
source = f"model_type='{model_type}'"
|
||||
|
||||
expected = "; ".join(f"{preset_name}: moe_layer_pattern='{preset.moe_layer_pattern}', "
|
||||
f"router='{preset.router_pattern}', experts='{preset.experts_pattern}'"
|
||||
for preset_name, preset in presets_to_try)
|
||||
raise ValueError(f"AutoEP: no MoE layers detected for {source}. "
|
||||
f"Expected MoE structure for selected preset(s): {expected}. "
|
||||
"This usually means the selected preset does not match the model implementation, "
|
||||
"or the installed Transformers version exposes a different structure. Choose a matching "
|
||||
"preset, upgrade Transformers, or provide custom AutoEP patterns.")
|
||||
|
||||
def _resolve_presets(self) -> list[tuple[str, MoEModelPreset]]:
|
||||
"""Determine which preset(s) to use for detection."""
|
||||
return resolve_preset_candidates(self.config, self.model_config)
|
||||
@@ -0,0 +1,348 @@
|
||||
# Copyright (c) DeepSpeed Team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""AutoEP configuration: config parsing, model presets, and validation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import (
|
||||
_UNSET,
|
||||
_raise_unsupported_load_balance_coeff,
|
||||
AutoEPConfig,
|
||||
MoELayerSpec,
|
||||
MoEModelPreset,
|
||||
)
|
||||
from deepspeed.module_inject.auto_ep_presets.registry import (
|
||||
PRESET_MODELS,
|
||||
available_preset_names,
|
||||
resolve_autoep_config_defaults,
|
||||
)
|
||||
from deepspeed.module_inject.auto_ep_folding import build_folding_spec, validate_folding_global
|
||||
from deepspeed.utils import logger
|
||||
|
||||
__all__ = [
|
||||
"_UNSET",
|
||||
"AutoEPConfig",
|
||||
"MoELayerSpec",
|
||||
"MoEModelPreset",
|
||||
"PRESET_MODELS",
|
||||
"parse_autoep_config",
|
||||
"resolve_autoep_config_defaults",
|
||||
"validate_autoep_config",
|
||||
"validate_autoep_post_detection",
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_autoep_config(param_dict: dict) -> AutoEPConfig:
|
||||
"""Parse the 'expert_parallel' section from DS config JSON."""
|
||||
if not param_dict:
|
||||
return AutoEPConfig()
|
||||
|
||||
config = AutoEPConfig()
|
||||
config.enabled = param_dict.get("enabled", False)
|
||||
config.autoep_size = param_dict.get("autoep_size", 1)
|
||||
config.expert_tensor_parallel_size = param_dict.get("expert_tensor_parallel_size", 1)
|
||||
config.validate_folding_routing = param_dict.get("validate_folding_routing", False)
|
||||
config.preset_model = param_dict.get("preset_model", None)
|
||||
config.moe_layer_pattern = param_dict.get("moe_layer_pattern", None)
|
||||
config.expert_pattern = param_dict.get("expert_pattern", None)
|
||||
config.router_pattern = param_dict.get("router_pattern", None)
|
||||
config.use_grouped_mm = param_dict.get("use_grouped_mm", True)
|
||||
config.route_norm = param_dict.get("route_norm", None)
|
||||
config.route_scale = param_dict.get("route_scale", 1.0)
|
||||
config.score_apply = param_dict.get("score_apply", "auto")
|
||||
config.combine_impl = param_dict.get("combine_impl", "auto")
|
||||
config.num_expert_groups = param_dict.get("num_expert_groups", None)
|
||||
config.num_limited_groups = param_dict.get("num_limited_groups", None)
|
||||
config.score_func = param_dict.get("score_func", "auto")
|
||||
config.top_k = param_dict.get("top_k", "auto")
|
||||
if "load_balance_coeff" in param_dict:
|
||||
value = param_dict["load_balance_coeff"]
|
||||
if value is not None:
|
||||
_raise_unsupported_load_balance_coeff(value)
|
||||
config.load_balance_coeff = None
|
||||
config._load_balance_coeff_explicit = True
|
||||
else:
|
||||
config.load_balance_coeff = None
|
||||
config._load_balance_coeff_explicit = False
|
||||
config.routed_scaling_factor = param_dict.get("routed_scaling_factor", "auto")
|
||||
config.expert_w1 = param_dict.get("expert_w1", None)
|
||||
config.expert_w2 = param_dict.get("expert_w2", None)
|
||||
# expert_w3: key absent → _UNSET (preset default); key present with null → None (fused); key present with string → custom name
|
||||
if "expert_w3" in param_dict:
|
||||
config.expert_w3 = param_dict["expert_w3"] # None or string
|
||||
else:
|
||||
config.expert_w3 = _UNSET
|
||||
config.num_experts_attr = param_dict.get("num_experts_attr", None)
|
||||
config.top_k_attr = param_dict.get("top_k_attr", None)
|
||||
config.has_shared_experts = param_dict.get("has_shared_experts", None)
|
||||
config.shared_experts_pattern = param_dict.get("shared_experts_pattern", None)
|
||||
config.shared_experts_gate_pattern = param_dict.get("shared_experts_gate_pattern", None)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Validation helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def validate_autoep_config(
|
||||
config: AutoEPConfig,
|
||||
world_size: int,
|
||||
pp_size: int,
|
||||
tp_size: int,
|
||||
sp_size: int,
|
||||
*,
|
||||
zero_stage: int = 0,
|
||||
deepcompile_enabled: bool = False,
|
||||
tp_preset_model: str | None = None,
|
||||
use_data_before_expert_parallel: bool = False,
|
||||
mpu=None,
|
||||
zero_offload_optimizer: bool = False,
|
||||
zero_offload_param: bool = False,
|
||||
) -> None:
|
||||
"""Validate config constraints. Raises ValueError on invalid config."""
|
||||
if config.load_balance_coeff is not None:
|
||||
_raise_unsupported_load_balance_coeff(config.load_balance_coeff)
|
||||
|
||||
if not isinstance(config.validate_folding_routing, bool):
|
||||
raise ValueError("expert_parallel.validate_folding_routing must be a boolean")
|
||||
|
||||
if not config.enabled:
|
||||
return
|
||||
|
||||
folding_spec = build_folding_spec(
|
||||
world_size=world_size,
|
||||
pp_size=pp_size,
|
||||
tp_size=max(tp_size, 1),
|
||||
ep_size=config.autoep_size,
|
||||
etp_size=config.expert_tensor_parallel_size,
|
||||
mp_mode="tp" if tp_size > 1 else "sp",
|
||||
)
|
||||
validate_folding_global(
|
||||
folding_spec,
|
||||
zero_stage=zero_stage,
|
||||
sp_size=sp_size,
|
||||
deepcompile_enabled=deepcompile_enabled,
|
||||
use_data_before_expert_parallel=use_data_before_expert_parallel,
|
||||
mpu=mpu,
|
||||
autoep_enabled=config.enabled,
|
||||
tp_preset=tp_preset_model,
|
||||
ep_preset=config.preset_model,
|
||||
zero_offload_optimizer=zero_offload_optimizer,
|
||||
zero_offload_param=zero_offload_param,
|
||||
)
|
||||
|
||||
# Validate preset_model if specified
|
||||
if config.preset_model is not None and config.preset_model not in PRESET_MODELS:
|
||||
raise ValueError(f"Unknown preset_model '{config.preset_model}'. "
|
||||
f"Available presets: {list(available_preset_names())}")
|
||||
|
||||
# Validate score_apply
|
||||
valid_score_apply = ("auto", "pre", "post")
|
||||
if config.score_apply not in valid_score_apply:
|
||||
raise ValueError(f"score_apply must be one of {valid_score_apply}, "
|
||||
f"got '{config.score_apply}'")
|
||||
|
||||
# Validate combine_impl
|
||||
valid_combine_impl = ("auto", "weighted_sum", "legacy_bmm")
|
||||
if config.combine_impl not in valid_combine_impl:
|
||||
raise ValueError(f"combine_impl must be one of {valid_combine_impl}, "
|
||||
f"got '{config.combine_impl}'")
|
||||
|
||||
# Validate score_func
|
||||
valid_score_func = ("auto", "softmax", "sigmoid")
|
||||
if config.score_func not in valid_score_func:
|
||||
raise ValueError(f"score_func must be one of {valid_score_func}, "
|
||||
f"got '{config.score_func}'")
|
||||
|
||||
# Validate group-limited routing constraints
|
||||
if config.num_limited_groups is not None:
|
||||
if config.num_limited_groups < 1:
|
||||
raise ValueError(f"num_limited_groups must be >= 1, got {config.num_limited_groups}")
|
||||
|
||||
if config.num_expert_groups is not None:
|
||||
if config.num_expert_groups < 1:
|
||||
raise ValueError(f"num_expert_groups must be >= 1, got {config.num_expert_groups}")
|
||||
if config.num_limited_groups is not None and config.num_limited_groups > config.num_expert_groups:
|
||||
raise ValueError(f"num_limited_groups ({config.num_limited_groups}) must be <= "
|
||||
f"num_expert_groups ({config.num_expert_groups})")
|
||||
logger.warning("num_expert_groups is set; interaction with EP topology "
|
||||
"is not yet optimized.")
|
||||
|
||||
# Warn if autoep_size == 1 (no EP needed)
|
||||
if config.autoep_size == 1:
|
||||
logger.warning("autoep_size=1 means every rank owns all experts with no AllToAll. "
|
||||
"AutoEP replacement remains enabled, but expert-parallel communication "
|
||||
"is bypassed because every rank owns all experts.")
|
||||
|
||||
# Helper validators (local to validate_autoep_config)
|
||||
def _validate_attr_name(field_name: str, value, *, allow_dot: bool = False) -> None:
|
||||
if value is None:
|
||||
return
|
||||
if not isinstance(value, str) or value == "":
|
||||
raise ValueError(f"{field_name} must be a non-empty string")
|
||||
if not allow_dot and "." in value:
|
||||
raise ValueError(f"{field_name} must be a direct attribute name (no dots)")
|
||||
|
||||
# Validate expert weight names
|
||||
_validate_attr_name("expert_w1", config.expert_w1)
|
||||
_validate_attr_name("expert_w2", config.expert_w2)
|
||||
if config.expert_w3 is not _UNSET and config.expert_w3 is not None:
|
||||
_validate_attr_name("expert_w3", config.expert_w3)
|
||||
|
||||
# Validate model.config attribute names
|
||||
_validate_attr_name("num_experts_attr", config.num_experts_attr)
|
||||
_validate_attr_name("top_k_attr", config.top_k_attr)
|
||||
|
||||
# Validate child-name fields (direct attribute names, not regex/path)
|
||||
_validate_attr_name("router_pattern", config.router_pattern)
|
||||
_validate_attr_name("expert_pattern", config.expert_pattern)
|
||||
_validate_attr_name("shared_experts_pattern", config.shared_experts_pattern)
|
||||
_validate_attr_name("shared_experts_gate_pattern", config.shared_experts_gate_pattern)
|
||||
|
||||
# Validate has_shared_experts type
|
||||
if config.has_shared_experts is not None and not isinstance(config.has_shared_experts, bool):
|
||||
raise ValueError("has_shared_experts must be a boolean when set")
|
||||
|
||||
# Warn if explicit top_k overrides top_k_attr
|
||||
if isinstance(config.top_k, int) and config.top_k_attr is not None:
|
||||
logger.warning("top_k is explicitly set; top_k_attr will be ignored.")
|
||||
|
||||
if config.routed_scaling_factor != "auto" and not isinstance(config.routed_scaling_factor, (int, float)):
|
||||
raise ValueError("routed_scaling_factor must be a number or 'auto'")
|
||||
|
||||
# Validate shared expert field pairing
|
||||
if config.has_shared_experts is True and not config.shared_experts_pattern:
|
||||
logger.warning("has_shared_experts=True but shared_experts_pattern is not set. "
|
||||
"Shared expert detection requires both fields.")
|
||||
if config.shared_experts_pattern and config.has_shared_experts is not True:
|
||||
logger.warning(f"shared_experts_pattern='{config.shared_experts_pattern}' is set "
|
||||
f"but has_shared_experts is not True. Pattern will be ignored.")
|
||||
if config.shared_experts_gate_pattern and config.has_shared_experts is not True:
|
||||
logger.warning(f"shared_experts_gate_pattern='{config.shared_experts_gate_pattern}' is set "
|
||||
f"but has_shared_experts is not True. Pattern will be ignored.")
|
||||
|
||||
# Warn if custom override fields are set alongside preset_model or auto-detect
|
||||
custom_fields_set = []
|
||||
if config.moe_layer_pattern is not None:
|
||||
custom_fields_set.append("moe_layer_pattern")
|
||||
if config.router_pattern is not None:
|
||||
custom_fields_set.append("router_pattern")
|
||||
if config.expert_pattern is not None:
|
||||
custom_fields_set.append("expert_pattern")
|
||||
if config.expert_w1 is not None:
|
||||
custom_fields_set.append("expert_w1")
|
||||
if config.expert_w2 is not None:
|
||||
custom_fields_set.append("expert_w2")
|
||||
if config.expert_w3 is not _UNSET:
|
||||
custom_fields_set.append("expert_w3")
|
||||
if config.num_experts_attr is not None:
|
||||
custom_fields_set.append("num_experts_attr")
|
||||
if config.top_k_attr is not None:
|
||||
custom_fields_set.append("top_k_attr")
|
||||
if config.has_shared_experts is not None:
|
||||
custom_fields_set.append("has_shared_experts")
|
||||
if config.shared_experts_pattern is not None:
|
||||
custom_fields_set.append("shared_experts_pattern")
|
||||
if config.shared_experts_gate_pattern is not None:
|
||||
custom_fields_set.append("shared_experts_gate_pattern")
|
||||
if custom_fields_set and config.preset_model is not None:
|
||||
logger.warning(f"Custom preset fields {custom_fields_set} are set alongside "
|
||||
f"preset_model='{config.preset_model}'. Custom fields will override "
|
||||
f"preset defaults during detection.")
|
||||
if custom_fields_set and config.preset_model is None and config.moe_layer_pattern is None:
|
||||
logger.warning(f"Custom preset fields {custom_fields_set} are set without preset_model or "
|
||||
f"moe_layer_pattern. Overrides will apply to auto-detected presets or try-all.")
|
||||
|
||||
|
||||
def validate_autoep_post_detection(
|
||||
config: AutoEPConfig,
|
||||
specs: list[MoELayerSpec],
|
||||
) -> None:
|
||||
"""Post-detection validation: ep_size vs num_experts constraints."""
|
||||
if not config.enabled or not specs:
|
||||
return
|
||||
|
||||
for spec in specs:
|
||||
# ep_size must not exceed num_experts
|
||||
if config.autoep_size > spec.num_experts:
|
||||
valid_divisors = _divisors(spec.num_experts)
|
||||
raise ValueError(f"autoep_size={config.autoep_size} exceeds num_experts="
|
||||
f"{spec.num_experts} in layer '{spec.moe_module_name}'. "
|
||||
f"Each rank must own at least one expert. "
|
||||
f"Valid autoep_size values (divisors of {spec.num_experts}): "
|
||||
f"{valid_divisors}")
|
||||
|
||||
# num_experts must be divisible by ep_size
|
||||
if spec.num_experts % config.autoep_size != 0:
|
||||
valid_sizes = [d for d in _divisors(spec.num_experts) if d <= spec.num_experts]
|
||||
raise ValueError(f"num_experts={spec.num_experts} in layer "
|
||||
f"'{spec.moe_module_name}' is not divisible by "
|
||||
f"autoep_size={config.autoep_size}. "
|
||||
f"Suggested autoep_size values: {valid_sizes}")
|
||||
|
||||
num_expert_groups = spec.num_expert_groups if spec.num_expert_groups is not None else config.num_expert_groups
|
||||
num_limited_groups = spec.num_limited_groups if spec.num_limited_groups is not None else config.num_limited_groups
|
||||
|
||||
# Validate group-limited routing constraints after layer-specific defaults.
|
||||
if num_limited_groups is not None and num_expert_groups is None:
|
||||
raise ValueError(f"num_limited_groups requires num_expert_groups to be set "
|
||||
f"in layer '{spec.moe_module_name}'")
|
||||
|
||||
if num_expert_groups is not None:
|
||||
if num_expert_groups < 1:
|
||||
raise ValueError(f"num_expert_groups must be >= 1 in layer '{spec.moe_module_name}', "
|
||||
f"got {num_expert_groups}")
|
||||
if spec.num_experts % num_expert_groups != 0:
|
||||
raise ValueError(f"num_expert_groups ({num_expert_groups}) must divide "
|
||||
f"num_experts ({spec.num_experts}) in layer "
|
||||
f"'{spec.moe_module_name}'")
|
||||
if num_limited_groups is None:
|
||||
raise ValueError(f"num_limited_groups must be set when num_expert_groups is set "
|
||||
f"in layer '{spec.moe_module_name}'")
|
||||
if num_limited_groups < 1:
|
||||
raise ValueError(f"num_limited_groups must be >= 1 in layer '{spec.moe_module_name}', "
|
||||
f"got {num_limited_groups}")
|
||||
if num_limited_groups > num_expert_groups:
|
||||
raise ValueError(f"num_limited_groups ({num_limited_groups}) must be <= "
|
||||
f"num_expert_groups ({num_expert_groups}) in layer "
|
||||
f"'{spec.moe_module_name}'")
|
||||
|
||||
|
||||
def _divisors(n: int) -> list[int]:
|
||||
"""Return sorted list of positive divisors of n."""
|
||||
divs = []
|
||||
for i in range(1, int(n**0.5) + 1):
|
||||
if n % i == 0:
|
||||
divs.append(i)
|
||||
if i != n // i:
|
||||
divs.append(n // i)
|
||||
return sorted(divs)
|
||||
|
||||
|
||||
def fill_autoep_config_from_hf(config: AutoEPConfig, model_config) -> None:
|
||||
"""Back-fill AutoEPConfig fields from HF model config when user hasn't set them.
|
||||
|
||||
HF field names (e.g. n_group, topk_group, routed_scaling_factor) differ from
|
||||
AutoEP's internal names, so we map them explicitly rather than relying on the
|
||||
user to duplicate these values in the DS config JSON.
|
||||
"""
|
||||
if model_config is None:
|
||||
return
|
||||
# n_group / topk_group: DeepSeek-style node-limited routing groups
|
||||
if config.num_expert_groups is None:
|
||||
config.num_expert_groups = getattr(model_config, 'n_group', None)
|
||||
if config.num_limited_groups is None:
|
||||
config.num_limited_groups = getattr(model_config, 'topk_group', None)
|
||||
# routed_scaling_factor: sigmoid score scaling (DeepSeek-V3 / Moonlight)
|
||||
if config.routed_scaling_factor == "auto":
|
||||
hf_scale = getattr(model_config, 'routed_scaling_factor', None)
|
||||
if hf_scale is not None:
|
||||
config.route_scale = float(hf_scale)
|
||||
@@ -0,0 +1,513 @@
|
||||
# Copyright (c) DeepSpeed Team.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""AutoEP + AutoTP folding topology helpers.
|
||||
|
||||
The functions in this module are pure topology math unless a caller passes
|
||||
runtime process-group handles into :class:`FoldingGroupHandles`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
AUTOEP_FOLDING_PARAM_FAMILY_ATTR = "ds_autoep_folding_param_family"
|
||||
AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM = "router_gate_replicated"
|
||||
AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM = "router_gate_partial"
|
||||
AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM = "sp_sharded_layernorm"
|
||||
AUTOEP_FOLDING_GRAD_CORRECTED_ATTR = "ds_autoep_folding_grad_corrected"
|
||||
AUTOEP_FOLDING_GRAD_REDUCE_SKIP = "skip"
|
||||
AUTOEP_FOLDING_GRAD_REDUCE_SUM = "sum"
|
||||
AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE = "average"
|
||||
# Divide by tp_size with NO TP all_reduce. Used for routed-expert parameters: the
|
||||
# folded forward all-gathers expert outputs into a replicated full view in
|
||||
# ``restore_combined``, whose backward injects a ``tp_size`` factor (same factor the
|
||||
# replicated router cancels via AVERAGE). Routed experts are not TP-replicated, so
|
||||
# they must not be TP all_reduced; they only need that spurious ``tp_size`` factor
|
||||
# divided out. The remaining data-parallel reduction is owned by the expert-data
|
||||
# -parallel (EDP) path, and ``/tp_size`` is linear so it composes with that EDP
|
||||
# all_reduce in either order.
|
||||
AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL = "expert_tp_cancel"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParallelFoldingSpec:
|
||||
world_size: int
|
||||
pp_size: int
|
||||
stage_size: int
|
||||
tp_size: int
|
||||
dp_size: int
|
||||
ep_size: int
|
||||
etp_size: int
|
||||
edp_size: int
|
||||
mp_mode: str = "tp"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FoldingGroupTables:
|
||||
tp_groups: tuple[tuple[int, ...], ...]
|
||||
dense_dp_groups: tuple[tuple[int, ...], ...]
|
||||
ep_groups: tuple[tuple[int, ...], ...]
|
||||
edp_groups: tuple[tuple[int, ...], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FoldingGroupHandles:
|
||||
spec: ParallelFoldingSpec
|
||||
tp_group: object
|
||||
dense_dp_group: object
|
||||
ep_group: object
|
||||
edp_group: object
|
||||
ep_group_name: str
|
||||
tp_ranks: tuple[int, ...]
|
||||
dense_dp_ranks: tuple[int, ...]
|
||||
ep_ranks: tuple[int, ...]
|
||||
edp_ranks: tuple[int, ...]
|
||||
|
||||
|
||||
def _divisors(value: int) -> list[int]:
|
||||
return [candidate for candidate in range(1, value + 1) if value % candidate == 0]
|
||||
|
||||
|
||||
def _require_positive(name: str, value: int) -> None:
|
||||
if not isinstance(value, int) or value < 1:
|
||||
raise ValueError(f"{name} must be a positive integer, got {value!r}")
|
||||
|
||||
|
||||
def build_folding_spec(
|
||||
*,
|
||||
world_size: int,
|
||||
pp_size: int,
|
||||
tp_size: int,
|
||||
ep_size: int,
|
||||
etp_size: int = 1,
|
||||
mp_mode: str = "tp",
|
||||
) -> ParallelFoldingSpec:
|
||||
"""Build the immutable per-stage folding spec from public config sizes."""
|
||||
for name, value in (
|
||||
("world_size", world_size),
|
||||
("pp_size", pp_size),
|
||||
("tensor_parallel.autotp_size", tp_size),
|
||||
("expert_parallel.autoep_size", ep_size),
|
||||
("expert_parallel.expert_tensor_parallel_size", etp_size),
|
||||
):
|
||||
_require_positive(name, value)
|
||||
|
||||
if world_size % pp_size != 0:
|
||||
raise ValueError(f"pp_size={pp_size} must divide world_size={world_size}. "
|
||||
f"Valid pp_size values: {_divisors(world_size)}")
|
||||
|
||||
stage_size = world_size // pp_size
|
||||
if stage_size % tp_size != 0:
|
||||
raise ValueError(f"tensor_parallel.autotp_size={tp_size} must divide the stage size "
|
||||
f"(world_size={world_size} / pp_size={pp_size} = {stage_size}). "
|
||||
f"Computed dp would be non-integral. Valid autotp_size values: {_divisors(stage_size)}")
|
||||
|
||||
expert_width = ep_size * etp_size
|
||||
if stage_size % expert_width != 0:
|
||||
raise ValueError(f"expert_parallel.autoep_size * expert_parallel.expert_tensor_parallel_size "
|
||||
f"({ep_size} * {etp_size} = {expert_width}) must divide the stage size "
|
||||
f"(world_size={world_size} / pp_size={pp_size} = {stage_size}). "
|
||||
f"Computed edp would be non-integral. Valid expert-width values: {_divisors(stage_size)}")
|
||||
|
||||
return ParallelFoldingSpec(
|
||||
world_size=world_size,
|
||||
pp_size=pp_size,
|
||||
stage_size=stage_size,
|
||||
tp_size=tp_size,
|
||||
dp_size=stage_size // tp_size,
|
||||
ep_size=ep_size,
|
||||
etp_size=etp_size,
|
||||
edp_size=stage_size // expert_width,
|
||||
mp_mode=mp_mode,
|
||||
)
|
||||
|
||||
|
||||
def expected_folding_group_tables(spec: ParallelFoldingSpec) -> FoldingGroupTables:
|
||||
"""Derive TP, dense-DP, EP, and EDP rank tables without process groups."""
|
||||
tp_groups: list[tuple[int, ...]] = []
|
||||
dense_dp_groups: list[tuple[int, ...]] = []
|
||||
ep_groups: list[tuple[int, ...]] = []
|
||||
edp_groups: list[tuple[int, ...]] = []
|
||||
|
||||
for stage_start in range(0, spec.world_size, spec.stage_size):
|
||||
stage_ranks = list(range(stage_start, stage_start + spec.stage_size))
|
||||
|
||||
for dp_idx in range(spec.dp_size):
|
||||
start = dp_idx * spec.tp_size
|
||||
tp_groups.append(tuple(stage_ranks[start:start + spec.tp_size]))
|
||||
for tp_lane in range(spec.tp_size):
|
||||
dense_dp_groups.append(tuple(stage_ranks[tp_lane::spec.tp_size]))
|
||||
|
||||
local_ep_groups = [
|
||||
tuple(stage_ranks[start:start + spec.ep_size]) for start in range(0, len(stage_ranks), spec.ep_size)
|
||||
]
|
||||
ep_groups.extend(local_ep_groups)
|
||||
for pos in range(spec.ep_size):
|
||||
edp_groups.append(tuple(group[pos] for group in local_ep_groups))
|
||||
|
||||
return FoldingGroupTables(
|
||||
tp_groups=tuple(tp_groups),
|
||||
dense_dp_groups=tuple(dense_dp_groups),
|
||||
ep_groups=tuple(ep_groups),
|
||||
edp_groups=tuple(edp_groups),
|
||||
)
|
||||
|
||||
|
||||
def local_folding_ranks(global_rank: int, spec: ParallelFoldingSpec) -> dict[str, tuple[int, ...]]:
|
||||
tables = expected_folding_group_tables(spec)
|
||||
result = {}
|
||||
for name, groups in (
|
||||
("tp", tables.tp_groups),
|
||||
("dense_dp", tables.dense_dp_groups),
|
||||
("ep", tables.ep_groups),
|
||||
("edp", tables.edp_groups),
|
||||
):
|
||||
result[name] = next(group for group in groups if global_rank in group)
|
||||
return result
|
||||
|
||||
|
||||
def _mpu_world_size(mpu, *names: str) -> int | None:
|
||||
if mpu is None:
|
||||
return None
|
||||
for name in names:
|
||||
getter = getattr(mpu, name, None)
|
||||
if getter is not None:
|
||||
return getter()
|
||||
return None
|
||||
|
||||
|
||||
def validate_folding_global(
|
||||
spec: ParallelFoldingSpec,
|
||||
*,
|
||||
zero_stage: int = 0,
|
||||
sp_size: int = 1,
|
||||
deepcompile_enabled: bool = False,
|
||||
use_data_before_expert_parallel: bool = False,
|
||||
mpu=None,
|
||||
autoep_enabled: bool = True,
|
||||
tp_preset: str | None = None,
|
||||
ep_preset: str | None = None,
|
||||
zero_offload_optimizer: bool = False,
|
||||
zero_offload_param: bool = False,
|
||||
) -> None:
|
||||
"""Validate global folding policy before any process group is created."""
|
||||
if not autoep_enabled:
|
||||
return
|
||||
|
||||
if deepcompile_enabled and spec.tp_size > 1:
|
||||
raise ValueError("DeepCompile with AutoEP+AutoTP folding is not supported; "
|
||||
"disable compile.deepcompile or use non-folded AutoEP with tensor_parallel.autotp_size=1.")
|
||||
|
||||
if spec.tp_size > 1 and spec.pp_size > 1:
|
||||
raise ValueError("AutoEP+AutoTP folding currently supports pp_size=1 only; "
|
||||
f"got pp_size={spec.pp_size}. Pipeline-parallel validation is planned separately.")
|
||||
|
||||
if spec.tp_size > 1 and sp_size > 1:
|
||||
raise ValueError("tensor_parallel.autotp_size and Ulysses sequence parallelism are mutually exclusive "
|
||||
f"for AutoEP folding (autotp_size={spec.tp_size}, sp_size={sp_size}).")
|
||||
|
||||
if spec.etp_size != 1:
|
||||
raise ValueError(f"expert_parallel.expert_tensor_parallel_size={spec.etp_size} is reserved for "
|
||||
"expert-internal tensor parallelism and is not supported yet. Use "
|
||||
"expert_tensor_parallel_size=1; ETP support is planned as follow-up work.")
|
||||
|
||||
# Cross-lane expert parallelism (expert_width = ep * etp need NOT be a subset of
|
||||
# the dense data-parallel size) is supported: ``expected_folding_group_tables``
|
||||
# lays EP groups across consecutive stage ranks while dense DP remains TP-lane
|
||||
# strided, so an EP group may span TP lanes and dense-DP ranks while preserving
|
||||
# node-local EP groups under node-contiguous rank mappings. The only structural
|
||||
# requirement is that the expert width tiles the stage cleanly, which
|
||||
# ``build_folding_spec`` already enforces (``stage_size % expert_width == 0``,
|
||||
# so ``edp`` is integral). The gradient convention holds across the pool
|
||||
# because each family's reduction is keyed to its replication structure, not
|
||||
# the EP layout: router/gate and dense/LayerNorm AVERAGE over the TP
|
||||
# (token-replication) group; routed experts cancel the restore ``tp_size``
|
||||
# factor (EXPERT_TP_CANCEL) and reduce data-parallel over
|
||||
# the EDP group. The earlier ``expert_width <= dp`` / ``dp % expert_width == 0``
|
||||
# fail-fast limitation is therefore removed; only genuinely non-tiling shapes are
|
||||
# rejected above (in ``build_folding_spec``).
|
||||
|
||||
if tp_preset is not None and ep_preset is not None and tp_preset != ep_preset:
|
||||
raise ValueError("tensor_parallel.preset_model and expert_parallel.preset_model must match when both "
|
||||
f"are set (tensor_parallel.preset_model={tp_preset!r}, "
|
||||
f"expert_parallel.preset_model={ep_preset!r}).")
|
||||
|
||||
if spec.tp_size > 1 and spec.ep_size == 1:
|
||||
raise ValueError("AutoEP+AutoTP folding requires expert_parallel.autoep_size > 1. "
|
||||
"The ep=1 local-computation path would duplicate routed-token gradients across TP lanes.")
|
||||
|
||||
if spec.tp_size > 1 and use_data_before_expert_parallel:
|
||||
raise ValueError("expert_parallel with use_data_before_expert_parallel_ is not supported with "
|
||||
"AutoEP+AutoTP folding. Disable use_data_before_expert_parallel_.")
|
||||
|
||||
if spec.tp_size > 1 and zero_stage == 3:
|
||||
raise ValueError("AutoEP+AutoTP with ZeRO stage 3 is reserved for the separate ZeRO-3 composition lane. "
|
||||
"Use ZeRO stage 0, 1, or 2 for this folding MVP.")
|
||||
|
||||
if spec.tp_size > 1 and (zero_offload_optimizer or zero_offload_param):
|
||||
raise ValueError("ZeRO optimizer/parameter offload with AutoEP+AutoTP folding is not validated yet. "
|
||||
"Disable offload or run a follow-up proof for per-family replica groups.")
|
||||
|
||||
mpu_tp = _mpu_world_size(mpu, "get_tensor_model_parallel_world_size", "get_model_parallel_world_size")
|
||||
if mpu_tp not in (None, 1, spec.tp_size):
|
||||
raise ValueError(f"mpu tensor/model parallel world size ({mpu_tp}) conflicts with "
|
||||
f"tensor_parallel.autotp_size={spec.tp_size}.")
|
||||
mpu_pp = _mpu_world_size(mpu, "get_pipeline_model_parallel_world_size", "get_pipeline_parallel_world_size")
|
||||
if mpu_pp not in (None, spec.pp_size):
|
||||
raise ValueError(f"mpu pipeline parallel world size ({mpu_pp}) conflicts with pp_size={spec.pp_size}.")
|
||||
|
||||
|
||||
def mark_autoep_folding_router_parameter(param) -> None:
|
||||
"""Tag a router/gate parameter as the *replicated* folded family (AVERAGE).
|
||||
|
||||
This is the ONLY family marker applied on the live forward path today:
|
||||
``AutoEPMoELayer.__init__`` marks every ``router.*`` parameter with it. The
|
||||
folded router runs redundantly on every TP peer (same tokens, same routing)
|
||||
and its gradient is reconstructed into a replicated full view by the restore
|
||||
all-gather (see ``deepspeed.moe.ep_tp_dispatch._AllGatherVariableRows`` and
|
||||
``restore_combined``). That all-gather backward scales each peer's slice by
|
||||
``tp_size``, so the extra TP reduction must AVERAGE (all_reduce then divide
|
||||
by ``tp_size``); SUM would leave the ``tp_size`` factor, i.e. the 2.0x
|
||||
parity regression the CPU/Gloo tests guard.
|
||||
"""
|
||||
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM)
|
||||
|
||||
|
||||
def mark_autoep_folding_partial_router_parameter(param) -> None:
|
||||
"""Tag a router/gate parameter as a *routed-token partial* family (SUM).
|
||||
|
||||
Forward-looking contract; NOT used on the current forward path -- only the
|
||||
unit tests in ``tests/unit/v1/moe/test_autoep_autotp_grad_parity.py`` set
|
||||
it. Use it only for a future design where the router's per-token work is
|
||||
genuinely partitioned across peers and the slices are NOT all-gathered back
|
||||
into a replicated full view, so each peer holds a real partial gradient that
|
||||
must be SUMed. Such a router is a SUM partial in any token-partitioned mode
|
||||
(``mp_mode in {"tp", "sp"}``) because its partition can ride the existing
|
||||
expert-dispatch all-to-all without changing the dense activation layout.
|
||||
Prove the SUM with a parity test (like the existing router/gate cases)
|
||||
before enabling it on a real forward path.
|
||||
"""
|
||||
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM)
|
||||
|
||||
|
||||
def mark_autoep_folding_sp_sharded_layernorm_parameter(param) -> None:
|
||||
"""Tag a LayerNorm parameter as *SP-sequence-sharded* family (SUM under SP).
|
||||
|
||||
Forward-looking contract; NOT used on the current forward path -- only the
|
||||
unit tests set it. Unlike the router, a LayerNorm has no adjacent dispatch
|
||||
all-to-all to ride on, so the only way to token-partition it is to shard the
|
||||
sequence dimension of the dense activations, which is Sequence Parallel by
|
||||
definition. It therefore becomes a SUM partial only when ``mp_mode == "sp"``
|
||||
and otherwise falls back to the replicated AVERAGE. Today ``tp_size > 1``
|
||||
with sequence parallelism is rejected in ``validate_folding_global``; this
|
||||
marker is the explicit contract for when that restriction is lifted, and
|
||||
must be backed by a parity test before use.
|
||||
"""
|
||||
setattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM)
|
||||
|
||||
|
||||
def _is_moe_param_marker(param) -> bool:
|
||||
return hasattr(param, "allreduce") and not param.allreduce
|
||||
|
||||
|
||||
def _is_model_parallel_param_marker(param) -> bool:
|
||||
return bool(getattr(param, "model_parallel", False) or getattr(param, "tensor_model_parallel", False))
|
||||
|
||||
|
||||
def _autoep_folding_param_family(param, *, param_name: str | None = None) -> str | None:
|
||||
"""Resolve a parameter's folded reduction family.
|
||||
|
||||
An explicit ``mark_autoep_folding_*`` tag always wins. The ``.router.`` name
|
||||
match is only a redundant safety net: ``AutoEPMoELayer`` already tags router
|
||||
params, so this fallback merely keeps the conservative *replicated* (AVERAGE)
|
||||
classification if some router param ever reaches the reducer untagged. It
|
||||
never returns a SUM family by name -- SUM families are opt-in via explicit
|
||||
markers only, so any unrecognized replicated/dense/LayerNorm param falls
|
||||
through to the AVERAGE default rather than being silently over-scaled.
|
||||
"""
|
||||
family = getattr(param, AUTOEP_FOLDING_PARAM_FAMILY_ATTR, None)
|
||||
if family is not None:
|
||||
return family
|
||||
if param_name is not None and ".router." in param_name:
|
||||
return AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM
|
||||
return None
|
||||
|
||||
|
||||
def autoep_folding_gradient_reduction_strategy(
|
||||
folding_spec: ParallelFoldingSpec | None,
|
||||
param,
|
||||
*,
|
||||
param_name: str | None = None,
|
||||
) -> str:
|
||||
"""Classify one folded TP/SP gradient as ``sum``, ``average``, or ``skip``.
|
||||
|
||||
TP means Tensor Parallel and SP means Sequence Parallel. The parallel mode
|
||||
alone is not a safe SUM-vs-AVG selector because different parameter
|
||||
families see different backward semantics:
|
||||
|
||||
- Router/gate parameters that are explicitly marked as routed-token
|
||||
partials in TP/SP token-partitioned modes receive one partial gradient per
|
||||
lane, so their TP/SP reduction is a SUM. The current AutoEP folded router
|
||||
gate is marked ``router_gate_replicated`` because the full-flow backward
|
||||
reaches this reducer as a lane-replicated gradient; that family uses the
|
||||
same AVERAGE normalization as other replicated parameters.
|
||||
- Dense and LayerNorm parameters that are merely replicated by TP folding
|
||||
are not routed-token partials; blindly SUMing them scales gradients by
|
||||
the TP size, so their extra TP reduction is an AVERAGE.
|
||||
- A true SP-sharded LayerNorm would be a partial-gradient parameter and
|
||||
should SUM. The current AutoEP folding path does not mark runtime
|
||||
LayerNorm parameters that way; the marker and strategy boundary exist so
|
||||
future SP support has an explicit contract instead of reusing the dense
|
||||
replicated default by accident.
|
||||
- Model-parallel (genuinely TP-sharded) parameters are SKIP because the
|
||||
TP-specific path owns their reduction.
|
||||
- Routed-expert parameters are EXPERT_TP_CANCEL: their data-parallel
|
||||
reduction is owned by the EP/EDP path, but the folded forward all-gathers
|
||||
their outputs into a replicated full view in ``restore_combined`` (whose
|
||||
backward injects a ``tp_size`` factor), so the expert-weight gradient
|
||||
reaches the optimizer ``tp_size`` times too large. Experts are not
|
||||
TP-replicated, so the fix is a plain ``/tp_size`` (no TP all_reduce), which
|
||||
is linear and composes with the EDP all_reduce in any order. Without this,
|
||||
folded expert gradients are over-scaled by ``tp_size`` -- invisible to
|
||||
scale-invariant Adam but real for SGD/Lion/Muon and for gradient clipping
|
||||
(it inflates the expert contribution to the global grad norm).
|
||||
|
||||
Underlying rule and mechanism: a folded parameter is replicated (AVERAGE)
|
||||
when the forward reconstructs its partitioned work into an identical full
|
||||
view inside the layer, and a genuine partial (SUM) only when the shard is
|
||||
kept all the way to the loss. Today the router/gate is partitioned across
|
||||
TP peers for dispatch but then all-gathered back by ``restore_combined``
|
||||
(see ``deepspeed.moe.ep_tp_dispatch``), whose backward scales each peer's
|
||||
gradient by ``tp_size``; the TP all_reduce then yields ``tp_size *
|
||||
full_grad`` and AVERAGE divides it out. Reducing with SUM would leave that
|
||||
factor -- the 2.0x router/gate parity regression the CPU/Gloo tests guard.
|
||||
The router can be a SUM partial in either ``tp`` or ``sp`` mode because its
|
||||
token partition can ride the existing dispatch all-to-all, whereas a
|
||||
LayerNorm becomes a partial only under true ``sp`` (sequence sharding): it
|
||||
has no adjacent all-to-all, so partitioning it requires changing the dense
|
||||
activation layout, which is Sequence Parallel by definition.
|
||||
|
||||
Both the DeepSpeedEngine path and the ZeRO-2 path call this helper so the
|
||||
policy cannot silently drift between optimizers.
|
||||
"""
|
||||
if folding_spec is None or getattr(folding_spec, "tp_size", 1) <= 1:
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
||||
if _is_model_parallel_param_marker(param):
|
||||
# Genuinely TP-sharded (column/row-parallel) params: the TP-specific path
|
||||
# owns their reduction. Not produced by the folded skip-partition MVP.
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
||||
if _is_moe_param_marker(param):
|
||||
# Routed-expert params. Their EP/EDP data-parallel reduction is owned by
|
||||
# the expert path, but the folded forward routes their outputs through the
|
||||
# ``restore_combined`` all-gather, whose backward leaves a ``tp_size``
|
||||
# factor on the expert-weight gradient (the same factor the replicated
|
||||
# router cancels with AVERAGE). Experts are NOT TP-replicated, so they must
|
||||
# not be TP all_reduced; the factor is cancelled with a plain ``/tp_size``.
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL
|
||||
|
||||
family = _autoep_folding_param_family(param, param_name=param_name)
|
||||
mp_mode = getattr(folding_spec, "mp_mode", "tp")
|
||||
token_partitioned_mode = mp_mode in ("tp", "sp")
|
||||
if family == AUTOEP_FOLDING_ROUTER_GATE_PARTIAL_PARAM:
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_SUM if token_partitioned_mode else AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
||||
if family == AUTOEP_FOLDING_ROUTER_GATE_REPLICATED_PARAM:
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
||||
if family == AUTOEP_FOLDING_SP_SHARDED_LAYERNORM_PARAM and mp_mode == "sp":
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_SUM
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE
|
||||
|
||||
|
||||
def reduce_autoep_folding_gradient(
|
||||
folding_spec: ParallelFoldingSpec | None,
|
||||
param,
|
||||
grad,
|
||||
*,
|
||||
tp_group,
|
||||
param_name: str | None = None,
|
||||
) -> str:
|
||||
strategy = autoep_folding_gradient_reduction_strategy(folding_spec, param, param_name=param_name)
|
||||
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_SKIP or grad is None or grad.data.is_sparse:
|
||||
return strategy
|
||||
|
||||
from deepspeed import comm as dist
|
||||
|
||||
grad_data = grad.data
|
||||
tp_world_size = dist.get_world_size(group=tp_group)
|
||||
|
||||
# Routed experts: cancel the ``tp_size`` factor the restore all-gather leaves,
|
||||
# WITHOUT a TP all_reduce (experts are not TP-replicated; cross-TP summation of
|
||||
# disjoint expert-token slices is owned by the EDP all_reduce). ``/tp_size`` is
|
||||
# linear, so it composes with that EDP reduction in either order.
|
||||
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_EXPERT_TP_CANCEL:
|
||||
if tp_world_size > 1:
|
||||
grad_data.div_(tp_world_size)
|
||||
return strategy
|
||||
|
||||
if grad_data.dtype != torch.float32:
|
||||
reduced = grad_data.float()
|
||||
dist.all_reduce(reduced, group=tp_group)
|
||||
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE:
|
||||
reduced.div_(tp_world_size)
|
||||
grad_data.copy_(reduced.to(grad_data.dtype))
|
||||
return strategy
|
||||
|
||||
dist.all_reduce(grad_data, group=tp_group)
|
||||
if strategy == AUTOEP_FOLDING_GRAD_REDUCE_AVERAGE:
|
||||
grad_data.div_(tp_world_size)
|
||||
return strategy
|
||||
|
||||
|
||||
def is_autoep_folding_gradient_corrected(param) -> bool:
|
||||
return bool(getattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, False))
|
||||
|
||||
|
||||
def clear_autoep_folding_gradient_corrected(param) -> None:
|
||||
if hasattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR):
|
||||
setattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, False)
|
||||
|
||||
|
||||
def apply_folding_correction_to_grad_buffer(
|
||||
folding_spec: ParallelFoldingSpec | None,
|
||||
param,
|
||||
grad,
|
||||
*,
|
||||
tp_group,
|
||||
param_name: str | None = None,
|
||||
use_correction_marker: bool = True,
|
||||
) -> str:
|
||||
if use_correction_marker and is_autoep_folding_gradient_corrected(param):
|
||||
return AUTOEP_FOLDING_GRAD_REDUCE_SKIP
|
||||
|
||||
strategy = reduce_autoep_folding_gradient(folding_spec, param, grad, tp_group=tp_group, param_name=param_name)
|
||||
if use_correction_marker and strategy != AUTOEP_FOLDING_GRAD_REDUCE_SKIP:
|
||||
setattr(param, AUTOEP_FOLDING_GRAD_CORRECTED_ATTR, True)
|
||||
return strategy
|
||||
|
||||
|
||||
def _normalize_rank_groups(groups: Iterable[Iterable[int]]) -> set[tuple[int, ...]]:
|
||||
return {tuple(int(rank) for rank in group) for group in groups}
|
||||
|
||||
|
||||
def assert_group_matches_spec(existing_rank_lists, spec: ParallelFoldingSpec, *, group_kind: str = "ep_edp") -> None:
|
||||
"""Ensure cached ``ep_size_N`` rank lists match the requested folding spec."""
|
||||
tables = expected_folding_group_tables(spec)
|
||||
expected_ep = _normalize_rank_groups(tables.ep_groups)
|
||||
expected_edp = _normalize_rank_groups(tables.edp_groups)
|
||||
|
||||
if isinstance(existing_rank_lists, dict):
|
||||
observed_ep = existing_rank_lists.get("ep", [])
|
||||
observed_edp = existing_rank_lists.get("edp", [])
|
||||
else:
|
||||
observed_ep, observed_edp = existing_rank_lists
|
||||
|
||||
for group in _normalize_rank_groups(observed_ep):
|
||||
if group not in expected_ep:
|
||||
raise RuntimeError(f"Cached expert-parallel group {group} does not match folding spec {spec}.")
|
||||
for group in _normalize_rank_groups(observed_edp):
|
||||
if group not in expected_edp:
|
||||
raise RuntimeError(f"Cached expert-data-parallel group {group} does not match folding spec {spec}.")
|
||||
@@ -0,0 +1,753 @@
|
||||
# Copyright (c) DeepSpeed Team.
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
|
||||
#
|
||||
# Portions of this file are derived from TorchTitan.
|
||||
# See THIRD_PARTY_NOTICES.md for the BSD-3-Clause notice.
|
||||
|
||||
# DeepSpeed Team
|
||||
"""AutoEP MoE Layer: drop-in replacement for HF MoE blocks with EP support.
|
||||
|
||||
Contains AutoEPMoELayer, compute_split_plan, _AllToAllV, and helper functions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults
|
||||
from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed.moe.ep_router import TokenChoiceTopKRouter
|
||||
from deepspeed.moe.ep_count import count_tokens_per_expert
|
||||
from deepspeed.moe.ep_experts import GroupedExperts
|
||||
from deepspeed.moe.ep_kernels import TokenReorderer
|
||||
from deepspeed.moe.ep_repack import _gather_source_zero_params, repack_expert_requires_grad_flags, repack_expert_weights
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Named tuples
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class RouterOutput(NamedTuple):
|
||||
top_scores: torch.Tensor # [T, K]
|
||||
selected_experts: torch.Tensor # [T, K]
|
||||
num_tokens_per_expert: torch.Tensor # [E_global]
|
||||
|
||||
|
||||
class SplitPlan(NamedTuple):
|
||||
input_splits: list[int] # len=ep_size
|
||||
output_splits: list[int] # len=ep_size
|
||||
local_counts: torch.Tensor # [E_local]
|
||||
local_counts_by_source: torch.Tensor # [ep_size, E_local]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def resolve_score_apply_mode(
|
||||
spec: MoELayerSpec,
|
||||
config_override: Literal["auto", "pre", "post"],
|
||||
) -> Literal["pre", "post"]:
|
||||
"""Resolve score-application mode from config override or preset default."""
|
||||
if config_override != "auto":
|
||||
return config_override
|
||||
return spec.score_apply
|
||||
|
||||
|
||||
def resolve_combine_impl(
|
||||
config_override: Literal["auto", "weighted_sum", "legacy_bmm"], ) -> Literal["weighted_sum", "legacy_bmm"]:
|
||||
"""Resolve combine implementation from config override or default."""
|
||||
if config_override != "auto":
|
||||
return config_override
|
||||
return "weighted_sum"
|
||||
|
||||
|
||||
def _copy_parameter_data(target: nn.Parameter, source: torch.Tensor) -> None:
|
||||
full_shape = torch.Size(getattr(source, "ds_shape", source.shape))
|
||||
with torch.no_grad():
|
||||
source_data = source.data
|
||||
if torch.Size(source_data.shape) != full_shape:
|
||||
raise RuntimeError("AutoEP source parameter must be gathered before copying: "
|
||||
f"expected full shape {tuple(full_shape)}, got {tuple(source_data.shape)}")
|
||||
if (torch.Size(target.data.shape) != full_shape or target.data.dtype != source_data.dtype
|
||||
or target.data.device != source_data.device):
|
||||
target.data = torch.empty(full_shape, dtype=source_data.dtype, device=source_data.device)
|
||||
target.data.copy_(source_data)
|
||||
|
||||
|
||||
def apply_scores_before_experts_if_enabled(
|
||||
routed_input: torch.Tensor,
|
||||
top_scores: torch.Tensor,
|
||||
score_apply: Literal["pre", "post"],
|
||||
) -> torch.Tensor:
|
||||
"""Pre-multiply token representations by router scores before expert compute."""
|
||||
if score_apply == "pre":
|
||||
return (routed_input.to(torch.float32) * top_scores.reshape(-1, 1)).to(routed_input.dtype)
|
||||
return routed_input
|
||||
|
||||
|
||||
def compute_split_plan(
|
||||
selected_experts: torch.Tensor, # [T, K]
|
||||
num_experts: int,
|
||||
ep_size: int,
|
||||
num_local_experts: int,
|
||||
ep_group: dist.ProcessGroup | None,
|
||||
) -> SplitPlan:
|
||||
"""Compute AllToAllV split sizes for token dispatch/combine.
|
||||
|
||||
Returns SplitPlan with input_splits, output_splits, local_counts, and
|
||||
local_counts_by_source.
|
||||
"""
|
||||
T_K = selected_experts.numel()
|
||||
|
||||
if ep_size == 1:
|
||||
# No dispatch needed - all tokens stay local
|
||||
num_tokens_per_expert = count_tokens_per_expert(
|
||||
selected_experts,
|
||||
num_experts,
|
||||
out_dtype=torch.int32,
|
||||
)
|
||||
return SplitPlan(
|
||||
input_splits=[T_K],
|
||||
output_splits=[T_K],
|
||||
local_counts=num_tokens_per_expert,
|
||||
local_counts_by_source=num_tokens_per_expert.view(1, num_local_experts),
|
||||
)
|
||||
|
||||
# Count tokens per expert globally
|
||||
num_tokens_per_expert = count_tokens_per_expert(
|
||||
selected_experts,
|
||||
num_experts,
|
||||
out_dtype=torch.int32,
|
||||
)
|
||||
|
||||
# Reshape to [ep_size, num_local_experts] to get per-rank counts
|
||||
count_matrix = num_tokens_per_expert.view(ep_size, num_local_experts)
|
||||
|
||||
# input_splits: how many tokens THIS rank sends to each destination rank
|
||||
input_splits = count_matrix.sum(dim=1).cpu().tolist()
|
||||
|
||||
# Exchange counts with all ranks to get output_splits
|
||||
# Each rank tells every other rank how many tokens it will send
|
||||
local_counts_tensor = count_matrix.sum(dim=1).clone() # [ep_size]
|
||||
remote_counts_tensor = torch.zeros_like(local_counts_tensor)
|
||||
|
||||
dist.all_to_all_single(
|
||||
remote_counts_tensor,
|
||||
local_counts_tensor,
|
||||
group=ep_group,
|
||||
)
|
||||
output_splits = remote_counts_tensor.cpu().tolist()
|
||||
|
||||
# local_counts: how many tokens this rank will process for each local expert
|
||||
# After receiving tokens, we need per-expert counts for this rank
|
||||
local_expert_counts = count_matrix[:, :].clone() # [ep_size, E_local]
|
||||
|
||||
# Exchange the detailed per-expert counts
|
||||
# Each rank needs to know, for its local experts, how many tokens come from each source
|
||||
local_expert_counts_flat = local_expert_counts.view(-1).contiguous() # [ep_size * E_local]
|
||||
received_counts_flat = torch.zeros_like(local_expert_counts_flat)
|
||||
|
||||
dist.all_to_all_single(
|
||||
received_counts_flat,
|
||||
local_expert_counts_flat,
|
||||
group=ep_group,
|
||||
)
|
||||
|
||||
# Sum over source ranks to get total per local expert
|
||||
received_counts = received_counts_flat.view(ep_size, num_local_experts)
|
||||
local_counts = received_counts.sum(dim=0) # [E_local]
|
||||
|
||||
return SplitPlan(
|
||||
input_splits=input_splits,
|
||||
output_splits=output_splits,
|
||||
local_counts=local_counts,
|
||||
local_counts_by_source=received_counts,
|
||||
)
|
||||
|
||||
|
||||
def compute_split_plan_from_expert_indices(
|
||||
expert_indices: torch.Tensor,
|
||||
num_experts: int,
|
||||
ep_size: int,
|
||||
num_local_experts: int,
|
||||
ep_group: dist.ProcessGroup | None,
|
||||
) -> SplitPlan:
|
||||
"""Compute EP AllToAllV splits for an already partitioned assignment list."""
|
||||
if ep_size == 1:
|
||||
counts = count_tokens_per_expert(expert_indices, num_experts, out_dtype=torch.int32)
|
||||
return SplitPlan([int(expert_indices.numel())], [int(expert_indices.numel())], counts,
|
||||
counts.view(1, num_local_experts))
|
||||
|
||||
counts = count_tokens_per_expert(expert_indices, num_experts, out_dtype=torch.int32)
|
||||
count_matrix = counts.view(ep_size, num_local_experts)
|
||||
input_splits = count_matrix.sum(dim=1).cpu().tolist()
|
||||
local_counts_tensor = count_matrix.sum(dim=1).clone()
|
||||
remote_counts_tensor = torch.zeros_like(local_counts_tensor)
|
||||
dist.all_to_all_single(remote_counts_tensor, local_counts_tensor, group=ep_group)
|
||||
output_splits = remote_counts_tensor.cpu().tolist()
|
||||
|
||||
local_expert_counts_flat = count_matrix.reshape(-1).contiguous()
|
||||
received_counts_flat = torch.zeros_like(local_expert_counts_flat)
|
||||
dist.all_to_all_single(received_counts_flat, local_expert_counts_flat, group=ep_group)
|
||||
received_counts = received_counts_flat.view(ep_size, num_local_experts)
|
||||
local_counts = received_counts.sum(dim=0)
|
||||
return SplitPlan(input_splits, output_splits, local_counts, received_counts)
|
||||
|
||||
|
||||
class _AllToAllV(torch.autograd.Function):
|
||||
"""Autograd-compatible all-to-all with variable split sizes."""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, group, x, input_splits, output_splits):
|
||||
ctx.group = group
|
||||
ctx.input_splits = input_splits
|
||||
ctx.output_splits = output_splits
|
||||
|
||||
output_size = sum(output_splits)
|
||||
output = torch.empty(
|
||||
(output_size, x.shape[1]),
|
||||
dtype=x.dtype,
|
||||
device=x.device,
|
||||
)
|
||||
|
||||
dist.all_to_all_single(
|
||||
output,
|
||||
x.contiguous(),
|
||||
output_split_sizes=output_splits,
|
||||
input_split_sizes=input_splits,
|
||||
group=group,
|
||||
)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out):
|
||||
# Reverse the splits for backward
|
||||
grad_out = grad_out.contiguous()
|
||||
input_size = sum(ctx.input_splits)
|
||||
grad_input = torch.empty(
|
||||
(input_size, grad_out.shape[1]),
|
||||
dtype=grad_out.dtype,
|
||||
device=grad_out.device,
|
||||
)
|
||||
|
||||
dist.all_to_all_single(
|
||||
grad_input,
|
||||
grad_out,
|
||||
output_split_sizes=ctx.input_splits,
|
||||
input_split_sizes=ctx.output_splits,
|
||||
group=ctx.group,
|
||||
)
|
||||
return None, grad_input, None, None
|
||||
|
||||
|
||||
def permute_by_local_expert(
|
||||
tokens: torch.Tensor,
|
||||
local_counts: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
||||
"""Reorder tokens so they are grouped contiguously by local expert ID.
|
||||
|
||||
Uses TorchTitan's Triton kernel for permutation index generation.
|
||||
|
||||
Returns:
|
||||
tokens_permuted: [N_padded, H] (alignment-padded)
|
||||
permuted_indices: [N_padded] (maps padded positions -> original positions)
|
||||
aligned_counts: [E_local] aligned token counts per expert (for expert computation)
|
||||
n_tokens: original token count before padding (for unpermute)
|
||||
"""
|
||||
from deepspeed.moe.ep_kernels import generate_permute_indices, TOKEN_GROUP_ALIGN_SIZE_M
|
||||
|
||||
if local_counts.ndim == 1:
|
||||
# [E_local]: already aggregated over sources (ep_degree=1)
|
||||
ep_degree = 1
|
||||
num_local_experts = local_counts.shape[0]
|
||||
local_counts_flat = local_counts
|
||||
elif local_counts.ndim == 2:
|
||||
# [ep_size, E_local]: preserve per-source layout for correct regrouping
|
||||
ep_degree, num_local_experts = local_counts.shape
|
||||
local_counts_flat = local_counts.reshape(-1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"local_counts must have shape [E_local] or [ep_degree, E_local], got {tuple(local_counts.shape)}")
|
||||
|
||||
n_tokens = tokens.shape[0]
|
||||
alignment = TOKEN_GROUP_ALIGN_SIZE_M
|
||||
|
||||
# Compute padded max length
|
||||
x_padded_per_expert = n_tokens + num_local_experts * alignment
|
||||
padded_max_len = ((x_padded_per_expert + alignment - 1) // alignment) * alignment
|
||||
|
||||
# Use the pure-PyTorch path for host tensors. The CPU accelerator reports
|
||||
# CPU tensors as "on accelerator", but Triton still requires a GPU driver.
|
||||
use_cpu = tokens.device.type == "cpu"
|
||||
counts_for_permute = local_counts_flat.cpu() if use_cpu else local_counts_flat
|
||||
with torch.no_grad():
|
||||
permuted_indices, m_sizes, _offsets = generate_permute_indices(
|
||||
counts_for_permute,
|
||||
num_local_experts,
|
||||
ep_degree,
|
||||
padded_max_len,
|
||||
alignment,
|
||||
use_cpu=use_cpu,
|
||||
)
|
||||
if not use_cpu:
|
||||
permuted_indices = permuted_indices.to(tokens.device)
|
||||
m_sizes = m_sizes.to(tokens.device)
|
||||
|
||||
# Add padding row for out-of-bounds indices (index n_tokens -> zero row)
|
||||
tokens_padded = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1], ))))
|
||||
tokens_permuted = tokens_padded[permuted_indices, :]
|
||||
|
||||
return tokens_permuted, permuted_indices, m_sizes, n_tokens
|
||||
|
||||
|
||||
def unpermute_by_local_expert(
|
||||
expert_output: torch.Tensor,
|
||||
permuted_indices: torch.Tensor,
|
||||
n_tokens: int,
|
||||
) -> torch.Tensor:
|
||||
"""Reverse permute_by_local_expert: restore original token order and strip padding.
|
||||
|
||||
Args:
|
||||
expert_output: [N_padded, H] from expert computation
|
||||
permuted_indices: [N_padded] index mapping from permute_by_local_expert
|
||||
n_tokens: original token count before alignment padding
|
||||
"""
|
||||
# Scatter expert outputs back to original positions.
|
||||
# permuted_indices values range 0..n_tokens, where n_tokens is the zero-padding row.
|
||||
out_unpermuted = expert_output.new_zeros((n_tokens + 1, expert_output.shape[-1]))
|
||||
out_unpermuted[permuted_indices, :] = expert_output
|
||||
# Strip the zero-padding row to get [n_tokens, H]
|
||||
return out_unpermuted[:-1]
|
||||
|
||||
|
||||
def combine_from_routed(
|
||||
expert_output: torch.Tensor, # [N, H]
|
||||
top_scores: torch.Tensor, # [T, K]
|
||||
token_indices_sorted: torch.Tensor, # [N]
|
||||
top_k: int,
|
||||
score_apply: Literal["pre", "post"],
|
||||
combine_impl: Literal["weighted_sum", "legacy_bmm"],
|
||||
shape: tuple[int, int, int], # (B, S, H)
|
||||
) -> torch.Tensor:
|
||||
"""Scatter-add expert outputs back to original token positions."""
|
||||
bsz, seqlen, hdim = shape
|
||||
T = bsz * seqlen
|
||||
|
||||
# Create output tensor
|
||||
output = torch.zeros(T * top_k, hdim, dtype=expert_output.dtype, device=expert_output.device)
|
||||
|
||||
# Place expert outputs back in unsorted order
|
||||
output[token_indices_sorted] = expert_output
|
||||
|
||||
# Reshape to [T, K, H]
|
||||
output = output.reshape(T, top_k, hdim)
|
||||
|
||||
if score_apply == "post":
|
||||
if combine_impl == "legacy_bmm":
|
||||
# Legacy reduction path retained as a debug option for model-family
|
||||
# verification. The weighted-sum path is the default.
|
||||
output = torch.bmm(
|
||||
top_scores.reshape(-1, 1, top_k).float(),
|
||||
output.float(),
|
||||
).to(expert_output.dtype).squeeze(1)
|
||||
else:
|
||||
# Match the runtime HF grouped-mm path: apply routing weights per
|
||||
# token-slot sample, then reduce over top-k.
|
||||
output = (output.float() * top_scores.reshape(T, top_k, 1).float()).sum(dim=1).to(expert_output.dtype)
|
||||
else:
|
||||
# Scores already applied pre-experts, just sum over top_k
|
||||
output = output.sum(dim=1)
|
||||
|
||||
return output.reshape(bsz, seqlen, hdim)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AutoEPMoELayer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AutoEPMoELayer(nn.Module):
|
||||
"""Drop-in replacement for HF MoE blocks with Expert Parallelism support."""
|
||||
|
||||
_is_autoep_layer = True # Marker for AutoTP skip handshake
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
spec: MoELayerSpec,
|
||||
source_module: nn.Module,
|
||||
ep_size: int,
|
||||
ep_rank: int,
|
||||
config: AutoEPConfig,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model_family = spec.model_family
|
||||
self.return_router_logits = spec.return_router_logits
|
||||
self.router_logits_capture_target = spec.router_logits_capture_target
|
||||
self.router_logits_capture_index = spec.router_logits_capture_index
|
||||
self.router_logits_capture_mode = spec.router_logits_capture_mode
|
||||
self.moe_output_shape = spec.moe_output_shape
|
||||
self.top_k = spec.top_k
|
||||
self.score_apply = resolve_score_apply_mode(spec, config.score_apply)
|
||||
self.combine_impl = resolve_combine_impl(config.combine_impl)
|
||||
route_norm = spec.route_norm if config.route_norm is None else config.route_norm
|
||||
self.ep_size = ep_size
|
||||
self.ep_rank = ep_rank
|
||||
self.num_experts = spec.num_experts
|
||||
self.num_local_experts = spec.num_experts // ep_size
|
||||
self.hidden_size = spec.hidden_size
|
||||
self.ep_group_name = f"ep_size_{ep_size}"
|
||||
self.ep_group = None # Set by set_deepspeed_parallelism()
|
||||
self.folding_group_handles = None
|
||||
self.tp_group = None
|
||||
resolved_config = resolve_autoep_config_defaults(config, spec.model_family)
|
||||
self.validate_folding_routing = bool(resolved_config.validate_folding_routing)
|
||||
|
||||
# Router: copy gate weights from source
|
||||
source_gate = getattr(source_module, spec.router_name)
|
||||
source_gate_bias = getattr(source_gate, 'bias', None)
|
||||
source_ecb = getattr(source_gate, 'e_score_correction_bias', None)
|
||||
unsupported_router_biases = [
|
||||
getattr(source_gate, bias_name, None) for bias_name in spec.unsupported_router_bias_names
|
||||
]
|
||||
if not spec.supports_expert_bias and resolved_config.load_balance_coeff is not None:
|
||||
raise ValueError(f"AutoEP preset '{spec.model_family}' does not support load_balance_coeff/expert_bias "
|
||||
"yet. Set load_balance_coeff=None.")
|
||||
with _gather_source_zero_params([source_gate.weight, source_gate_bias, source_ecb,
|
||||
*unsupported_router_biases]):
|
||||
for bias_name, router_bias in zip(spec.unsupported_router_bias_names, unsupported_router_biases):
|
||||
if router_bias is None:
|
||||
continue
|
||||
if torch.is_tensor(router_bias) and torch.count_nonzero(router_bias.detach()).item() == 0:
|
||||
continue
|
||||
raise ValueError(f"AutoEP preset '{spec.model_family}' does not support nonzero router bias "
|
||||
f"'{bias_name}' yet.")
|
||||
self.router = TokenChoiceTopKRouter(
|
||||
dim=spec.hidden_size,
|
||||
num_experts=spec.num_experts,
|
||||
num_expert_groups=spec.num_expert_groups,
|
||||
num_limited_groups=spec.num_limited_groups,
|
||||
top_k=spec.top_k,
|
||||
score_func=spec.score_func,
|
||||
route_norm=route_norm,
|
||||
route_scale=spec.route_scale,
|
||||
gate_bias=spec.gate_bias,
|
||||
group_score_func=spec.group_score_func,
|
||||
)
|
||||
# Copy gate weights
|
||||
_copy_parameter_data(self.router.gate.weight, source_gate.weight)
|
||||
self.router.gate.weight.requires_grad_(source_gate.weight.requires_grad)
|
||||
if spec.gate_bias and source_gate_bias is not None:
|
||||
_copy_parameter_data(self.router.gate.bias, source_gate_bias)
|
||||
self.router.gate.bias.requires_grad_(source_gate_bias.requires_grad)
|
||||
|
||||
# Copy pre-trained score correction bias (DeepSeek-V3/Moonlight noaux_tc routing)
|
||||
if source_ecb is not None and isinstance(source_ecb, nn.Parameter):
|
||||
self.router.e_score_correction_bias = nn.Parameter(source_ecb.data.clone(),
|
||||
requires_grad=source_ecb.requires_grad)
|
||||
logger.info('AutoEP: copied e_score_correction_bias from source gate '
|
||||
'(shape=%s)', source_ecb.shape)
|
||||
|
||||
# Alias router under the name OutputRecorder expects (layer_name if provided),
|
||||
# but only when OutputRecorder captures from the router child and the alias is safe.
|
||||
alias_target = spec.router_logits_capture_layer_name or spec.router_name
|
||||
if spec.router_logits_capture_target == "router" and alias_target != "router":
|
||||
if "." in alias_target or alias_target in ("experts", "shared_experts") or hasattr(self, alias_target):
|
||||
logger.warning(f"Skipping router alias '{alias_target}' to avoid name collision.")
|
||||
else:
|
||||
setattr(self, alias_target, self.router)
|
||||
|
||||
# Experts: extract local expert weights
|
||||
w1, w2, w3 = repack_expert_weights(
|
||||
experts_source=getattr(source_module, spec.experts_name),
|
||||
spec=spec,
|
||||
ep_rank=ep_rank,
|
||||
ep_size=ep_size,
|
||||
)
|
||||
w1_requires_grad, w2_requires_grad, w3_requires_grad = repack_expert_requires_grad_flags(
|
||||
experts_source=getattr(source_module, spec.experts_name),
|
||||
spec=spec,
|
||||
ep_rank=ep_rank,
|
||||
ep_size=ep_size,
|
||||
)
|
||||
self.experts = GroupedExperts(
|
||||
dim=spec.hidden_size,
|
||||
hidden_dim=spec.ffn_hidden_size,
|
||||
num_experts=self.num_local_experts,
|
||||
use_grouped_mm=config.use_grouped_mm,
|
||||
)
|
||||
_copy_parameter_data(self.experts.w1, w1)
|
||||
_copy_parameter_data(self.experts.w2, w2)
|
||||
_copy_parameter_data(self.experts.w3, w3)
|
||||
self.experts.w1.requires_grad_(w1_requires_grad)
|
||||
self.experts.w2.requires_grad_(w2_requires_grad)
|
||||
self.experts.w3.requires_grad_(w3_requires_grad)
|
||||
|
||||
self.reorderer = TokenReorderer(num_experts=self.num_experts, top_k=self.top_k)
|
||||
self.shared_experts = getattr(source_module, spec.shared_experts_name,
|
||||
None) if spec.has_shared_experts else None
|
||||
self.shared_experts_gate = getattr(source_module, spec.shared_experts_gate_name,
|
||||
None) if spec.shared_experts_gate_name else None
|
||||
|
||||
# Mark expert params for EDP gradient reduction
|
||||
for param in self.experts.parameters():
|
||||
param.allreduce = False
|
||||
param.group_name = self.ep_group_name
|
||||
param.ds_zero_placement_family = "autoep_expert"
|
||||
param.ds_zero_partition_group_name = self.ep_group_name
|
||||
|
||||
# Mark shared expert and router params for global DP reduction.
|
||||
# The router runs redundantly on every TP peer and its gradient is
|
||||
# rebuilt into a replicated full view by the restore all-gather, so it
|
||||
# is tagged as the replicated family (AVERAGE TP reduction); a SUM would
|
||||
# double it under tp_size=2. See mark_autoep_folding_router_parameter.
|
||||
for param in self.router.parameters():
|
||||
param.allreduce = True
|
||||
mark_autoep_folding_router_parameter(param)
|
||||
param.ds_zero_placement_family = "replicated"
|
||||
if self.shared_experts is not None:
|
||||
for param in self.shared_experts.parameters():
|
||||
param.allreduce = True
|
||||
param.ds_zero_placement_family = "replicated"
|
||||
if self.shared_experts_gate is not None:
|
||||
for param in self.shared_experts_gate.parameters():
|
||||
param.allreduce = True
|
||||
param.ds_zero_placement_family = "replicated"
|
||||
|
||||
# Load balancing buffers
|
||||
self.load_balance_coeff = resolved_config.load_balance_coeff
|
||||
buf_device = source_gate.weight.device
|
||||
if self.load_balance_coeff is not None:
|
||||
self.register_buffer(
|
||||
"expert_bias",
|
||||
torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
|
||||
persistent=True,
|
||||
)
|
||||
else:
|
||||
self.expert_bias = None
|
||||
self.register_buffer(
|
||||
"tokens_per_expert",
|
||||
torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
# Router-logit cache
|
||||
self._cached_router_logits = None
|
||||
self._register_logit_hook()
|
||||
|
||||
def _register_logit_hook(self):
|
||||
"""Register a forward hook that caches gate logits for OutputRecorder capture."""
|
||||
if self.router_logits_capture_target != "router":
|
||||
return
|
||||
|
||||
def hook_fn(module, input, output):
|
||||
x = input[0] # [T, H]
|
||||
logits = module.gate(x) # [T, E_global]
|
||||
if self.router_logits_capture_mode == "post_score":
|
||||
if self.router.score_func == "softmax":
|
||||
logits = torch.softmax(logits.float(), dim=-1).to(logits.dtype)
|
||||
elif self.router.score_func == "sigmoid":
|
||||
logits = torch.sigmoid(logits.float()).to(logits.dtype)
|
||||
self._cached_router_logits = logits
|
||||
|
||||
self.router.register_forward_hook(hook_fn)
|
||||
|
||||
def set_deepspeed_parallelism(
|
||||
self,
|
||||
use_data_before_expert_parallel_: bool = False,
|
||||
folding_group_handles=None,
|
||||
) -> None:
|
||||
"""Bind EP group handle to this module."""
|
||||
from deepspeed.utils import groups
|
||||
from deepspeed.utils.bwc import bwc_pipeline_parallel_world_size
|
||||
|
||||
if folding_group_handles is not None:
|
||||
self.folding_group_handles = folding_group_handles
|
||||
self.ep_group_name = folding_group_handles.ep_group_name
|
||||
self.ep_group = folding_group_handles.ep_group
|
||||
self.tp_group = folding_group_handles.tp_group
|
||||
self.ep_rank = dist.get_rank(group=self.ep_group)
|
||||
return
|
||||
|
||||
if self.ep_group_name not in groups._get_expert_parallel_group_dict():
|
||||
mp_size = max(
|
||||
getattr(groups, '_get_model_parallel_world_size', lambda: 1)(),
|
||||
getattr(groups, '_get_sequence_parallel_world_size', lambda: 1)(),
|
||||
)
|
||||
mp_mode = "tp" if getattr(groups, '_get_model_parallel_world_size', lambda: 1)() > 1 else "sp"
|
||||
pp_size = 1 if groups.mpu is None else bwc_pipeline_parallel_world_size(groups.mpu)
|
||||
groups._create_expert_and_data_parallel(
|
||||
expert_parallel_size_=self.ep_size,
|
||||
mp_size=mp_size,
|
||||
pp_size=pp_size,
|
||||
mp_mode=mp_mode,
|
||||
use_data_before_expert_parallel_=use_data_before_expert_parallel_,
|
||||
)
|
||||
self.ep_group = groups._get_expert_parallel_group(self.ep_group_name)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
hidden_states: [B, S, H]
|
||||
|
||||
Returns:
|
||||
[B, S, H] or ([B, S, H], [T, E]) if return_router_logits.
|
||||
Some HF MoE contracts return ([T, H], [T, E]) instead.
|
||||
"""
|
||||
bsz, seqlen, hdim = hidden_states.shape
|
||||
x = hidden_states.reshape(-1, hdim) # [T, H]
|
||||
|
||||
# Router
|
||||
ro: RouterOutput = RouterOutput(*self.router(x, self.expert_bias))
|
||||
|
||||
# Accumulate expert utilization
|
||||
with torch.no_grad():
|
||||
self.tokens_per_expert.add_(ro.num_tokens_per_expert)
|
||||
|
||||
# Reorder tokens by expert
|
||||
top_scores_sorted, token_indices_sorted, _ = self.reorderer(ro.top_scores, ro.selected_experts)
|
||||
expert_indices_sorted = ro.selected_experts.reshape(-1).index_select(0, token_indices_sorted)
|
||||
|
||||
folded_tp = self.folding_group_handles is not None and self.folding_group_handles.spec.tp_size > 1
|
||||
restore_ctx = None
|
||||
if folded_tp:
|
||||
from deepspeed.moe.ep_tp_dispatch import (
|
||||
RoutedAssignmentPayload,
|
||||
assignment_ordinals_by_expert,
|
||||
assert_tp_payload_consistent,
|
||||
dispatch_counters,
|
||||
partition_assignments,
|
||||
restore_combined,
|
||||
)
|
||||
payload = RoutedAssignmentPayload(
|
||||
token_indices=(token_indices_sorted // self.top_k).to(torch.long),
|
||||
expert_indices=expert_indices_sorted.to(torch.long),
|
||||
assignment_indices=assignment_ordinals_by_expert(expert_indices_sorted.to(torch.long)),
|
||||
capacity_slots=(token_indices_sorted % self.top_k).to(torch.long),
|
||||
combine_weights=top_scores_sorted
|
||||
if self.score_apply == "post" else torch.ones_like(top_scores_sorted),
|
||||
drop_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
|
||||
pad_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
|
||||
input_splits=[0 for _ in range(self.ep_size)],
|
||||
output_splits=[0 for _ in range(self.ep_size)],
|
||||
extra={
|
||||
"destination_ranks": (expert_indices_sorted // self.num_local_experts).to(torch.long),
|
||||
"top_scores": top_scores_sorted,
|
||||
"num_tokens": torch.tensor(bsz * seqlen, device=hidden_states.device, dtype=torch.long),
|
||||
},
|
||||
)
|
||||
if self.validate_folding_routing:
|
||||
assert_tp_payload_consistent(payload,
|
||||
tp_group=self.tp_group,
|
||||
tp_size=self.folding_group_handles.spec.tp_size)
|
||||
tp_rank = dist.get_rank(group=self.tp_group)
|
||||
local_payload, restore_ctx = partition_assignments(payload,
|
||||
tp_group=self.tp_group,
|
||||
tp_rank=tp_rank,
|
||||
tp_size=self.folding_group_handles.spec.tp_size)
|
||||
token_indices_for_compute = token_indices_sorted.index_select(0, restore_ctx.local_indices)
|
||||
top_scores_for_compute = top_scores_sorted.index_select(0, restore_ctx.local_indices)
|
||||
expert_indices_for_plan = local_payload.expert_indices
|
||||
else:
|
||||
token_indices_for_compute = token_indices_sorted
|
||||
top_scores_for_compute = top_scores_sorted
|
||||
expert_indices_for_plan = expert_indices_sorted
|
||||
|
||||
routed_input = x[token_indices_for_compute // self.top_k] # [N, H]
|
||||
routed_input = apply_scores_before_experts_if_enabled(routed_input,
|
||||
top_scores_for_compute,
|
||||
score_apply=self.score_apply)
|
||||
|
||||
if self.ep_size == 1:
|
||||
# No AllToAll needed - local computation only
|
||||
local_counts = count_tokens_per_expert(
|
||||
ro.selected_experts,
|
||||
self.num_local_experts,
|
||||
out_dtype=torch.int32,
|
||||
)
|
||||
|
||||
routed_input_permuted, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
|
||||
routed_input, local_counts)
|
||||
expert_output = self.experts(routed_input_permuted, aligned_counts)
|
||||
expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)
|
||||
else:
|
||||
# EP dispatch/compute/combine
|
||||
if folded_tp:
|
||||
plan = compute_split_plan_from_expert_indices(
|
||||
expert_indices=expert_indices_for_plan,
|
||||
num_experts=self.num_experts,
|
||||
ep_size=self.ep_size,
|
||||
num_local_experts=self.num_local_experts,
|
||||
ep_group=self.ep_group,
|
||||
)
|
||||
else:
|
||||
plan = compute_split_plan(
|
||||
selected_experts=ro.selected_experts,
|
||||
num_experts=self.num_experts,
|
||||
ep_size=self.ep_size,
|
||||
num_local_experts=self.num_local_experts,
|
||||
ep_group=self.ep_group,
|
||||
)
|
||||
|
||||
routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits)
|
||||
|
||||
routed_input, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
|
||||
routed_input, plan.local_counts_by_source)
|
||||
expert_output = self.experts(routed_input, aligned_counts)
|
||||
expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)
|
||||
|
||||
expert_output = _AllToAllV.apply(self.ep_group, expert_output, plan.output_splits, plan.input_splits)
|
||||
|
||||
if folded_tp:
|
||||
output = restore_combined(expert_output,
|
||||
restore_ctx,
|
||||
tp_group=self.tp_group,
|
||||
validate_coverage=self.validate_folding_routing).reshape(bsz, seqlen, hdim)
|
||||
self._last_folding_dispatch_counters = dispatch_counters(restore_ctx)
|
||||
else:
|
||||
output = combine_from_routed(
|
||||
expert_output,
|
||||
top_scores=ro.top_scores,
|
||||
token_indices_sorted=token_indices_sorted,
|
||||
top_k=self.top_k,
|
||||
score_apply=self.score_apply,
|
||||
combine_impl=self.combine_impl,
|
||||
shape=(bsz, seqlen, hdim),
|
||||
)
|
||||
|
||||
if self.moe_output_shape == "flat":
|
||||
output = output.reshape(-1, hdim)
|
||||
shared_expert_input = x
|
||||
elif self.shared_experts_gate is not None:
|
||||
shared_expert_input = x
|
||||
else:
|
||||
shared_expert_input = hidden_states
|
||||
|
||||
if self.shared_experts is not None:
|
||||
shared_expert_output = self.shared_experts(shared_expert_input)
|
||||
if self.shared_experts_gate is not None:
|
||||
shared_expert_gate = torch.sigmoid(self.shared_experts_gate(shared_expert_input))
|
||||
shared_expert_output = shared_expert_gate * shared_expert_output
|
||||
if shared_expert_output.shape != output.shape:
|
||||
shared_expert_output = shared_expert_output.reshape_as(output)
|
||||
output = output + shared_expert_output
|
||||
|
||||
if self.return_router_logits:
|
||||
logits = self._cached_router_logits
|
||||
self._cached_router_logits = None
|
||||
return output, logits
|
||||
|
||||
self._cached_router_logits = None
|
||||
return output
|
||||
@@ -0,0 +1,25 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""Compatibility shim for AutoEP preset adapter APIs."""
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import (
|
||||
AutoEPPresetAdapter,
|
||||
ForwardContract,
|
||||
GroupRoutingConfig,
|
||||
TransformersTopLevelRouterLogitsAdapter,
|
||||
)
|
||||
from deepspeed.module_inject.auto_ep_presets.deepseek_v2 import DeepSeekV2PresetAdapter
|
||||
from deepspeed.module_inject.auto_ep_presets.deepseek_v3 import DeepSeekV3PresetAdapter
|
||||
from deepspeed.module_inject.auto_ep_presets.qwen3_5_moe import Qwen35MoePresetAdapter
|
||||
from deepspeed.module_inject.auto_ep_presets.registry import get_preset_adapter
|
||||
|
||||
__all__ = [
|
||||
"AutoEPPresetAdapter",
|
||||
"DeepSeekV2PresetAdapter",
|
||||
"DeepSeekV3PresetAdapter",
|
||||
"ForwardContract",
|
||||
"GroupRoutingConfig",
|
||||
"Qwen35MoePresetAdapter",
|
||||
"TransformersTopLevelRouterLogitsAdapter",
|
||||
"get_preset_adapter",
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""AutoEP built-in preset registry package."""
|
||||
@@ -0,0 +1,433 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""Shared AutoEP preset dataclasses and adapter interface."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Any, Callable, Literal, NoReturn
|
||||
|
||||
import torch.nn as nn
|
||||
from packaging.version import InvalidVersion, Version
|
||||
|
||||
# Sentinel for "not specified in config, use preset default".
|
||||
# Unlike None (which means "fused gate+up, no separate w3"), _UNSET means
|
||||
# the user did not set the field at all. Compare with `is _UNSET`.
|
||||
_UNSET = object()
|
||||
|
||||
|
||||
def _raise_unsupported_load_balance_coeff(value: object) -> NoReturn:
|
||||
raise ValueError(f"load_balance_coeff={value!r} is not supported in this AutoEP build "
|
||||
"(would register expert_bias and route through unsupported "
|
||||
"auxiliary-loss-free load balancing). Set load_balance_coeff to null "
|
||||
"or omit the key.")
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoEModelPreset:
|
||||
"""Preset configuration for a known MoE model family."""
|
||||
|
||||
moe_layer_pattern: str
|
||||
router_pattern: str
|
||||
experts_pattern: str
|
||||
expert_storage: Literal["fused_3d", "module_list"]
|
||||
expert_w1: str
|
||||
expert_w2: str
|
||||
expert_w3: str | None
|
||||
num_experts_attr: str
|
||||
top_k_attr: str
|
||||
score_func: Literal["softmax", "sigmoid"]
|
||||
score_apply: Literal["pre", "post"]
|
||||
route_norm: bool
|
||||
gate_bias: bool
|
||||
has_shared_experts: bool = False
|
||||
shared_experts_pattern: str = ""
|
||||
shared_experts_gate_pattern: str = ""
|
||||
autoep_config_defaults: dict[str, Any] = field(default_factory=dict)
|
||||
supports_expert_bias: bool = True
|
||||
unsupported_router_bias_names: tuple[str, ...] = ()
|
||||
preset_adapter: str = "default"
|
||||
hf_model_types: tuple[str, ...] = ()
|
||||
unsupported_hf_model_type_notes: dict[str, str] = field(default_factory=dict)
|
||||
min_transformers_version: str | None = None
|
||||
validated_transformers_versions: str = ""
|
||||
docs_support_notes: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoELayerSpec:
|
||||
"""Detected MoE layer specification for a single module in the model."""
|
||||
|
||||
moe_module_name: str
|
||||
model_family: str
|
||||
router_name: str
|
||||
experts_name: str
|
||||
expert_storage: Literal["fused_3d", "module_list"]
|
||||
expert_w1_name: str
|
||||
expert_w2_name: str
|
||||
expert_w3_name: str | None
|
||||
num_experts: int
|
||||
top_k: int
|
||||
hidden_size: int
|
||||
ffn_hidden_size: int
|
||||
score_func: Literal["softmax", "sigmoid"]
|
||||
score_apply: Literal["pre", "post"]
|
||||
route_norm: bool
|
||||
gate_bias: bool
|
||||
return_router_logits: bool
|
||||
router_logits_capture_target: Literal["moe_block", "router", "none"]
|
||||
router_logits_capture_index: int | None
|
||||
router_logits_capture_layer_name: str | None
|
||||
has_shared_experts: bool
|
||||
shared_experts_name: str
|
||||
shared_experts_gate_name: str = ""
|
||||
route_scale: float = 1.0
|
||||
num_expert_groups: int | None = None
|
||||
num_limited_groups: int | None = None
|
||||
group_score_func: Literal["max", "top2_sum"] = "top2_sum"
|
||||
supports_expert_bias: bool = True
|
||||
unsupported_router_bias_names: tuple[str, ...] = ()
|
||||
preset_adapter: str = "default"
|
||||
router_logits_capture_mode: Literal["raw", "post_score"] = "post_score"
|
||||
moe_output_shape: Literal["batched", "flat"] = "batched"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEPConfig:
|
||||
"""User-facing configuration parsed from DS config JSON."""
|
||||
|
||||
enabled: bool = False
|
||||
autoep_size: int = 1
|
||||
expert_tensor_parallel_size: int = 1
|
||||
validate_folding_routing: bool = False
|
||||
preset_model: str | None = None
|
||||
moe_layer_pattern: str | None = None
|
||||
expert_pattern: str | None = None
|
||||
router_pattern: str | None = None
|
||||
use_grouped_mm: bool = True
|
||||
route_norm: bool | None = None
|
||||
route_scale: float = 1.0
|
||||
score_apply: Literal["auto", "pre", "post"] = "auto"
|
||||
combine_impl: Literal["auto", "weighted_sum", "legacy_bmm"] = "auto"
|
||||
num_expert_groups: int | None = None
|
||||
num_limited_groups: int | None = None
|
||||
score_func: Literal["auto", "softmax", "sigmoid"] = "auto"
|
||||
top_k: int | str = "auto"
|
||||
load_balance_coeff: float | None | object = _UNSET
|
||||
routed_scaling_factor: float | str = "auto"
|
||||
expert_w1: str | None = None
|
||||
expert_w2: str | None = None
|
||||
expert_w3: object = _UNSET
|
||||
num_experts_attr: str | None = None
|
||||
top_k_attr: str | None = None
|
||||
has_shared_experts: bool | None = None
|
||||
shared_experts_pattern: str | None = None
|
||||
shared_experts_gate_pattern: str | None = None
|
||||
_load_balance_coeff_explicit: bool = field(default=False, init=False, repr=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.load_balance_coeff is _UNSET:
|
||||
self.load_balance_coeff = None
|
||||
self._load_balance_coeff_explicit = False
|
||||
else:
|
||||
self._load_balance_coeff_explicit = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GroupRoutingConfig:
|
||||
num_expert_groups: int | None
|
||||
num_limited_groups: int | None
|
||||
group_score_func: Literal["max", "top2_sum"] = "top2_sum"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ForwardContract:
|
||||
return_router_logits: bool = False
|
||||
capture_target: Literal["moe_block", "router", "none"] = "none"
|
||||
capture_index: int | None = None
|
||||
capture_layer_name: str | None = None
|
||||
router_logits_capture_mode: Literal["raw", "post_score"] = "post_score"
|
||||
moe_output_shape: Literal["batched", "flat"] = "batched"
|
||||
|
||||
|
||||
class AutoEPPresetAdapter:
|
||||
"""Default behavior shared by presets without model-specific parser rules."""
|
||||
|
||||
def validate_compatibility(
|
||||
self,
|
||||
preset_name: str,
|
||||
preset: MoEModelPreset,
|
||||
model_config,
|
||||
) -> None:
|
||||
"""Validate public HF compatibility metadata for a selected preset."""
|
||||
model_type = getattr(model_config, "model_type", None) if model_config is not None else None
|
||||
self._validate_hf_model_type(preset_name, preset, model_type)
|
||||
self._validate_transformers_version(preset_name, preset, model_type)
|
||||
|
||||
def _validate_hf_model_type(
|
||||
self,
|
||||
preset_name: str,
|
||||
preset: MoEModelPreset,
|
||||
model_type: str | None,
|
||||
) -> None:
|
||||
if model_type is None:
|
||||
return
|
||||
|
||||
unsupported_note = preset.unsupported_hf_model_type_notes.get(model_type)
|
||||
if unsupported_note is None:
|
||||
return
|
||||
|
||||
supported = ", ".join(repr(value) for value in preset.hf_model_types) or "none"
|
||||
raise ValueError(f"AutoEP preset '{preset_name}' does not support model_type='{model_type}'. "
|
||||
f"{unsupported_note} Supported HF model_type value(s): {supported}.")
|
||||
|
||||
def _validate_transformers_version(
|
||||
self,
|
||||
preset_name: str,
|
||||
preset: MoEModelPreset,
|
||||
model_type: str | None,
|
||||
) -> None:
|
||||
min_version = preset.min_transformers_version
|
||||
if min_version is None or model_type is None:
|
||||
return
|
||||
if not self._requires_transformers_version_validation():
|
||||
return
|
||||
if model_type not in preset.hf_model_types and model_type not in preset.unsupported_hf_model_type_notes:
|
||||
return
|
||||
|
||||
try:
|
||||
installed_version = self._installed_transformers_version()
|
||||
except Exception as exc:
|
||||
raise ValueError(f"AutoEP preset '{preset_name}' for model_type='{model_type}' requires "
|
||||
f"Transformers >= {min_version}, but transformers could not be imported: {exc}.") from exc
|
||||
|
||||
try:
|
||||
installed = Version(installed_version)
|
||||
minimum = Version(min_version)
|
||||
except InvalidVersion as exc:
|
||||
raise ValueError(f"AutoEP preset '{preset_name}' for model_type='{model_type}' requires "
|
||||
f"Transformers >= {min_version}, but the installed Transformers version "
|
||||
f"'{installed_version}' could not be parsed.") from exc
|
||||
|
||||
if installed < minimum:
|
||||
raise ValueError(f"AutoEP preset '{preset_name}' for model_type='{model_type}' requires "
|
||||
f"Transformers >= {min_version}, but installed transformers=={installed_version}. "
|
||||
"Upgrade Transformers or choose a preset/model combination supported by the "
|
||||
"installed Transformers version.")
|
||||
|
||||
def _installed_transformers_version(self) -> str:
|
||||
import transformers
|
||||
return getattr(transformers, "__version__", "unknown")
|
||||
|
||||
def _requires_transformers_version_validation(self) -> bool:
|
||||
# The default adapter also covers non-HF/mock/custom-compatible configs;
|
||||
# specialized HF-only adapters opt in to minimum Transformers checks.
|
||||
return False
|
||||
|
||||
def resolve_route_norm(
|
||||
self,
|
||||
config: AutoEPConfig,
|
||||
preset: MoEModelPreset,
|
||||
model_config,
|
||||
) -> bool:
|
||||
if config.route_norm is not None:
|
||||
return config.route_norm
|
||||
|
||||
cfg_norm = getattr(model_config, 'norm_topk_prob', None)
|
||||
if cfg_norm is not None:
|
||||
return bool(cfg_norm)
|
||||
return preset.route_norm
|
||||
|
||||
def resolve_group_routing(
|
||||
self,
|
||||
config: AutoEPConfig,
|
||||
model_config,
|
||||
) -> GroupRoutingConfig:
|
||||
return GroupRoutingConfig(
|
||||
num_expert_groups=config.num_expert_groups,
|
||||
num_limited_groups=config.num_limited_groups,
|
||||
)
|
||||
|
||||
def resolve_expert_layout(
|
||||
self,
|
||||
experts_module: nn.Module,
|
||||
preset: MoEModelPreset,
|
||||
) -> MoEModelPreset:
|
||||
return preset
|
||||
|
||||
def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract:
|
||||
return contract
|
||||
|
||||
def retarget_transformers_output_recorders(
|
||||
self,
|
||||
model: nn.Module,
|
||||
spec: MoELayerSpec,
|
||||
replacement: nn.Module,
|
||||
retargeted_keys: set[str],
|
||||
remove_output_capture_hooks: Callable[[nn.Module], int],
|
||||
) -> None:
|
||||
return
|
||||
|
||||
|
||||
_MISSING_REGISTRY_ENTRY = object()
|
||||
|
||||
|
||||
def _restore_transformers_output_capture_registry(
|
||||
registry: dict[str, Any],
|
||||
original_entries: dict[str, object],
|
||||
) -> None:
|
||||
for registry_key, original_entry in original_entries.items():
|
||||
if original_entry is _MISSING_REGISTRY_ENTRY:
|
||||
registry.pop(registry_key, None)
|
||||
else:
|
||||
registry[registry_key] = original_entry
|
||||
|
||||
|
||||
def _install_instance_transformers_output_recorders(
|
||||
model: nn.Module,
|
||||
registry_entries: dict[str, dict[str, Any]],
|
||||
output_capturing: Any,
|
||||
remove_output_capture_hooks: Callable[[nn.Module], int],
|
||||
) -> bool:
|
||||
maybe_install_capturing_hooks = getattr(output_capturing, "maybe_install_capturing_hooks", None)
|
||||
registry = getattr(output_capturing, "_CAN_RECORD_REGISTRY", None)
|
||||
if not callable(maybe_install_capturing_hooks) or not isinstance(registry, dict):
|
||||
return False
|
||||
|
||||
remove_output_capture_hooks(model)
|
||||
for module in model.modules():
|
||||
if hasattr(module, "_output_capturing_hooks_installed"):
|
||||
module._output_capturing_hooks_installed = False
|
||||
model._output_capturing_hooks_installed = False
|
||||
|
||||
original_entries = {
|
||||
registry_key: registry.get(registry_key, _MISSING_REGISTRY_ENTRY)
|
||||
for registry_key in registry_entries
|
||||
}
|
||||
try:
|
||||
registry.update(registry_entries)
|
||||
maybe_install_capturing_hooks(model)
|
||||
finally:
|
||||
_restore_transformers_output_capture_registry(registry, original_entries)
|
||||
return True
|
||||
|
||||
|
||||
def _retarget_transformers_output_recorders_for_modules(
|
||||
*,
|
||||
model: nn.Module,
|
||||
display_name: str,
|
||||
recorder_key: str,
|
||||
retargeted_keys: set[str],
|
||||
remove_output_capture_hooks: Callable[[nn.Module], int],
|
||||
module_matches: Callable[[nn.Module], bool],
|
||||
make_output_recorder: Callable[[Any], Any],
|
||||
) -> int:
|
||||
try:
|
||||
from transformers.utils import output_capturing
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
registry = getattr(output_capturing, "_CAN_RECORD_REGISTRY", None)
|
||||
if not isinstance(registry, dict):
|
||||
return 0
|
||||
|
||||
registry_entries: dict[str, dict[str, Any]] = {}
|
||||
retargeted = 0
|
||||
for module in model.modules():
|
||||
if not module_matches(module):
|
||||
continue
|
||||
|
||||
registry_key = str(module.__class__)
|
||||
record_outputs = getattr(module, "_can_record_outputs", None)
|
||||
registry_outputs = registry.get(registry_key)
|
||||
base_outputs = record_outputs if isinstance(record_outputs, dict) else registry_outputs
|
||||
if not isinstance(base_outputs, dict) or "router_logits" not in base_outputs:
|
||||
continue
|
||||
|
||||
retargeted_outputs = dict(base_outputs)
|
||||
retargeted_outputs["router_logits"] = make_output_recorder(output_capturing.OutputRecorder)
|
||||
module._can_record_outputs = retargeted_outputs
|
||||
registry_entries[registry_key] = retargeted_outputs
|
||||
retargeted += 1
|
||||
|
||||
if retargeted == 0:
|
||||
from deepspeed.utils import logger
|
||||
logger.warning(f"AutoEP: {display_name} conversion did not find a HF output-capture registry "
|
||||
"entry for router_logits.")
|
||||
return 0
|
||||
|
||||
if _install_instance_transformers_output_recorders(
|
||||
model,
|
||||
registry_entries,
|
||||
output_capturing,
|
||||
remove_output_capture_hooks,
|
||||
):
|
||||
return retargeted
|
||||
|
||||
if recorder_key in retargeted_keys:
|
||||
return retargeted
|
||||
retargeted_keys.add(recorder_key)
|
||||
registry.update(registry_entries)
|
||||
if getattr(model, "_output_capturing_hooks_installed", False):
|
||||
remove_output_capture_hooks(model)
|
||||
model._output_capturing_hooks_installed = False
|
||||
return retargeted
|
||||
|
||||
|
||||
class TransformersTopLevelRouterLogitsAdapter(AutoEPPresetAdapter):
|
||||
"""Retarget Transformers model-level router-logit recorders to AutoEP."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
display_name: str,
|
||||
hf_model_types: tuple[str, ...],
|
||||
class_name_fragments: tuple[str, ...],
|
||||
) -> None:
|
||||
self.display_name = display_name
|
||||
self.hf_model_types = hf_model_types
|
||||
self.class_name_fragments = class_name_fragments
|
||||
|
||||
def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract:
|
||||
# Mixtral/Qwen3/Qwen2 capture raw router logits through Transformers'
|
||||
# model-level OutputRecorder hooks. AutoEP keeps the MoE block tensor
|
||||
# return contract intact and retargets the recorder to router.gate.
|
||||
return replace(
|
||||
contract,
|
||||
return_router_logits=False,
|
||||
capture_target="router",
|
||||
capture_index=0,
|
||||
router_logits_capture_mode="raw",
|
||||
)
|
||||
|
||||
def retarget_transformers_output_recorders(
|
||||
self,
|
||||
model: nn.Module,
|
||||
spec: MoELayerSpec,
|
||||
replacement: nn.Module,
|
||||
retargeted_keys: set[str],
|
||||
remove_output_capture_hooks: Callable[[nn.Module], int],
|
||||
) -> None:
|
||||
recorder_key = f"{spec.model_family}:{replacement.__class__.__module__}.{replacement.__class__.__qualname__}"
|
||||
|
||||
router_gate = getattr(getattr(replacement, "router", None), "gate", None)
|
||||
if router_gate is None:
|
||||
return
|
||||
|
||||
def module_matches(module: nn.Module) -> bool:
|
||||
module_config = getattr(module, "config", None)
|
||||
model_type = getattr(module_config, "model_type", None)
|
||||
class_name = module.__class__.__name__
|
||||
return (model_type in self.hf_model_types
|
||||
or any(fragment in class_name for fragment in self.class_name_fragments))
|
||||
|
||||
_retarget_transformers_output_recorders_for_modules(
|
||||
model=model,
|
||||
display_name=self.display_name,
|
||||
recorder_key=recorder_key,
|
||||
retargeted_keys=retargeted_keys,
|
||||
remove_output_capture_hooks=remove_output_capture_hooks,
|
||||
module_matches=module_matches,
|
||||
make_output_recorder=lambda OutputRecorder: OutputRecorder(
|
||||
router_gate.__class__, index=0, layer_name="router.gate"),
|
||||
)
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""DeepSeek-V2 AutoEP preset and parser adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import (
|
||||
AutoEPConfig,
|
||||
AutoEPPresetAdapter,
|
||||
GroupRoutingConfig,
|
||||
MoEModelPreset,
|
||||
)
|
||||
|
||||
PRESET_NAME = "deepseek_v2"
|
||||
|
||||
PRESET = MoEModelPreset(
|
||||
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
|
||||
router_pattern="gate",
|
||||
experts_pattern="experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1="gate_up_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3=None,
|
||||
num_experts_attr="n_routed_experts",
|
||||
top_k_attr="num_experts_per_tok",
|
||||
score_func="softmax",
|
||||
score_apply="post",
|
||||
route_norm=False,
|
||||
gate_bias=False,
|
||||
has_shared_experts=True,
|
||||
shared_experts_pattern="shared_experts",
|
||||
autoep_config_defaults={"load_balance_coeff": None},
|
||||
supports_expert_bias=False,
|
||||
preset_adapter="deepseek_v2",
|
||||
hf_model_types=("deepseek_v2", ),
|
||||
min_transformers_version="5.0.0",
|
||||
docs_support_notes=("load_balance_coeff / expert-bias auxiliary-loss-free load balancing "
|
||||
"is not currently supported; non-null values are rejected."),
|
||||
)
|
||||
|
||||
|
||||
class DeepSeekV2PresetAdapter(AutoEPPresetAdapter):
|
||||
"""DeepSeek-V2 keeps native top-k normalization and optional group-limited routing."""
|
||||
|
||||
def _requires_transformers_version_validation(self) -> bool:
|
||||
return True
|
||||
|
||||
def resolve_route_norm(
|
||||
self,
|
||||
config: AutoEPConfig,
|
||||
preset: MoEModelPreset,
|
||||
model_config,
|
||||
) -> bool:
|
||||
if config.route_norm is not None:
|
||||
return config.route_norm
|
||||
return preset.route_norm
|
||||
|
||||
def resolve_group_routing(
|
||||
self,
|
||||
config: AutoEPConfig,
|
||||
model_config,
|
||||
) -> GroupRoutingConfig:
|
||||
group_routing = super().resolve_group_routing(config, model_config)
|
||||
if getattr(model_config, 'topk_method', None) != "group_limited_greedy":
|
||||
return group_routing
|
||||
|
||||
return GroupRoutingConfig(
|
||||
num_expert_groups=group_routing.num_expert_groups or getattr(model_config, 'n_group', None),
|
||||
num_limited_groups=group_routing.num_limited_groups or getattr(model_config, 'topk_group', None),
|
||||
group_score_func="max",
|
||||
)
|
||||
|
||||
|
||||
PRESET_ADAPTERS = {
|
||||
"deepseek_v2": DeepSeekV2PresetAdapter(),
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""DeepSeek-V3 AutoEP preset and parser adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import AutoEPConfig, AutoEPPresetAdapter, GroupRoutingConfig, MoEModelPreset
|
||||
|
||||
PRESET_NAME = "deepseek_v3"
|
||||
|
||||
PRESET = MoEModelPreset(
|
||||
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
|
||||
router_pattern="gate",
|
||||
experts_pattern="experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1="gate_up_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3=None,
|
||||
num_experts_attr="n_routed_experts",
|
||||
top_k_attr="num_experts_per_tok",
|
||||
score_func="sigmoid",
|
||||
score_apply="post",
|
||||
route_norm=False,
|
||||
gate_bias=False,
|
||||
has_shared_experts=True,
|
||||
shared_experts_pattern="shared_experts",
|
||||
autoep_config_defaults={"load_balance_coeff": None},
|
||||
supports_expert_bias=False,
|
||||
unsupported_router_bias_names=(),
|
||||
preset_adapter="deepseek_v3",
|
||||
hf_model_types=("deepseek_v3", ),
|
||||
min_transformers_version="5.0.0",
|
||||
docs_support_notes=("load_balance_coeff / expert-bias auxiliary-loss-free load balancing "
|
||||
"is not currently supported; non-null values are rejected."),
|
||||
)
|
||||
|
||||
|
||||
class DeepSeekV3PresetAdapter(AutoEPPresetAdapter):
|
||||
"""DeepSeek-V3 always carries group-limited routing fields when present."""
|
||||
|
||||
def _requires_transformers_version_validation(self) -> bool:
|
||||
return True
|
||||
|
||||
def resolve_group_routing(
|
||||
self,
|
||||
config: AutoEPConfig,
|
||||
model_config,
|
||||
) -> GroupRoutingConfig:
|
||||
group_routing = super().resolve_group_routing(config, model_config)
|
||||
return GroupRoutingConfig(
|
||||
num_expert_groups=group_routing.num_expert_groups or getattr(model_config, 'n_group', None),
|
||||
num_limited_groups=group_routing.num_limited_groups or getattr(model_config, 'topk_group', None),
|
||||
group_score_func=group_routing.group_score_func,
|
||||
)
|
||||
|
||||
def resolve_expert_layout(
|
||||
self,
|
||||
experts_module: nn.Module,
|
||||
preset: MoEModelPreset,
|
||||
) -> MoEModelPreset:
|
||||
if not isinstance(experts_module, nn.ModuleList) or len(experts_module) == 0:
|
||||
return preset
|
||||
|
||||
default_fused_layout = (preset.expert_storage == "fused_3d" and preset.expert_w1 == "gate_up_proj"
|
||||
and preset.expert_w2 == "down_proj" and preset.expert_w3 is None)
|
||||
if not default_fused_layout:
|
||||
return preset
|
||||
|
||||
expert0 = experts_module[0]
|
||||
if not all(_has_expert_projection(expert0, name) for name in ("gate_proj", "up_proj", "down_proj")):
|
||||
return preset
|
||||
|
||||
return replace(
|
||||
preset,
|
||||
expert_storage="module_list",
|
||||
expert_w1="gate_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3="up_proj",
|
||||
)
|
||||
|
||||
|
||||
def _has_expert_projection(expert_module: nn.Module, name: str) -> bool:
|
||||
projection = getattr(expert_module, name, None)
|
||||
if projection is None:
|
||||
return False
|
||||
if isinstance(projection, (nn.Linear, nn.Parameter)):
|
||||
return True
|
||||
return hasattr(projection, "weight")
|
||||
|
||||
|
||||
PRESET_ADAPTERS = {
|
||||
"deepseek_v3": DeepSeekV3PresetAdapter(),
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""Mixtral AutoEP preset."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import MoEModelPreset, TransformersTopLevelRouterLogitsAdapter
|
||||
|
||||
PRESET_NAME = "mixtral"
|
||||
|
||||
PRESET = MoEModelPreset(
|
||||
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
|
||||
router_pattern="gate",
|
||||
experts_pattern="experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1="gate_up_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3=None,
|
||||
num_experts_attr="num_local_experts",
|
||||
top_k_attr="num_experts_per_tok",
|
||||
score_func="softmax",
|
||||
score_apply="post",
|
||||
route_norm=True,
|
||||
gate_bias=False,
|
||||
preset_adapter="mixtral",
|
||||
hf_model_types=("mixtral", ),
|
||||
min_transformers_version="5.0.0",
|
||||
)
|
||||
|
||||
PRESET_ADAPTERS = {
|
||||
"mixtral":
|
||||
TransformersTopLevelRouterLogitsAdapter(
|
||||
display_name="Mixtral",
|
||||
hf_model_types=("mixtral", ),
|
||||
class_name_fragments=("Mixtral", ),
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""Qwen3.5-MoE AutoEP preset and parser adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
from typing import Callable
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import (
|
||||
AutoEPPresetAdapter,
|
||||
ForwardContract,
|
||||
MoELayerSpec,
|
||||
MoEModelPreset,
|
||||
_retarget_transformers_output_recorders_for_modules,
|
||||
)
|
||||
|
||||
PRESET_NAME = "qwen3_5_moe"
|
||||
|
||||
PRESET = MoEModelPreset(
|
||||
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
|
||||
router_pattern="gate",
|
||||
experts_pattern="experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1="gate_up_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3=None,
|
||||
num_experts_attr="num_experts",
|
||||
top_k_attr="num_experts_per_tok",
|
||||
score_func="softmax",
|
||||
score_apply="post",
|
||||
route_norm=True,
|
||||
gate_bias=False,
|
||||
has_shared_experts=True,
|
||||
shared_experts_pattern="shared_expert",
|
||||
shared_experts_gate_pattern="shared_expert_gate",
|
||||
preset_adapter="qwen3_5_moe",
|
||||
hf_model_types=("qwen3_5_moe_text", ),
|
||||
unsupported_hf_model_type_notes={
|
||||
"qwen3_5_moe": ("AutoEP supports the Qwen3.5 text backbone preset path; pass the "
|
||||
"text-backbone model/config with model_type='qwen3_5_moe_text'.")
|
||||
},
|
||||
min_transformers_version="5.2.0",
|
||||
docs_support_notes="Requires the Qwen3.5 text-backbone qwen3_5_moe_text model type.",
|
||||
)
|
||||
|
||||
|
||||
class Qwen35MoePresetAdapter(AutoEPPresetAdapter):
|
||||
"""Qwen3.5 MoE exposes router logits through HF output recording."""
|
||||
|
||||
def _requires_transformers_version_validation(self) -> bool:
|
||||
return True
|
||||
|
||||
def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract:
|
||||
# HF records Qwen3.5 router output on Qwen3_5MoeTopKRouter. AutoEP replaces
|
||||
# the owning MoE block, so replacement output index 1 is used for recorder retargeting.
|
||||
return replace(
|
||||
contract,
|
||||
return_router_logits=True,
|
||||
capture_target="router",
|
||||
capture_index=1,
|
||||
)
|
||||
|
||||
def retarget_transformers_output_recorders(
|
||||
self,
|
||||
model: nn.Module,
|
||||
spec: MoELayerSpec,
|
||||
replacement: nn.Module,
|
||||
retargeted_keys: set[str],
|
||||
remove_output_capture_hooks: Callable[[nn.Module], int],
|
||||
) -> None:
|
||||
recorder_key = f"{spec.model_family}:{replacement.__class__.__module__}.{replacement.__class__.__qualname__}"
|
||||
|
||||
replacement_cls = replacement.__class__
|
||||
|
||||
def module_matches(module: nn.Module) -> bool:
|
||||
module_config = getattr(module, "config", None)
|
||||
model_type = getattr(module_config, "model_type", None)
|
||||
class_name = module.__class__.__name__
|
||||
return model_type == "qwen3_5_moe_text" or "Qwen3_5Moe" in class_name
|
||||
|
||||
_retarget_transformers_output_recorders_for_modules(
|
||||
model=model,
|
||||
display_name="Qwen3.5 AutoEP",
|
||||
recorder_key=recorder_key,
|
||||
retargeted_keys=retargeted_keys,
|
||||
remove_output_capture_hooks=remove_output_capture_hooks,
|
||||
module_matches=module_matches,
|
||||
make_output_recorder=lambda OutputRecorder: OutputRecorder(replacement_cls, index=1),
|
||||
)
|
||||
|
||||
|
||||
PRESET_ADAPTERS = {
|
||||
"qwen3_5_moe": Qwen35MoePresetAdapter(),
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""Qwen3-MoE AutoEP preset."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import MoEModelPreset, TransformersTopLevelRouterLogitsAdapter
|
||||
|
||||
PRESET_NAME = "qwen3_moe"
|
||||
|
||||
PRESET = MoEModelPreset(
|
||||
moe_layer_pattern=r"model\.layers\.\d+\.mlp",
|
||||
router_pattern="gate",
|
||||
experts_pattern="experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1="gate_up_proj",
|
||||
expert_w2="down_proj",
|
||||
expert_w3=None,
|
||||
num_experts_attr="num_experts",
|
||||
top_k_attr="num_experts_per_tok",
|
||||
score_func="softmax",
|
||||
score_apply="post",
|
||||
route_norm=True,
|
||||
gate_bias=False,
|
||||
has_shared_experts=True,
|
||||
shared_experts_pattern="shared_expert",
|
||||
shared_experts_gate_pattern="shared_expert_gate",
|
||||
preset_adapter="qwen3_moe",
|
||||
hf_model_types=("qwen3_moe", "qwen2_moe"),
|
||||
min_transformers_version="5.0.0",
|
||||
docs_support_notes=("Also covers Qwen2-MoE when the installed Transformers build uses the "
|
||||
"validated fused expert layout."),
|
||||
)
|
||||
|
||||
PRESET_ADAPTERS = {
|
||||
"qwen3_moe":
|
||||
TransformersTopLevelRouterLogitsAdapter(
|
||||
display_name="Qwen3-MoE/Qwen2-MoE",
|
||||
hf_model_types=("qwen3_moe", "qwen2_moe"),
|
||||
class_name_fragments=("Qwen3Moe", "Qwen2Moe"),
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DeepSpeed Team
|
||||
"""AutoEP preset registry and config override helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from dataclasses import replace
|
||||
|
||||
from deepspeed.module_inject.auto_ep_presets.base import (
|
||||
_UNSET,
|
||||
AutoEPConfig,
|
||||
AutoEPPresetAdapter,
|
||||
MoEModelPreset,
|
||||
)
|
||||
from deepspeed.module_inject.auto_ep_presets import deepseek_v2, deepseek_v3, mixtral, qwen3_5_moe, qwen3_moe
|
||||
from deepspeed.utils import logger
|
||||
|
||||
_PRESET_MODULES = (
|
||||
mixtral,
|
||||
qwen3_moe,
|
||||
qwen3_5_moe,
|
||||
deepseek_v2,
|
||||
deepseek_v3,
|
||||
)
|
||||
|
||||
PRESET_MODELS: dict[str, MoEModelPreset] = {module.PRESET_NAME: module.PRESET for module in _PRESET_MODULES}
|
||||
|
||||
_PRESET_ADAPTERS: dict[str, AutoEPPresetAdapter] = {
|
||||
"default": AutoEPPresetAdapter(),
|
||||
}
|
||||
for _module in _PRESET_MODULES:
|
||||
_PRESET_ADAPTERS.update(getattr(_module, "PRESET_ADAPTERS", {}))
|
||||
|
||||
|
||||
def _validate_registered_preset_adapters(
|
||||
preset_models: dict[str, MoEModelPreset] | None = None,
|
||||
preset_adapters: dict[str, AutoEPPresetAdapter] | None = None,
|
||||
) -> None:
|
||||
"""Fail fast if a registered preset references an adapter that is not registered."""
|
||||
preset_models = PRESET_MODELS if preset_models is None else preset_models
|
||||
preset_adapters = _PRESET_ADAPTERS if preset_adapters is None else preset_adapters
|
||||
|
||||
missing_presets = []
|
||||
for preset_name, preset in preset_models.items():
|
||||
if preset.preset_adapter not in preset_adapters:
|
||||
missing_presets.append((preset_name, preset.preset_adapter))
|
||||
|
||||
if not missing_presets:
|
||||
return
|
||||
|
||||
details = ", ".join(f"{preset_name}:{adapter_name}" for preset_name, adapter_name in missing_presets)
|
||||
raise RuntimeError(f"AutoEP preset registry is inconsistent; missing preset_adapter registration(s): {details}")
|
||||
|
||||
|
||||
_validate_registered_preset_adapters()
|
||||
|
||||
_PRESET_DEFAULT_EXPLICIT_FLAGS = {
|
||||
"load_balance_coeff": "_load_balance_coeff_explicit",
|
||||
}
|
||||
|
||||
|
||||
def available_preset_names() -> tuple[str, ...]:
|
||||
"""Return built-in AutoEP preset names in registry order."""
|
||||
return tuple(PRESET_MODELS.keys())
|
||||
|
||||
|
||||
def get_preset(preset_name: str) -> MoEModelPreset:
|
||||
"""Return a registered AutoEP preset by name."""
|
||||
preset = PRESET_MODELS.get(preset_name)
|
||||
if preset is None:
|
||||
raise ValueError(f"Unknown preset_model '{preset_name}'. Available presets: {list(available_preset_names())}")
|
||||
return preset
|
||||
|
||||
|
||||
def get_preset_adapter(adapter_name: str) -> AutoEPPresetAdapter:
|
||||
"""Return a registered AutoEP preset adapter by name."""
|
||||
adapter = _PRESET_ADAPTERS.get(adapter_name)
|
||||
if adapter is None:
|
||||
raise ValueError(f"Unknown AutoEP preset adapter '{adapter_name}'")
|
||||
return adapter
|
||||
|
||||
|
||||
def preset_name_for_hf_model_type(model_type: str) -> str | None:
|
||||
"""Return the AutoEP preset name for a supported HF model_type."""
|
||||
for preset_name, preset in PRESET_MODELS.items():
|
||||
if model_type in preset.hf_model_types:
|
||||
return preset_name
|
||||
return None
|
||||
|
||||
|
||||
def unsupported_preset_for_hf_model_type(model_type: str) -> tuple[str, MoEModelPreset] | None:
|
||||
"""Return a preset carrying an actionable diagnostic for an unsupported HF model_type."""
|
||||
for preset_name, preset in PRESET_MODELS.items():
|
||||
if model_type in preset.unsupported_hf_model_type_notes:
|
||||
return preset_name, preset
|
||||
return None
|
||||
|
||||
|
||||
def resolve_autoep_config_defaults(config: AutoEPConfig, preset_name: str | None) -> AutoEPConfig:
|
||||
"""Return config with preset-level AutoEP defaults applied where the user did not override."""
|
||||
if preset_name is None or preset_name not in PRESET_MODELS:
|
||||
return config
|
||||
|
||||
preset_defaults = PRESET_MODELS[preset_name].autoep_config_defaults
|
||||
if not preset_defaults:
|
||||
return config
|
||||
|
||||
resolved = copy.copy(config)
|
||||
for field_name, default_value in preset_defaults.items():
|
||||
explicit_flag = _PRESET_DEFAULT_EXPLICIT_FLAGS.get(field_name)
|
||||
if explicit_flag is None:
|
||||
continue
|
||||
if not getattr(config, explicit_flag, False):
|
||||
setattr(resolved, field_name, default_value)
|
||||
return resolved
|
||||
|
||||
|
||||
def apply_config_overrides(config: AutoEPConfig, preset: MoEModelPreset) -> MoEModelPreset:
|
||||
"""Apply explicit AutoEP config overrides to a preset.
|
||||
|
||||
Return the original preset object when there are no overrides. When overrides
|
||||
are present, return a dataclass copy so the registered preset remains unchanged.
|
||||
"""
|
||||
overrides = {}
|
||||
if config.moe_layer_pattern is not None:
|
||||
overrides["moe_layer_pattern"] = config.moe_layer_pattern
|
||||
if config.router_pattern is not None:
|
||||
overrides["router_pattern"] = config.router_pattern
|
||||
if config.expert_pattern is not None:
|
||||
overrides["experts_pattern"] = config.expert_pattern
|
||||
if config.expert_w1 is not None:
|
||||
overrides["expert_w1"] = config.expert_w1
|
||||
if config.expert_w2 is not None:
|
||||
overrides["expert_w2"] = config.expert_w2
|
||||
if config.expert_w3 is not _UNSET:
|
||||
overrides["expert_w3"] = config.expert_w3
|
||||
if config.num_experts_attr is not None:
|
||||
overrides["num_experts_attr"] = config.num_experts_attr
|
||||
if config.top_k_attr is not None:
|
||||
overrides["top_k_attr"] = config.top_k_attr
|
||||
if config.has_shared_experts is not None:
|
||||
overrides["has_shared_experts"] = config.has_shared_experts
|
||||
if config.shared_experts_pattern is not None:
|
||||
overrides["shared_experts_pattern"] = config.shared_experts_pattern
|
||||
if config.shared_experts_gate_pattern is not None:
|
||||
overrides["shared_experts_gate_pattern"] = config.shared_experts_gate_pattern
|
||||
if not overrides:
|
||||
return preset
|
||||
return replace(preset, **overrides)
|
||||
|
||||
|
||||
def resolve_preset_candidates(
|
||||
config: AutoEPConfig,
|
||||
model_config,
|
||||
) -> list[tuple[str, MoEModelPreset]]:
|
||||
"""Resolve ordered preset candidates for AutoEP detection."""
|
||||
if config.preset_model is not None:
|
||||
preset = apply_config_overrides(config, get_preset(config.preset_model))
|
||||
_validate_preset_compatibility(config.preset_model, preset, model_config)
|
||||
return [(config.preset_model, preset)]
|
||||
|
||||
if model_config is not None:
|
||||
model_type = getattr(model_config, 'model_type', None)
|
||||
if model_type:
|
||||
preset_name = preset_name_for_hf_model_type(model_type)
|
||||
if preset_name is not None:
|
||||
logger.info(f"AutoEP: auto-detected model_type='{model_type}', using preset '{preset_name}'")
|
||||
preset = apply_config_overrides(config, get_preset(preset_name))
|
||||
_validate_preset_compatibility(preset_name, preset, model_config)
|
||||
return [(preset_name, preset)]
|
||||
|
||||
unsupported_preset = unsupported_preset_for_hf_model_type(model_type)
|
||||
if unsupported_preset is not None:
|
||||
preset_name, preset = unsupported_preset
|
||||
_validate_preset_compatibility(preset_name, preset, model_config)
|
||||
|
||||
if config.moe_layer_pattern:
|
||||
return [("custom", _build_custom_preset(config))]
|
||||
|
||||
return [(name, apply_config_overrides(config, preset)) for name, preset in PRESET_MODELS.items()]
|
||||
|
||||
|
||||
def _validate_preset_compatibility(
|
||||
preset_name: str,
|
||||
preset: MoEModelPreset,
|
||||
model_config,
|
||||
) -> None:
|
||||
adapter = get_preset_adapter(preset.preset_adapter)
|
||||
adapter.validate_compatibility(preset_name, preset, model_config)
|
||||
|
||||
|
||||
def _build_custom_preset(config: AutoEPConfig) -> MoEModelPreset:
|
||||
return MoEModelPreset(
|
||||
moe_layer_pattern=config.moe_layer_pattern,
|
||||
router_pattern=config.router_pattern or "gate",
|
||||
experts_pattern=config.expert_pattern or "experts",
|
||||
expert_storage="fused_3d",
|
||||
expert_w1=config.expert_w1 or "gate_up_proj",
|
||||
expert_w2=config.expert_w2 or "down_proj",
|
||||
expert_w3=(None if config.expert_w3 is _UNSET else config.expert_w3),
|
||||
num_experts_attr=config.num_experts_attr or "num_local_experts",
|
||||
top_k_attr=config.top_k_attr or "num_experts_per_tok",
|
||||
score_func=(config.score_func if config.score_func != "auto" else "softmax"),
|
||||
score_apply=(config.score_apply if config.score_apply != "auto" else "post"),
|
||||
route_norm=(config.route_norm if config.route_norm is not None else True),
|
||||
gate_bias=False,
|
||||
has_shared_experts=(config.has_shared_experts if config.has_shared_experts is not None else False),
|
||||
shared_experts_pattern=config.shared_experts_pattern or "",
|
||||
shared_experts_gate_pattern=config.shared_experts_gate_pattern or "",
|
||||
)
|
||||
Executable
+660
@@ -0,0 +1,660 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Automatic Tensor Parallelism
|
||||
import re
|
||||
|
||||
from torch import nn
|
||||
from .replace_policy import replace_policies
|
||||
from typing import Optional
|
||||
import torch
|
||||
from deepspeed import comm as dist
|
||||
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.utils import groups
|
||||
from deepspeed.module_inject.layers import is_autotp_training_mode
|
||||
from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType
|
||||
|
||||
|
||||
def move(tensor, device, copy=True):
|
||||
if tensor.is_meta:
|
||||
return torch.empty_like(tensor, device=device)
|
||||
else:
|
||||
# Using new tensors help in freeing memory (after split for example) was done before by calling clone().
|
||||
# Using copy=True instead of clone() will help in case of cpu --> cpu.
|
||||
# Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced.
|
||||
return tensor.to(device, copy=copy)
|
||||
|
||||
|
||||
class ReplaceWithTensorSlicing:
|
||||
|
||||
def __init__(self, mp_group=None, mp_size=1, out_dim=1, in_dim=0):
|
||||
if mp_group is not None:
|
||||
self.gpu_index = dist.get_rank(group=mp_group)
|
||||
else:
|
||||
self.gpu_index = 0
|
||||
self.out_dim = out_dim
|
||||
self.in_dim = in_dim
|
||||
self.mp_size = mp_size
|
||||
|
||||
def merge_assert(self, dim1, dim2):
|
||||
assert dim1 > dim2, \
|
||||
'Merging tensors is not allowed here! Please use deepspeed load_checkpoint\
|
||||
for merging your checkpoints before replacing the transformer layer with\
|
||||
inference-kernels'
|
||||
|
||||
def strided_copy(self,
|
||||
dst: Optional[torch.Tensor],
|
||||
src: Optional[torch.Tensor],
|
||||
num_splits: int,
|
||||
int8: bool = False,
|
||||
allocate_tensor: bool = False):
|
||||
if src is None:
|
||||
return src
|
||||
src_shape = src.shape
|
||||
dst_shape = dst.shape
|
||||
|
||||
outer_dim = 0 if int8 else -1
|
||||
|
||||
if allocate_tensor:
|
||||
dst = torch.empty_like(dst)
|
||||
|
||||
src_split = torch.split(src.data, src.shape[outer_dim] // num_splits, dim=outer_dim)
|
||||
if (len(src_shape) == 2 and len(dst_shape) == 2):
|
||||
if src_shape[outer_dim] == dst_shape[self.out_dim]:
|
||||
try:
|
||||
dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape)
|
||||
except Exception:
|
||||
print(dst.shape, src.shape)
|
||||
exit()
|
||||
dst = torch.nn.parameter.Parameter(dst, requires_grad=False)
|
||||
if hasattr(src, 'scale'):
|
||||
dst.scale = src.scale
|
||||
return dst
|
||||
self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim])
|
||||
qkv_size = dst_shape[self.out_dim] // num_splits
|
||||
qkv_split = [torch.split(src_s, qkv_size, dim=outer_dim) for src_s in src_split]
|
||||
weight_split = [
|
||||
torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=outer_dim) for i in range(len(qkv_split[0]))
|
||||
]
|
||||
dst = dst.reshape(-1).data.copy_(weight_split[self.gpu_index].contiguous().reshape(-1)).reshape(
|
||||
weight_split[self.gpu_index].shape)
|
||||
else:
|
||||
if src_shape[0] == dst_shape[0]:
|
||||
return torch.nn.parameter.Parameter(src)
|
||||
qkv_size = dst_shape[0] // num_splits
|
||||
qkv_split = [torch.split(src_s, qkv_size, dim=0) for src_s in src_split]
|
||||
bias_split = [torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=0) for i in range(len(qkv_split[0]))]
|
||||
dst.data.copy_(bias_split[self.gpu_index].contiguous())
|
||||
|
||||
dst = torch.nn.parameter.Parameter(dst, requires_grad=False)
|
||||
if hasattr(src, 'scale'):
|
||||
dst.scale = src.scale
|
||||
return dst
|
||||
|
||||
def copy(self, dst, src, int8=False, allocate_tensor=False):
|
||||
if src is None:
|
||||
return src
|
||||
assert not dst.data.is_meta # the torch.Tensor.copy_ method used below will silently fail on meta tensors
|
||||
if allocate_tensor:
|
||||
dst = torch.empty_like(dst)
|
||||
outer_dim = 0 if int8 else 1
|
||||
inner_dim = 1 if int8 else 0
|
||||
src_shape = src.shape
|
||||
dst_shape = dst.shape
|
||||
if (len(src_shape) == 2 and len(dst_shape) == 2):
|
||||
|
||||
if src_shape[inner_dim] == dst_shape[self.in_dim] and src_shape[outer_dim] == dst_shape[self.out_dim]:
|
||||
dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape)
|
||||
else:
|
||||
if src_shape[inner_dim] != dst_shape[self.in_dim]:
|
||||
self.merge_assert(src_shape[inner_dim], dst_shape[self.in_dim])
|
||||
dst.data.copy_(src[:, self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim]] if inner_dim == 1 else \
|
||||
src[self.gpu_index * dst_shape[self.in_dim]: (self.gpu_index + 1) * dst_shape[self.in_dim], :])
|
||||
else:
|
||||
self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim])
|
||||
dst.data.copy_(src[:, self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim]] if outer_dim == 1 else \
|
||||
src[self.gpu_index * dst_shape[self.out_dim]: (self.gpu_index + 1) * dst_shape[self.out_dim], :])
|
||||
else:
|
||||
if src_shape[0] == dst_shape[0]:
|
||||
dst = src if src.dtype == dst.dtype else dst.data.copy_(src)
|
||||
else:
|
||||
dst.data.copy_(src[self.gpu_index * dst_shape[-1]:(self.gpu_index + 1) * dst_shape[-1]])
|
||||
dst = torch.nn.parameter.Parameter(dst, requires_grad=False)
|
||||
if hasattr(src, 'scale'):
|
||||
dst.scale = src.scale
|
||||
return dst
|
||||
|
||||
|
||||
class Loading():
|
||||
|
||||
def is_load_module(module):
|
||||
load_layers = [nn.Linear, nn.Embedding, nn.LayerNorm]
|
||||
load_layer_names = [
|
||||
"LPLayerNorm", "SharedEmbedding", "OPTLearnedPositionalEmbedding", "LlamaRMSNorm", "FalconLinear",
|
||||
"MistralRMSNorm", "T5LayerNorm", "MixtralRMSNorm", "Phi3RotaryEmbedding", "Phi3SuScaledRotaryEmbedding",
|
||||
"Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm",
|
||||
"Qwen3RMSNorm", "Qwen3MoeRMSNorm", "DeepseekV2RMSNorm", "DeepseekV3RMSNorm",
|
||||
"DeepseekV2YarnRotaryEmbedding", "DeepseekV3YarnRotaryEmbedding", "MoEGate"
|
||||
]
|
||||
return module.__class__ in load_layers or module._get_name() in load_layer_names
|
||||
|
||||
def load_buffer(module, state_dict, prefix):
|
||||
for name in module._buffers.keys():
|
||||
if module._buffers[name].data.is_meta:
|
||||
module._buffers[name] = torch.nn.parameter.Parameter(
|
||||
data=torch.empty_like(module._buffers[name].data, device="cpu"),
|
||||
requires_grad=module._buffers[name].data.requires_grad)
|
||||
if prefix + name in state_dict.keys():
|
||||
module._buffers[name].data.copy_(state_dict[prefix + name])
|
||||
|
||||
def load(module, state_dict, prefix, mp_group=None):
|
||||
mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group)
|
||||
if hasattr(module, 'weight'):
|
||||
if module.weight.data.is_meta:
|
||||
# meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
|
||||
module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, device="cpu"),
|
||||
requires_grad=module.weight.data.requires_grad)
|
||||
if 'query_key_value' in prefix:
|
||||
module.weight = mp_replace.strided_copy(module.weight.data,
|
||||
state_dict[prefix + 'weight'],
|
||||
num_splits=3)
|
||||
else:
|
||||
module.weight = mp_replace.copy(module.weight.data, state_dict[prefix + 'weight'])
|
||||
else:
|
||||
if hasattr(module, 'norm') and hasattr(module.norm, 'weight'):
|
||||
if module.norm.weight.data.is_meta:
|
||||
# meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
|
||||
module.norm.weight = torch.nn.parameter.Parameter(
|
||||
data=torch.empty_like(module.norm.weight.data, device="cpu"),
|
||||
requires_grad=module.norm.weight.data.requires_grad)
|
||||
module.norm.weight = mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight'])
|
||||
|
||||
if prefix + 'bias' in state_dict.keys():
|
||||
if hasattr(module, 'bias'):
|
||||
if module.bias.data.is_meta:
|
||||
# meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
|
||||
module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"),
|
||||
requires_grad=module.bias.data.requires_grad)
|
||||
module.bias = mp_replace.copy(module.bias, state_dict[prefix + 'bias'])
|
||||
else:
|
||||
if hasattr(module, 'norm') and hasattr(module.norm, 'bias'):
|
||||
if module.norm.bias.data.is_meta:
|
||||
# meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
|
||||
module.norm.bias = torch.nn.parameter.Parameter(
|
||||
data=torch.empty_like(module.norm.bias.data, device="cpu"),
|
||||
requires_grad=module.norm.bias.data.requires_grad)
|
||||
module.norm.bias = mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias'])
|
||||
|
||||
|
||||
class AutoTP():
|
||||
|
||||
def __init__(self,
|
||||
module,
|
||||
all_reduce_linears,
|
||||
prefix,
|
||||
state_dict,
|
||||
linear_layer_setting,
|
||||
orig_layer_impl,
|
||||
keep_module_on_host=False,
|
||||
partition_config: Optional[AutoTPConfig] = None):
|
||||
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
|
||||
self.linear_layer_setting = linear_layer_setting
|
||||
self.orig_layer_impl = orig_layer_impl
|
||||
self.linear_policies = None
|
||||
self.conv_linear_layer = False
|
||||
self.partition_config = partition_config
|
||||
TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host)
|
||||
|
||||
def in_module_list(module, module_list):
|
||||
for item in module_list:
|
||||
if type(item).__name__ == type(module).__name__:
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_module_list(model):
|
||||
mlist = []
|
||||
for child in model.children():
|
||||
if isinstance(child, nn.ModuleList):
|
||||
for module in child.children():
|
||||
if not mlist:
|
||||
mlist = [module]
|
||||
elif not AutoTP.in_module_list(module, mlist):
|
||||
mlist = mlist + [module]
|
||||
else:
|
||||
mlist = mlist + AutoTP.get_module_list(child)
|
||||
return mlist
|
||||
|
||||
def supported(model):
|
||||
unsupported = ['deberta', 'flaubert', 'fsmt', 'gpt2', 'led', 'longformer', 'xlm', 'xlnet']
|
||||
model = str(model)
|
||||
key = re.search(r": (.*?)Model", model)
|
||||
if key is None:
|
||||
key = re.search(r": (.*?)Stack", model)
|
||||
if key is None:
|
||||
key = re.match(r"(.*?)Model", model)
|
||||
assert key is not None, "Not able to determine model policy automatically. Please provide policy."
|
||||
if key.group(1).lower() in unsupported:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_layers(parent, module):
|
||||
layer_list = []
|
||||
for key, submodule in module._modules.items():
|
||||
if isinstance(submodule, nn.Linear):
|
||||
layer_list = layer_list + [parent + "." + key]
|
||||
elif isinstance(submodule, nn.LayerNorm) or key == 'LayerNorm' or key == 'layer_norm':
|
||||
layer_list = layer_list + ["ln"]
|
||||
else:
|
||||
layer_list = layer_list + AutoTP.get_layers(key, submodule)
|
||||
return layer_list
|
||||
|
||||
def update_policy_list(policy_list, new_module, new_gems):
|
||||
if len(policy_list):
|
||||
for i, policy in enumerate(policy_list):
|
||||
# if module already exists in policy, combine gems and remove duplicates
|
||||
if policy[0] == type(new_module):
|
||||
new_gems = set(new_gems + policy[1])
|
||||
policy_list[i] = tuple([type(new_module), new_gems])
|
||||
return policy_list
|
||||
policy_list.append(tuple([type(new_module), new_gems]))
|
||||
return policy_list
|
||||
|
||||
def kernel_supported(module_list):
|
||||
policy = []
|
||||
for plcy in replace_policies:
|
||||
# instantiate a throw-away policy in order to populate the _orig_layer_class
|
||||
_ = plcy(None)
|
||||
if isinstance(plcy._orig_layer_class, list):
|
||||
for orig_layer_class in plcy._orig_layer_class:
|
||||
policy.append(orig_layer_class)
|
||||
elif plcy._orig_layer_class is not None:
|
||||
policy.append(plcy._orig_layer_class)
|
||||
for child in module_list:
|
||||
if child.__class__ in policy:
|
||||
return True
|
||||
return False
|
||||
|
||||
def tp_parser(model):
|
||||
policy_list = []
|
||||
module_list = []
|
||||
layer_list = []
|
||||
gem_list = []
|
||||
|
||||
module_list = AutoTP.get_module_list(model)
|
||||
assert AutoTP.supported(model), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \
|
||||
if AutoTP.kernel_supported(module_list) else "AutoTP not supported for model. Please provide policy."
|
||||
norm_layer_name_list = ['LayerNorm', 'layer_norm', 'ln_1', 'ln_2']
|
||||
#ln_1 , ln_2 for Qwen
|
||||
for module in module_list:
|
||||
for key, submodule in module._modules.items():
|
||||
if isinstance(submodule, nn.Linear):
|
||||
layer_list = layer_list + ["." + key]
|
||||
elif isinstance(submodule, nn.LayerNorm) or key in norm_layer_name_list:
|
||||
layer_list = layer_list + ["ln"]
|
||||
else:
|
||||
layer_list = layer_list + AutoTP.get_layers(key, submodule)
|
||||
for i, layer in enumerate(layer_list):
|
||||
if layer == 'ln':
|
||||
if layer_list[i - 1] != 'ln':
|
||||
gem_list = gem_list + [layer_list[i - 1]]
|
||||
elif 'out_proj' in layer:
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'o_proj' in layer:
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'down_proj' in layer:
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'attention.dense' in layer and 'GPTNeoX' in str(model):
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'self_attention.dense' in layer and 'falcon' in str(
|
||||
type(module)): # this is a hack to get the right linear layer for this model!
|
||||
gem_list = gem_list + [layer]
|
||||
# Mixtral-7x8b used w2*act(w1*w3) linear. need to replace w2 to linearallreduce.
|
||||
elif 'w2' in layer and 'Mixtral' in str(type(module)):
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'self_attn.dense' in layer and 'Phi' in str(type(module)):
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'self_attention.dense' in layer and 'ChatGLM' in str(model):
|
||||
gem_list = gem_list + [layer]
|
||||
elif 'dense_4h_to_h' in layer and 'ChatGLM' in str(model):
|
||||
gem_list = gem_list + [layer]
|
||||
|
||||
layer_list = []
|
||||
if gem_list != []:
|
||||
gem_list = list(set(gem_list))
|
||||
policy_list = AutoTP.update_policy_list(policy_list, module, gem_list)
|
||||
gem_list = []
|
||||
assert len(policy_list), "AutoTP not supported for model. Please use kernel injection since container policy for model exists." \
|
||||
if AutoTP.kernel_supported(module_list) else "Not able to determine model policy automatically. Please provide policy."
|
||||
return policy_list
|
||||
|
||||
def set_tensor_parallel_config(self, mp_size, mp_group):
|
||||
|
||||
if is_autotp_training_mode():
|
||||
self.mp_group = groups.get_tensor_model_parallel_group()
|
||||
self.mp_size = groups.get_tensor_model_parallel_world_size()
|
||||
return
|
||||
|
||||
self.mp_size = mp_size
|
||||
self.mp_group = mp_group
|
||||
|
||||
def _replace(self, child, name, conv_linear_layer):
|
||||
# This function should clearly define the routing rules for specific layers
|
||||
# and avoid any complex shard-related logic.
|
||||
if getattr(child, "replaced", False) == True:
|
||||
return
|
||||
|
||||
# Skip AutoEP-managed modules (expert weights are EP-sharded, not TP-sharded)
|
||||
if getattr(child, "_is_autoep_layer", False):
|
||||
return child
|
||||
|
||||
weight_shape = child.weight.shape
|
||||
mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
|
||||
|
||||
# If partition_config is provided, use the new configurable API
|
||||
if self.partition_config is not None:
|
||||
return self._replace_with_config(child, name)
|
||||
|
||||
# For TP layer skip, e.g., MoE gate, deepseek low rank layer skip
|
||||
if "mlp.gate" == name or "q_a_proj" in name or "kv_a_proj_with_mqa" in name or name == "block_sparse_moe.gate" or (
|
||||
('mlp.shared_expert_gate' == name or 'mlp.gate' == name) and 'qwen2_moe' in str(type(self.module))):
|
||||
return child
|
||||
# For Yuan model
|
||||
if 'Yuan' in str(self.module):
|
||||
if 'v_proj' in name:
|
||||
return Yuan_LinearLayer(child, self.mp_group)
|
||||
|
||||
elif 'o_proj' in name:
|
||||
return Yuan_LinearAllreduce(child, self.mp_group)
|
||||
|
||||
# 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)
|
||||
# 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:
|
||||
arctic_w2_all_reduce_linear = True
|
||||
# For MoE MLP model, e.g., deepseek and jamba
|
||||
down_proj = False
|
||||
if 'down_proj' in name:
|
||||
down_proj = True
|
||||
if name in self.all_reduce_linears or arctic_w2_all_reduce_linear or down_proj:
|
||||
|
||||
setattr(child, "replaced", True)
|
||||
if self.conv_linear_layer:
|
||||
return Conv_LinearALlreduce(child, self.mp_group, name=name)
|
||||
elif name == "lm_head" or name == 'embed_out':
|
||||
return LmHeadLinearAllreduce(child, self.mp_group)
|
||||
|
||||
return LinearAllreduce(child, self.mp_group, name=name)
|
||||
else:
|
||||
|
||||
setattr(child, "replaced", True)
|
||||
if self.conv_linear_layer:
|
||||
conv_LinearLayer(child, self.mp_group)
|
||||
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 LinearLayer(child, self.mp_group, name=name)
|
||||
|
||||
def _replace_with_config(self, child, name):
|
||||
"""
|
||||
Replace layer using the new configurable AutoTP API.
|
||||
|
||||
This method uses TPLayerSpec to determine how to partition the layer.
|
||||
"""
|
||||
if getattr(child, "replaced", False) == True:
|
||||
return child
|
||||
|
||||
# Build the full parameter name for pattern matching
|
||||
param_name = name + ".weight" if not name.endswith(".weight") else name
|
||||
|
||||
# Find matching spec
|
||||
model_type = self._get_model_type()
|
||||
spec = self.partition_config.find_matching_spec(param_name, model_type)
|
||||
|
||||
if spec is None:
|
||||
# No matching spec found
|
||||
if self.partition_config.strict_mode:
|
||||
raise ValueError(f"No matching spec for {param_name}")
|
||||
# With partition_config, rely only on explicit specs and skip unmatched layers.
|
||||
return child
|
||||
|
||||
setattr(child, "replaced", True)
|
||||
|
||||
if spec.partition_type == PartitionType.SKIP:
|
||||
return child
|
||||
|
||||
if spec.partition_type == PartitionType.ROW:
|
||||
return self._create_row_parallel_layer(child, spec, name)
|
||||
else:
|
||||
return self._create_column_parallel_layer(child, spec, 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)
|
||||
# Check for lm_head / embed_out
|
||||
if name == "lm_head" or name == 'embed_out':
|
||||
return LmHeadLinearAllreduce(module, self.mp_group)
|
||||
|
||||
if spec.shape is not None:
|
||||
return SubParamLinearAllreduce(
|
||||
module,
|
||||
self.mp_group,
|
||||
shape=spec.shape,
|
||||
partition_dim=spec.get_partition_dim(),
|
||||
name=name,
|
||||
)
|
||||
return LinearAllreduce(module, self.mp_group, name=name)
|
||||
|
||||
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)
|
||||
# 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)
|
||||
if spec.shape is not None:
|
||||
return SubParamLinearLayer(
|
||||
module,
|
||||
self.mp_group,
|
||||
shape=spec.shape,
|
||||
partition_dim=spec.get_partition_dim(),
|
||||
name=name,
|
||||
)
|
||||
return LinearLayer(module, self.mp_group, name=name)
|
||||
|
||||
def _get_model_type(self) -> Optional[str]:
|
||||
"""Extract model type from module config or class name."""
|
||||
config = getattr(self.module, "config", None)
|
||||
if config is not None:
|
||||
model_type = getattr(config, "model_type", None)
|
||||
if model_type:
|
||||
return str(model_type).lower()
|
||||
module_str = str(type(self.module))
|
||||
# Try to extract model type from class name (e.g., "LlamaDecoderLayer" -> "llama")
|
||||
patterns = [
|
||||
r"(\w+)DecoderLayer",
|
||||
r"(\w+)Block",
|
||||
r"(\w+)Layer",
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, module_str)
|
||||
if match:
|
||||
return match.group(1).lower()
|
||||
return None
|
||||
|
||||
def _slice_embedding(self, child, name, conv_linear_layer):
|
||||
if getattr(child, "replaced", False) == True:
|
||||
return
|
||||
|
||||
mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
|
||||
|
||||
if hasattr(child.weight, 'ds_tensor'):
|
||||
data = child.weight.ds_tensor.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size), dim=1)
|
||||
else:
|
||||
data = child.weight.data.split(get_shard_size_list(child.weight.shape[1], self.mp_size, name), dim=1)
|
||||
data = data[mp_replace.gpu_index].to(get_accelerator().current_device_name())
|
||||
data = torch.nn.parameter.Parameter(data, requires_grad=False)
|
||||
|
||||
new_embedding = nn.Embedding(child.weight.shape[0], get_shard_size(child.weight.shape[1], self.mp_size, name))
|
||||
new_embedding.weight.data.copy_(data)
|
||||
setattr(child, "replaced", True)
|
||||
return new_embedding
|
||||
|
||||
def update_mp_params(self, child):
|
||||
if getattr(child, "replaced", False) == True:
|
||||
return
|
||||
param_list = [
|
||||
"n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", "all_head_size",
|
||||
"embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads", "kv_n_heads", "d_model",
|
||||
"num_attention_heads_per_partition", "num_multi_query_groups_per_partition", "hidden_size_per_partition"
|
||||
]
|
||||
for param in param_list:
|
||||
if "Yuan" in str(child) and 'embed_dim' in param_list:
|
||||
param_list.remove('embed_dim')
|
||||
if hasattr(child, param):
|
||||
param_val = getattr(child, param)
|
||||
setattr(child, param, get_shard_size(param_val, self.mp_size))
|
||||
setattr(child, "replaced", True)
|
||||
|
||||
def update_linear_policies(self):
|
||||
self.conv_linear_layer = False
|
||||
if self.linear_layer_setting is not None:
|
||||
self.linear_policies = {self.linear_layer_setting[0]: self._replace}
|
||||
if len(self.linear_layer_setting) == 2:
|
||||
self.linear_policies.update({self.linear_layer_setting[1]: self._slice_embedding})
|
||||
else:
|
||||
import transformers
|
||||
if self.orig_layer_impl is transformers.models.gpt2.modeling_gpt2.GPT2Block:
|
||||
try:
|
||||
self.conv_linear_layer = True
|
||||
self.linear_policies = {transformers.pytorch_utils.Conv1D: self._replace}
|
||||
except ImportError:
|
||||
self.linear_policies = {nn.Linear: self._replace}
|
||||
else:
|
||||
self.linear_policies = {nn.Linear: self._replace, nn.Embedding: self._slice_embedding}
|
||||
|
||||
def _replace_autoep_shared_experts(self, autoep_layer, autoep_name):
|
||||
for child_name in ("shared_experts", "shared_experts_gate"):
|
||||
child = getattr(autoep_layer, child_name, None)
|
||||
if child is None:
|
||||
continue
|
||||
full_name = f"{autoep_name}.{child_name}" if autoep_name else child_name
|
||||
if self.partition_config is not None and hasattr(child, "weight") and getattr(
|
||||
child.weight, "dim", lambda: 0)() == 2:
|
||||
new_child = self._replace_with_config(child, full_name)
|
||||
if new_child is not None:
|
||||
setattr(autoep_layer, child_name, new_child)
|
||||
elif child.__class__ in self.linear_policies:
|
||||
setattr(autoep_layer, child_name, self.linear_policies[child.__class__](child, full_name,
|
||||
self.conv_linear_layer))
|
||||
elif any(isinstance(child, lp) for lp in self.linear_policies):
|
||||
key = next(lp for lp in self.linear_policies if isinstance(child, lp))
|
||||
setattr(autoep_layer, child_name, self.linear_policies[key](child, full_name, self.conv_linear_layer))
|
||||
else:
|
||||
self.update_mp_params(child)
|
||||
self._replace_module(child, full_name, "")
|
||||
|
||||
def _replace_module(self, r_module, prev_name='', prev_class_name=''):
|
||||
for name, child in r_module.named_children():
|
||||
if getattr(child, "_is_autoep_layer", False):
|
||||
full_name = prev_name + '.' + name if prev_name else name
|
||||
self._replace_autoep_shared_experts(child, full_name)
|
||||
continue
|
||||
|
||||
if prev_class_name == "":
|
||||
class_name = prev_name
|
||||
elif prev_name == "":
|
||||
class_name = prev_class_name
|
||||
else:
|
||||
class_name = prev_class_name + '.' + prev_name
|
||||
checking_key = self.prefix + '.' + class_name + '.' + name + '.' if class_name != "" else self.prefix + '.' + name + '.'
|
||||
if Loading.is_load_module(child) and self.state_dict is not None:
|
||||
if any(checking_key in item for item in self.state_dict):
|
||||
Loading.load(child, self.state_dict, checking_key, self.mp_group)
|
||||
else:
|
||||
continue
|
||||
if len(child._buffers) != 0 and self.state_dict is not None:
|
||||
Loading.load_buffer(child, self.state_dict, checking_key)
|
||||
|
||||
# When using partition_config (custom patterns/presets), use pattern-based routing
|
||||
# instead of linear_policies. This keeps all pattern logic centralized here.
|
||||
if self.partition_config is not None:
|
||||
full_name = class_name + '.' + name if class_name else name
|
||||
if isinstance(child, nn.Embedding):
|
||||
# Check if embedding matches any pattern
|
||||
param_name = full_name + ".weight"
|
||||
model_type = self._get_model_type()
|
||||
spec = self.partition_config.find_matching_spec(param_name, model_type)
|
||||
if spec is not None and spec.partition_type != PartitionType.SKIP:
|
||||
new_child = self._slice_embedding(child, full_name, False)
|
||||
if new_child is not None:
|
||||
setattr(r_module, name, new_child)
|
||||
# If no pattern matched or skip, leave embedding unchanged
|
||||
elif hasattr(child, "weight") and getattr(child.weight, "dim", lambda: 0)() == 2:
|
||||
new_child = self._replace_with_config(child, full_name)
|
||||
if new_child is not None:
|
||||
setattr(r_module, name, new_child)
|
||||
else:
|
||||
self.update_mp_params(child)
|
||||
self._replace_module(child, name, class_name)
|
||||
# Traditional path: use linear_policies for type-based routing
|
||||
elif child.__class__ in self.linear_policies:
|
||||
setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name,
|
||||
self.conv_linear_layer))
|
||||
elif any(isinstance(child, lp) for lp in self.linear_policies):
|
||||
# Added for falcon model support
|
||||
# Note: isinstance will account for class inheritance, child.__class__ does not
|
||||
key = None
|
||||
for lp in self.linear_policies:
|
||||
if isinstance(child, lp):
|
||||
key = lp
|
||||
break
|
||||
assert key is not None
|
||||
setattr(r_module, name, self.linear_policies[key](child, prev_name + '.' + name,
|
||||
self.conv_linear_layer))
|
||||
else:
|
||||
self.update_mp_params(child)
|
||||
self._replace_module(child, name, class_name)
|
||||
return r_module
|
||||
|
||||
def get_model_num_kv_heads(self, 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"
|
||||
child = r_module.lm_head
|
||||
elif hasattr(r_module, "embed_out"):
|
||||
name = "embed_out"
|
||||
child = r_module.embed_out
|
||||
else:
|
||||
return r_module
|
||||
if child.__class__ in self.linear_policies:
|
||||
setattr(r_module, name, self.linear_policies[child.__class__](child, name, self.conv_linear_layer))
|
||||
return r_module
|
||||
@@ -0,0 +1,104 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed import comm as dist
|
||||
import torch
|
||||
from typing import Optional
|
||||
from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
|
||||
|
||||
|
||||
def build_bloom_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
|
||||
"""
|
||||
Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it
|
||||
relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value
|
||||
`softmax(l+a) = softmax(l)`. Based on
|
||||
https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742
|
||||
TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly.
|
||||
|
||||
Args:
|
||||
Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)
|
||||
attention_mask (`torch.Tensor`):
|
||||
Token-wise attention mask, this should be of shape (batch_size, max_seq_len).
|
||||
num_heads (`int`, *required*):
|
||||
number of heads
|
||||
dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`):
|
||||
dtype of the output tensor
|
||||
"""
|
||||
import math
|
||||
batch_size, seq_length = attention_mask.shape
|
||||
closest_power_of_2 = 2**math.floor(math.log2(num_heads))
|
||||
base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))),
|
||||
device=attention_mask.device,
|
||||
dtype=torch.float32)
|
||||
powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
|
||||
slopes = torch.pow(base, powers)
|
||||
|
||||
if closest_power_of_2 != num_heads:
|
||||
extra_base = torch.tensor(2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
|
||||
device=attention_mask.device,
|
||||
dtype=torch.float32)
|
||||
num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
|
||||
extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
|
||||
slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
|
||||
|
||||
# Note: alibi will added to the attention bias that will be applied to the query, key product of attention
|
||||
# => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)
|
||||
# => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)
|
||||
# => the query_length dimension will then be broadcasted correctly
|
||||
# This is more or less identical to T5's relative position bias:
|
||||
# https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527
|
||||
arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
|
||||
alibi = slopes[..., None] * arange_tensor
|
||||
if dist.is_initialized():
|
||||
num_heads_per_rank = get_shard_size(num_heads, dist.get_world_size())
|
||||
offset = sum(get_shard_size_list(num_heads, dist.get_world_size())[0:dist.get_rank()])
|
||||
alibi = alibi.view(batch_size, num_heads, 1, seq_length)
|
||||
alibi = alibi[:, offset:num_heads_per_rank + offset, :, :]
|
||||
return alibi.reshape(batch_size * num_heads_per_rank, 1, seq_length).to(dtype)
|
||||
else:
|
||||
return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)
|
||||
|
||||
|
||||
def get_alibi_mask(self, tensor, seq_length_with_past):
|
||||
mask = self.get_alibi_mask_orig(tensor, seq_length_with_past)
|
||||
if not self.training and dist.is_initialized():
|
||||
num_heads_per_rank = get_shard_size(self.n_head, dist.get_world_size())
|
||||
offset = sum(get_shard_size_list(self.n_head, dist.get_world_size())[0:dist.get_rank()])
|
||||
mask = mask[offset:num_heads_per_rank + offset, :seq_length_with_past, :seq_length_with_past]
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def build_mpt_atten_bias_tensor(self,
|
||||
device,
|
||||
dtype,
|
||||
attention_mask: Optional[torch.ByteTensor] = None,
|
||||
prefix_mask: Optional[torch.ByteTensor] = None,
|
||||
sequence_id: Optional[torch.LongTensor] = None):
|
||||
(attn_bias, attention_mask) = self._attn_bias_orig(device,
|
||||
dtype,
|
||||
attention_mask=attention_mask,
|
||||
prefix_mask=prefix_mask,
|
||||
sequence_id=sequence_id)
|
||||
if dist.is_initialized():
|
||||
num_heads_per_rank = get_shard_size(self.config.n_heads, dist.get_world_size())
|
||||
offset = sum(get_shard_size_list(self.config.n_heads, dist.get_world_size())[0:dist.get_rank()])
|
||||
attn_bias = attn_bias[:, offset:num_heads_per_rank + offset, :, :]
|
||||
return attn_bias, attention_mask
|
||||
|
||||
|
||||
def build_mpt_alibi_tensor(self, num_heads, sequence_length, alibi_bias_max=8, device=None) -> torch.Tensor:
|
||||
r"""
|
||||
Link to paper: https://arxiv.org/abs/2108.12409 - Alibi tensor is not causal as the original paper mentions, it
|
||||
relies on a translation invariance of softmax for quick implementation. This implementation has been copied from
|
||||
the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi:
|
||||
https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292
|
||||
"""
|
||||
alibi = self.build_mpt_alibi_tensor_orig(num_heads, sequence_length, alibi_bias_max, device)
|
||||
if dist.is_initialized():
|
||||
num_heads_per_rank = int(num_heads / dist.get_world_size())
|
||||
offset = dist.get_rank() * num_heads_per_rank
|
||||
alibi = alibi[offset:num_heads_per_rank + offset, :, :]
|
||||
return alibi
|
||||
@@ -0,0 +1,569 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""
|
||||
Configurable AutoTP API
|
||||
|
||||
This module provides a unified specification for tensor parallel layer partitioning.
|
||||
The design is inspired by Universal Checkpointing's SubparamShape and provides
|
||||
a single, well-defined format that users can easily understand, customize, and extend.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Tuple, Union, Optional
|
||||
from enum import Enum
|
||||
from deepspeed.utils.logging import warning_once
|
||||
|
||||
|
||||
class PartitionType(Enum):
|
||||
"""How the layer should be partitioned for tensor parallelism."""
|
||||
COLUMN = "column" # Partition output dim, AllReduce in backward
|
||||
ROW = "row" # Partition input dim, AllReduce in forward
|
||||
SKIP = "skip" # Do not partition this layer
|
||||
|
||||
|
||||
@dataclass
|
||||
class TPLayerSpec:
|
||||
"""
|
||||
Unified specification for tensor parallel layer partitioning.
|
||||
|
||||
This is inspired by Universal Checkpointing's SubparamShape but extended
|
||||
for AutoTP's needs (forward/backward communication patterns).
|
||||
|
||||
The `shape` parameter supports at most 1-level nesting at the partition dimension:
|
||||
- (3, -1) -> 3 equal-size sub-params
|
||||
- ((q, k, v), -1) -> 3 unequal-size sub-params (1-level nesting)
|
||||
|
||||
Examples:
|
||||
# Simple row-parallel layer (e.g., o_proj, down_proj)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.o_proj$", ".*\\.down_proj$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
)
|
||||
|
||||
# Simple column-parallel layer (e.g., q_proj, k_proj, v_proj)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.[qkv]_proj$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
)
|
||||
|
||||
# Fused QKV - GLM style [Q, K, V] concatenated on dim 0
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.query_key_value\\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(3, -1), # 3 equal sub-params, -1 = infer
|
||||
partition_dim=0,
|
||||
)
|
||||
|
||||
# Fused QKV - Bloom style [q1,k1,v1,q2,k2,v2,...]
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.query_key_value\\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
# No reshape needed, just split along dim 0
|
||||
)
|
||||
|
||||
# GQA with different Q/K/V sizes (1-level nesting)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.qkv_proj\\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=((q_size, k_size, v_size), -1), # Unequal sub-params
|
||||
partition_dim=0,
|
||||
)
|
||||
|
||||
# Chunked MLP (gate_up_proj)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.gate_up_proj\\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(2, -1), # [gate, up] packed
|
||||
partition_dim=0,
|
||||
)
|
||||
|
||||
# MoE FFN with expert dimension
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.experts\\..*\\.w1\\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(num_experts, -1, hidden_in), # View as 3D
|
||||
partition_dim=1, # Partition the hidden_out dimension
|
||||
)
|
||||
|
||||
# Skip layer (e.g., MoE gate)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.gate$", ".*\\.router$"],
|
||||
partition_type=PartitionType.SKIP,
|
||||
)
|
||||
"""
|
||||
|
||||
# Layer identification - regex patterns to match parameter names
|
||||
patterns: List[str]
|
||||
|
||||
# Partition type determines communication pattern
|
||||
partition_type: PartitionType = PartitionType.COLUMN
|
||||
|
||||
# Optional: logical shape for partitioning
|
||||
# - Use -1 for dimensions that should be inferred
|
||||
# - Use tuple of ints at partition_dim for unequal sub-params (1-level nesting only)
|
||||
# Examples:
|
||||
# (3, -1) -> 3 equal sub-params
|
||||
# ((4096, 1024, 1024), -1) -> 3 unequal sub-params (GQA)
|
||||
# (n_experts, -1, hidden) -> MoE reshape
|
||||
shape: Optional[Tuple[Union[int, Tuple[int, ...]], ...]] = None
|
||||
|
||||
# Which dimension to partition (after optional reshape)
|
||||
# Default: 0 for COLUMN, 1 for ROW (standard 2D weight matrix)
|
||||
partition_dim: Optional[int] = None
|
||||
|
||||
# Optional: model type constraint (only apply for specific models)
|
||||
model_types: Optional[List[str]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if isinstance(self.partition_type, str):
|
||||
self.partition_type = PartitionType(self.partition_type.lower())
|
||||
if self.shape is not None:
|
||||
self.shape = self._normalize_shape(self.shape)
|
||||
self._validate_shape_format()
|
||||
|
||||
@staticmethod
|
||||
def _normalize_shape(shape):
|
||||
if isinstance(shape, list):
|
||||
return tuple(TPLayerSpec._normalize_shape(item) for item in shape)
|
||||
if isinstance(shape, tuple):
|
||||
return tuple(TPLayerSpec._normalize_shape(item) if isinstance(item, list) else item for item in shape)
|
||||
return shape
|
||||
|
||||
def _validate_shape_format(self):
|
||||
if not isinstance(self.shape, tuple):
|
||||
raise ValueError("AutoTP shape must be a tuple of ints or a tuple at partition_dim.")
|
||||
partition_dim = self.get_partition_dim()
|
||||
if partition_dim < 0 or partition_dim >= len(self.shape):
|
||||
raise ValueError(
|
||||
f"AutoTP partition_dim {partition_dim} is out of range for shape length {len(self.shape)}.")
|
||||
nested_tuple_seen = False
|
||||
for idx, dim in enumerate(self.shape):
|
||||
if isinstance(dim, tuple):
|
||||
if idx != partition_dim:
|
||||
raise ValueError(
|
||||
f"AutoTP shape nested tuple only allowed at partition_dim={partition_dim}, got at {idx}.")
|
||||
if nested_tuple_seen:
|
||||
raise ValueError("AutoTP shape supports only 1-level nesting at partition_dim.")
|
||||
nested_tuple_seen = True
|
||||
if len(dim) == 0:
|
||||
raise ValueError("AutoTP shape nested tuple cannot be empty.")
|
||||
for val in dim:
|
||||
if isinstance(val, tuple):
|
||||
raise ValueError("AutoTP shape supports only 1-level nesting at partition_dim.")
|
||||
if not isinstance(val, int) or val <= 0:
|
||||
raise ValueError("AutoTP nested sub-parameter sizes must be positive integers.")
|
||||
elif isinstance(dim, int):
|
||||
if dim == 0 or dim < -1:
|
||||
raise ValueError("AutoTP shape dimensions must be positive integers or -1.")
|
||||
else:
|
||||
raise ValueError("AutoTP shape must contain only integers or a tuple at partition_dim.")
|
||||
|
||||
def get_partition_dim(self) -> int:
|
||||
"""Get effective partition dimension."""
|
||||
if self.partition_dim is not None:
|
||||
return self.partition_dim
|
||||
# Default based on partition type for 2D weight matrices
|
||||
return 0 if self.partition_type == PartitionType.COLUMN else 1
|
||||
|
||||
def has_unequal_sub_params(self) -> bool:
|
||||
"""Check if this spec has unequal sub-parameters (nested tuple at partition_dim)."""
|
||||
if self.shape is None:
|
||||
return False
|
||||
dim = self.get_partition_dim()
|
||||
if dim >= len(self.shape):
|
||||
return False
|
||||
return isinstance(self.shape[dim], tuple)
|
||||
|
||||
def get_sub_param_sizes(self) -> Optional[Tuple[int, ...]]:
|
||||
"""Get sub-parameter sizes if using unequal sub-params."""
|
||||
if not self.has_unequal_sub_params():
|
||||
return None
|
||||
return self.shape[self.get_partition_dim()]
|
||||
|
||||
def get_num_sub_params(self) -> Optional[int]:
|
||||
"""Get the number of sub-parameters."""
|
||||
if self.shape is None:
|
||||
return None
|
||||
dim = self.get_partition_dim()
|
||||
if dim >= len(self.shape):
|
||||
return None
|
||||
if isinstance(self.shape[dim], tuple):
|
||||
return len(self.shape[dim])
|
||||
elif isinstance(self.shape[dim], int) and self.shape[dim] > 0:
|
||||
return self.shape[dim]
|
||||
return None
|
||||
|
||||
def matches(self, param_name: str, model_type: Optional[str] = None) -> bool:
|
||||
"""Check if this spec matches the given parameter."""
|
||||
# Check model type constraint
|
||||
if self.model_types:
|
||||
if model_type is None:
|
||||
return False
|
||||
model_type_norm = str(model_type).lower()
|
||||
model_types_norm = [str(mt).lower() for mt in self.model_types]
|
||||
if model_type_norm not in model_types_norm:
|
||||
return False
|
||||
# Check pattern match
|
||||
return any(re.match(pattern, param_name) for pattern in self.patterns)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoTPConfig:
|
||||
"""
|
||||
Configuration for Automatic Tensor Parallelism.
|
||||
|
||||
Example usage:
|
||||
config = AutoTPConfig(
|
||||
tp_size=4,
|
||||
layer_specs=[
|
||||
# Row-parallel layers (AllReduce after forward)
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.o_proj", ".*\\.down_proj"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
# Column-parallel layers
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.[qkv]_proj", ".*\\.up_proj", ".*\\.gate_proj"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
# Skip MoE gates
|
||||
TPLayerSpec(
|
||||
patterns=[".*\\.gate$"],
|
||||
partition_type=PartitionType.SKIP,
|
||||
),
|
||||
],
|
||||
)
|
||||
"""
|
||||
|
||||
tp_size: int = 1
|
||||
|
||||
# Unified layer specifications
|
||||
layer_specs: List[TPLayerSpec] = field(default_factory=list)
|
||||
|
||||
# Embedding configuration
|
||||
embedding_partition_dim: int = 1 # Usually partition vocab dim
|
||||
|
||||
# LM head configuration
|
||||
lm_head_patterns: List[str] = field(default_factory=lambda: ["lm_head", "embed_out"])
|
||||
|
||||
# Behavior flags
|
||||
use_default_specs: bool = True # Merge with built-in specs
|
||||
strict_mode: bool = False # Fail if unmatched Linear layers found
|
||||
|
||||
def find_matching_spec(self, param_name: str, model_type: Optional[str] = None) -> Optional[TPLayerSpec]:
|
||||
"""Find the first matching spec for a parameter."""
|
||||
matches = [spec for spec in self.layer_specs if spec.matches(param_name, model_type)]
|
||||
if not matches:
|
||||
return None
|
||||
if len(matches) > 1:
|
||||
matched_patterns = [spec.patterns for spec in matches]
|
||||
warning_once(f"AutoTPConfig: parameter {param_name} matched multiple layer_specs {matched_patterns}; "
|
||||
"using the first match.")
|
||||
return matches[0]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, config_dict: dict) -> "AutoTPConfig":
|
||||
"""Create config from dictionary (JSON config)."""
|
||||
layer_specs = []
|
||||
for spec_dict in config_dict.get("layer_specs", []):
|
||||
# Convert partition_type string to enum
|
||||
partition_type_str = spec_dict.get("partition_type", "column")
|
||||
if isinstance(partition_type_str, str):
|
||||
partition_type = PartitionType(partition_type_str.lower())
|
||||
else:
|
||||
partition_type = partition_type_str
|
||||
|
||||
# Convert shape from list to tuple if necessary
|
||||
shape = spec_dict.get("shape")
|
||||
if shape is not None:
|
||||
shape = cls._convert_shape(shape)
|
||||
|
||||
layer_specs.append(
|
||||
TPLayerSpec(
|
||||
patterns=spec_dict.get("patterns", []),
|
||||
partition_type=partition_type,
|
||||
shape=shape,
|
||||
partition_dim=spec_dict.get("partition_dim"),
|
||||
model_types=spec_dict.get("model_types"),
|
||||
))
|
||||
|
||||
return cls(
|
||||
tp_size=config_dict.get("tp_size", 1),
|
||||
layer_specs=layer_specs,
|
||||
embedding_partition_dim=config_dict.get("embedding_partition_dim", 1),
|
||||
lm_head_patterns=config_dict.get("lm_head_patterns", ["lm_head", "embed_out"]),
|
||||
use_default_specs=config_dict.get("use_default_specs", True),
|
||||
strict_mode=config_dict.get("strict_mode", False),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _convert_shape(shape):
|
||||
"""Convert shape from list to tuple, handling nested structures."""
|
||||
if isinstance(shape, list):
|
||||
return tuple(AutoTPConfig._convert_shape(item) if isinstance(item, list) else item for item in shape)
|
||||
return shape
|
||||
|
||||
|
||||
class AutoTPPresets:
|
||||
"""Built-in presets for common model architectures."""
|
||||
|
||||
@staticmethod
|
||||
def llama() -> AutoTPConfig:
|
||||
"""LLaMA-style models (separate Q, K, V projections)."""
|
||||
return AutoTPConfig(layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.[qkv]_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.(up|gate)_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def llama_gqa(num_heads: int, num_kv_heads: int, head_dim: int) -> AutoTPConfig:
|
||||
"""LLaMA with Grouped Query Attention (fused QKV variant)."""
|
||||
q_size = num_heads * head_dim
|
||||
kv_size = num_kv_heads * head_dim
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
# Fused QKV with unequal sizes (GQA)
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.qkv_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=((q_size, kv_size, kv_size), -1), # 1-level nesting
|
||||
partition_dim=0,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.(up|gate)_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def bloom() -> AutoTPConfig:
|
||||
"""BLOOM-style models (fused QKV with interleaved heads)."""
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attention\.dense\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attention\.query_key_value\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
# Bloom style: [q1,k1,v1,q2,k2,v2,...] - no reshape needed
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.dense_4h_to_h\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.dense_h_to_4h\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def chatglm() -> AutoTPConfig:
|
||||
"""ChatGLM-style models (GLM-style fused QKV)."""
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attention\.dense\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attention\.query_key_value\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(3, -1), # [Q, K, V] concatenated
|
||||
partition_dim=0,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.dense_4h_to_h\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.dense_h_to_4h\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(2, -1), # [gate, up] packed
|
||||
partition_dim=0,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def mixtral() -> AutoTPConfig:
|
||||
"""Mixtral MoE model."""
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.[qkv]_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
# MoE experts
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.block_sparse_moe\.experts\.\d+\.w2\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.block_sparse_moe\.experts\.\d+\.w[13]\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
# Skip MoE gate
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.block_sparse_moe\.gate\.weight$"],
|
||||
partition_type=PartitionType.SKIP,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def deepseek_v2() -> AutoTPConfig:
|
||||
"""DeepSeek-V2 with MLA (Multi-head Latent Attention)."""
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
# Standard attention output
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
# MLA uses compressed KV, skip low-rank projections
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.(q_a_proj|kv_a_proj_with_mqa)\.weight$"],
|
||||
partition_type=PartitionType.SKIP,
|
||||
),
|
||||
# Q/K/V projections from latent
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.(q_b_proj|kv_b_proj)\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
# MoE experts
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.experts\.\d+\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.experts\.\d+\.(up|gate)_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
# Skip MoE gate
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.gate\.weight$"],
|
||||
partition_type=PartitionType.SKIP,
|
||||
),
|
||||
# Shared expert
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.shared_experts\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.shared_experts\.(up|gate)_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def qwen2() -> AutoTPConfig:
|
||||
"""Qwen2 model."""
|
||||
return AutoTPConfig(layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.[qkv]_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.(up|gate)_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def phi3() -> AutoTPConfig:
|
||||
"""Phi3 model with fused QKV and chunked MLP."""
|
||||
return AutoTPConfig(
|
||||
layer_specs=[
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.o_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
# Phi3 has fused qkv_proj
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.self_attn\.qkv_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(3, -1), # [Q, K, V] concatenated
|
||||
partition_dim=0,
|
||||
),
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.down_proj\.weight$"],
|
||||
partition_type=PartitionType.ROW,
|
||||
),
|
||||
# Phi3 has gate_up_proj fused
|
||||
TPLayerSpec(
|
||||
patterns=[r".*\.mlp\.gate_up_proj\.weight$"],
|
||||
partition_type=PartitionType.COLUMN,
|
||||
shape=(2, -1), # [gate, up] packed
|
||||
partition_dim=0,
|
||||
),
|
||||
], )
|
||||
|
||||
@staticmethod
|
||||
def get_preset(model_type: str) -> Optional[AutoTPConfig]:
|
||||
"""Get a preset configuration by model type name."""
|
||||
presets = {
|
||||
"llama": AutoTPPresets.llama,
|
||||
"bloom": AutoTPPresets.bloom,
|
||||
"chatglm": AutoTPPresets.chatglm,
|
||||
"mixtral": AutoTPPresets.mixtral,
|
||||
"deepseek_v2": AutoTPPresets.deepseek_v2,
|
||||
"qwen2": AutoTPPresets.qwen2,
|
||||
"phi3": AutoTPPresets.phi3,
|
||||
}
|
||||
preset_fn = presets.get(model_type.lower())
|
||||
if preset_fn:
|
||||
return preset_fn()
|
||||
return None
|
||||
|
||||
|
||||
def merge_autotp_configs(base: AutoTPConfig, override: AutoTPConfig) -> AutoTPConfig:
|
||||
"""Merge two AutoTP configs, with override taking precedence."""
|
||||
# Combine layer specs - override specs come first (higher priority)
|
||||
merged_specs = list(override.layer_specs) + list(base.layer_specs)
|
||||
|
||||
return AutoTPConfig(
|
||||
tp_size=override.tp_size if override.tp_size > 1 else base.tp_size,
|
||||
layer_specs=merged_specs,
|
||||
embedding_partition_dim=override.embedding_partition_dim,
|
||||
lm_head_patterns=override.lm_head_patterns or base.lm_head_patterns,
|
||||
use_default_specs=override.use_default_specs,
|
||||
strict_mode=override.strict_mode,
|
||||
)
|
||||
@@ -0,0 +1,21 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .bert import DS_BERTContainer, HFBertLayerPolicy
|
||||
from .bloom import DS_BloomContainer, BLOOMLayerPolicy, supported_models
|
||||
from .distil_bert import DS_DistilBERTContainer, HFDistilBertLayerPolicy
|
||||
from .gpt2 import DS_GPT2Container, HFGPT2LayerPolicy
|
||||
from .gptj import DS_GPTJContainer, HFGPTJLayerPolicy
|
||||
from .gptneo import DS_GPTNEOContainer, HFGPTNEOLayerPolicy
|
||||
from .gptneox import DS_GPTNEOXContainer, GPTNEOXLayerPolicy
|
||||
from .llama import DS_LLAMAContainer, LLAMALayerPolicy
|
||||
from .llama2 import LLAMA2LayerPolicy, DS_LLAMA2Container
|
||||
from .internlm import DS_InternLMContainer, InternLMLayerPolicy
|
||||
from .megatron_gpt import DS_MegatronGPTContainer, MegatronLayerPolicy
|
||||
from .megatron_gpt_moe import DS_MegatronGPTMoEContainer, MegatronMoELayerPolicy
|
||||
from .opt import DS_OPTContainer, HFOPTLayerPolicy
|
||||
from .clip import DS_CLIPContainer, HFCLIPLayerPolicy
|
||||
from .unet import UNetPolicy
|
||||
from .vae import VAEPolicy
|
||||
@@ -0,0 +1,322 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
from abc import ABC
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed
|
||||
from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
|
||||
# If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE
|
||||
# it is assumed the intermediate size is 4x the embedding dimension
|
||||
DEFAULT_INTERMEDIATE_SIZE = -1
|
||||
|
||||
|
||||
class BaseConvolutionContainer(ABC):
|
||||
# not implemented
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
class BaseTransformerContainer(ABC):
|
||||
|
||||
def __init__(self, policy, config, model_config, layer_id, child):
|
||||
self.policy = policy
|
||||
self.config = config
|
||||
self.model_config = model_config
|
||||
self.layer_id = layer_id
|
||||
self.child = child
|
||||
|
||||
self.megatron_v2 = self.policy.is_megatron_v2
|
||||
self.scale_attention = self.policy.scale_attention
|
||||
self.ckpt_load_enabled = False
|
||||
|
||||
# configuration for models. todo: can this be moved to a pydantic model config?
|
||||
self.hidden_size = None
|
||||
self.intermediate_size = None
|
||||
self.num_attention_heads = None
|
||||
self.mp_size = self.config.tensor_parallel.tp_size
|
||||
self.pre_layer_norm = self.model_config.do_layer_norm_before if \
|
||||
hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm
|
||||
self.dtype = self.config.dtype
|
||||
self.attn_linear_layer = self.policy.linear_layer
|
||||
self.mlp_linear_layer = self.policy.linear_layer
|
||||
self.return_tuple = self.config.return_tuple
|
||||
self.triangular_masking = True
|
||||
self.local_attention = ((self.model_config.attention_layers[self.layer_id] == "local") if hasattr(
|
||||
self.model_config, 'attention_layers') else False)
|
||||
self.window_size = getattr(self.model_config, "window_size", 1)
|
||||
self.mlp_act_func_type = self.policy.mlp_act_func_type
|
||||
self.norm_type = self.policy.norm_type
|
||||
self.training_mp_size = self.config.training_mp_size
|
||||
self.bigscience_bloom = False
|
||||
self.max_out_tokens = self.config.max_out_tokens
|
||||
self.min_out_tokens = self.config.min_out_tokens
|
||||
self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False)
|
||||
self.use_mup = self.policy.use_mup
|
||||
self.return_single_tuple = False
|
||||
self.rotary_dim = self.get_rotary_dim()
|
||||
self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0)
|
||||
|
||||
# Attention tensors
|
||||
self.qkvw = None
|
||||
self.qkvb = None
|
||||
self.dense_w = None
|
||||
self.dense_b = None
|
||||
# MLP tensors
|
||||
self._h4h_w = None
|
||||
self._h4h_b = None
|
||||
self._4hh_w = None
|
||||
self._4hh_b = None
|
||||
# LayerNorm tensors
|
||||
self.attn_nw = None
|
||||
self.attn_nb = None
|
||||
self.input_nw = None
|
||||
self.input_nb = None
|
||||
|
||||
self.mp_group = None
|
||||
self.use_triton = False
|
||||
|
||||
# Triton
|
||||
self.use_triton = config.use_triton and deepspeed.HAS_TRITON
|
||||
|
||||
def create_ds_model_config(self):
|
||||
self.set_hidden_heads(*self.policy.get_hidden_heads())
|
||||
assert self.num_attention_heads % self.mp_size == 0,\
|
||||
"To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\
|
||||
"This is because the attention computation is partitioned evenly among the parallel GPUs."
|
||||
|
||||
self.ds_model_config = DeepSpeedInferenceConfig(
|
||||
hidden_size=self.hidden_size,
|
||||
intermediate_size=self.intermediate_size,
|
||||
heads=self.num_attention_heads,
|
||||
layer_norm_eps=self.layernorm_epsilon,
|
||||
dtype=self.dtype,
|
||||
pre_layer_norm=self.pre_layer_norm,
|
||||
norm_type=self.norm_type,
|
||||
mp_size=self.mp_size,
|
||||
return_tuple=self.return_tuple,
|
||||
triangular_masking=self.triangular_masking,
|
||||
local_attention=self.local_attention,
|
||||
window_size=self.window_size,
|
||||
rotary_dim=self.rotary_dim,
|
||||
mlp_after_attn=self.mlp_after_attn,
|
||||
mlp_act_func_type=self.mlp_act_func_type,
|
||||
training_mp_size=self.training_mp_size,
|
||||
bigscience_bloom=self.bigscience_bloom,
|
||||
max_out_tokens=self.max_out_tokens,
|
||||
min_out_tokens=self.min_out_tokens,
|
||||
scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx,
|
||||
use_mup=self.use_mup,
|
||||
return_single_tuple=self.return_single_tuple,
|
||||
set_empty_params=self.config.set_empty_params,
|
||||
transposed_mode=self.config.transposed_mode,
|
||||
use_triton=self.use_triton,
|
||||
triton_autotune=self.config.triton_autotune)
|
||||
|
||||
if self.use_triton and deepspeed.HAS_TRITON:
|
||||
from .bert import DS_BERTContainer
|
||||
if not isinstance(self, DS_BERTContainer):
|
||||
raise NotImplementedError("Triton kernels are only for BERT-like models yet")
|
||||
|
||||
if not self.config.triton_autotune:
|
||||
from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul
|
||||
fp16_matmul.skip_autotune()
|
||||
|
||||
return self.ds_model_config
|
||||
|
||||
def check_meta_tensor_support(self):
|
||||
if hasattr(self.qkvw, 'is_meta'):
|
||||
if self.qkvw.is_meta:
|
||||
assert self.ckpt_load_enabled, "Meta tensors are not supported for this model currently."
|
||||
else:
|
||||
raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+")
|
||||
|
||||
def initialize_tensors(self, enable_training=False):
|
||||
# Set the tensors from policy (user module) to container (DS module)
|
||||
self.set_attention(*self.policy.attention(enable_training=enable_training))
|
||||
self.set_mlp(*self.policy.mlp(enable_training=enable_training))
|
||||
self.set_layernorm(*self.policy.layernorm())
|
||||
#self.check_meta_tensor_support()
|
||||
|
||||
def convert_to_required_dtype(self):
|
||||
# Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy
|
||||
if self.dtype in [torch.half, torch.bfloat16]:
|
||||
for k, v in self.__dict__.items():
|
||||
# The list comprehension is used for MoE tensor lists
|
||||
if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \
|
||||
or isinstance(tensor, torch.nn.Parameter)) for tensor in v):
|
||||
self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v]
|
||||
|
||||
if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter):
|
||||
self.__dict__[k] = v.to(self.dtype)
|
||||
|
||||
def get_rotary_dim(self):
|
||||
if hasattr(self.model_config, 'rotary_dim'):
|
||||
return self.model_config.rotary_dim
|
||||
if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'):
|
||||
return self.child.attention.rotary_ndims
|
||||
return -1
|
||||
|
||||
def set_moe(self, moe=False):
|
||||
self.moe = moe
|
||||
|
||||
def set_tensor_parallel_config(self, mp_size, mp_group):
|
||||
self.mp_size = mp_size
|
||||
self.mp_group = mp_group
|
||||
|
||||
def set_quantization_config(self, quantizer):
|
||||
self.quantizer = quantizer
|
||||
|
||||
def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size):
|
||||
"""
|
||||
Args:
|
||||
hidden_size: embedding dimension of the model
|
||||
num_attention_heads: number of attention heads in the model
|
||||
epsilon: epsilon value for layer norm (same value used for all norms)
|
||||
intermediate_size: Size of MLP projection. If `DEFAULT_INTERMEDIATE_SIZE` is passed
|
||||
it is assumed to be `4 * hidden_size`
|
||||
"""
|
||||
self.hidden_size = hidden_size
|
||||
if intermediate_size == DEFAULT_INTERMEDIATE_SIZE:
|
||||
self.intermediate_size = 4 * hidden_size
|
||||
else:
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.layernorm_epsilon = epsilon
|
||||
|
||||
def set_attention(self, qkvw, qkvb, dense_w, dense_b):
|
||||
self.qkvw = qkvw
|
||||
self.qkvb = qkvb
|
||||
self.dense_w = dense_w
|
||||
self.dense_b = dense_b
|
||||
|
||||
def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b):
|
||||
self._h4h_w = _h4h_w
|
||||
self._h4h_b = _h4h_b
|
||||
self._4hh_w = _4hh_w
|
||||
self._4hh_b = _4hh_b
|
||||
|
||||
def set_layernorm(self, attn_nw, attn_nb, input_nw, input_nb):
|
||||
self.attn_nw = attn_nw
|
||||
self.attn_nb = attn_nb
|
||||
self.input_nw = input_nw
|
||||
self.input_nb = input_nb
|
||||
|
||||
def apply_weight_quantization(self):
|
||||
# quantize attention weights
|
||||
self.attention_quantization()
|
||||
|
||||
# quantize mlp weights
|
||||
self.mlp_quantization()
|
||||
|
||||
def attention_quantization(self):
|
||||
self.module.attention.attn_qkvw = self.quantizer.quantize(self.module.attention.attn_qkvw)
|
||||
self.module.attention.attn_ow = self.quantizer.quantize(self.module.attention.attn_ow)
|
||||
|
||||
def mlp_quantization(self):
|
||||
self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w)
|
||||
self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w)
|
||||
|
||||
def apply_tensor_parallelism(self, mp_replace):
|
||||
# setup the new Attention module
|
||||
self.attention_qkv_mp(mp_replace)
|
||||
self.attention_o_mp(mp_replace)
|
||||
|
||||
# setup the new MLP module
|
||||
self.mlp_inter_mp(mp_replace)
|
||||
self.mlp_output_mp(mp_replace)
|
||||
|
||||
# Apply weight quantization
|
||||
# TODO(cmikeh2): Re-enable this once verified
|
||||
#self.apply_weight_quantization()
|
||||
|
||||
def attention_qkv_mp(self, mp_replace, reversed_dim=False):
|
||||
self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw,
|
||||
self.qkvw,
|
||||
num_splits=3,
|
||||
int8=reversed_dim)
|
||||
self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb,
|
||||
self.qkvb,
|
||||
num_splits=3,
|
||||
int8=reversed_dim)
|
||||
|
||||
def attention_o_mp(self, mp_replace, reversed_dim=False):
|
||||
self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim)
|
||||
self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob,
|
||||
self.dense_b,
|
||||
int8=reversed_dim,
|
||||
allocate_tensor=reversed_dim)
|
||||
|
||||
def mlp_inter_mp(self, mp_replace, reversed_dim=False):
|
||||
self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim)
|
||||
self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim)
|
||||
|
||||
def mlp_output_mp(self, mp_replace, reversed_dim=False):
|
||||
self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim)
|
||||
self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b,
|
||||
self._4hh_b,
|
||||
int8=reversed_dim,
|
||||
allocate_tensor=reversed_dim)
|
||||
|
||||
def copy_data_to_new_module(self):
|
||||
params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb}
|
||||
for key in params:
|
||||
if params[key] is None:
|
||||
setattr(self.module.mlp, key, None)
|
||||
else:
|
||||
setattr(self.module.mlp, key,
|
||||
torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name())))
|
||||
|
||||
params = {'norm_w': self.input_nw, 'norm_b': self.input_nb}
|
||||
for key in params:
|
||||
if params[key] is None:
|
||||
setattr(self.module, key, None)
|
||||
else:
|
||||
setattr(self.module, key,
|
||||
torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name())))
|
||||
|
||||
def transpose(self):
|
||||
self.transpose_attention()
|
||||
self.transpose_mlp()
|
||||
|
||||
def transpose_attention(self):
|
||||
if self.attn_linear_layer:
|
||||
self.qkvw = self.transpose_impl(self.qkvw.data)
|
||||
self.dense_w = self.transpose_impl(self.dense_w.data)
|
||||
|
||||
def transpose_mlp(self):
|
||||
if self.mlp_linear_layer:
|
||||
self._h4h_w = self.transpose_impl(self._h4h_w.data)
|
||||
self._4hh_w = self.transpose_impl(self._4hh_w.data)
|
||||
|
||||
def transpose_impl(self, data):
|
||||
data = data.contiguous()
|
||||
data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1))
|
||||
data = data.reshape(data.shape[-1], data.shape[-2])
|
||||
data.to(get_accelerator().current_device_name())
|
||||
return data
|
||||
|
||||
def get_all_params(self):
|
||||
params = [
|
||||
self.attn_nw,
|
||||
self.attn_nb,
|
||||
self.input_nw,
|
||||
self.input_nb,
|
||||
]
|
||||
|
||||
params.extend(self.get_attn_params())
|
||||
params.extend(self.get_mlp_params())
|
||||
|
||||
return params
|
||||
|
||||
def get_attn_params(self):
|
||||
return [self.qkvw, self.qkvb, self.dense_w, self.dense_b]
|
||||
|
||||
def get_mlp_params(self):
|
||||
return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b]
|
||||
@@ -0,0 +1,130 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Create a container object to save model-specific tensors using the policy file above.
|
||||
from .base import *
|
||||
from deepspeed import comm as dist
|
||||
import deepspeed.ops.transformer as transformer_inference
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
|
||||
|
||||
class BaseTransformerMoEContainer(BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
# Call the init function of the parent class to initialize the tensors and configs from parent class
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.num_experts = self.policy.get_num_experts()
|
||||
self.ep_world_size = dist.get_world_size()
|
||||
self.local_ep_size = 1 if self.num_experts < self.ep_world_size else self.num_experts // self.ep_world_size
|
||||
|
||||
self.layer_norm_eps = self.config.layer_norm_eps if hasattr(self.config, 'layer_norm_eps') else 1e-12,
|
||||
|
||||
# MoE models will have a list of mlp related tensors
|
||||
self._h4h_w = []
|
||||
self._h4h_b = []
|
||||
self._4hh_w = []
|
||||
self._4hh_b = []
|
||||
|
||||
# Residual MoE needs extra parameters
|
||||
self._res_h4h_w = None
|
||||
self._res_h4h_b = None
|
||||
self._res_4hh_w = None
|
||||
self._res_4hh_b = None
|
||||
self._res_coef = None
|
||||
|
||||
def create_ds_model_config(self):
|
||||
self.set_hidden_heads(*self.policy.get_hidden_heads())
|
||||
assert self.num_attention_heads % self.mp_size == 0,\
|
||||
"To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\
|
||||
"This is because the attention computation is partitioned evenly among the parallel GPUs."
|
||||
|
||||
self.ds_model_config = transformer_inference.DeepSpeedMoEInferenceConfig(
|
||||
hidden_size=self.hidden_size,
|
||||
heads=self.num_attention_heads,
|
||||
layer_norm_eps=self.layer_norm_eps,
|
||||
fp16=self.fp16,
|
||||
pre_layer_norm=self.pre_layer_norm,
|
||||
mp_size=self.mp_size,
|
||||
q_int8=self.quantize,
|
||||
moe_experts=self.local_ep_size,
|
||||
global_experts=self.num_experts,
|
||||
mlp_type=self.config.moe.type,
|
||||
scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx,
|
||||
)
|
||||
|
||||
return self.ds_model_config
|
||||
|
||||
def initialize_tensors(self):
|
||||
# Set the tensors from policy (user module) to container (DS module)
|
||||
self.set_attention(*self.policy.attention())
|
||||
self.set_mlp(self.config.moe.type)
|
||||
self.set_layernorm(*self.policy.layernorm())
|
||||
|
||||
def set_mlp(self, config_moe_type):
|
||||
if config_moe_type == 'standard':
|
||||
self._h4h_w, self._h4h_b, \
|
||||
self._4hh_w, self._4hh_b = self.policy.mlp()
|
||||
else:
|
||||
self._h4h_w, self._h4h_b, self._4hh_w, \
|
||||
self._4hh_b, self._res_h4h_w, self._res_h4h_b, \
|
||||
self._res_4hh_w, self._res_4hh_b, \
|
||||
self._res_coef = self.policy.mlp(config_moe_type)
|
||||
|
||||
def transpose(self):
|
||||
self.transpose_attention()
|
||||
self.transpose_mlp()
|
||||
|
||||
if self.config.moe.type == 'residual':
|
||||
self.transpose_residual()
|
||||
|
||||
def transpose_mlp(self):
|
||||
self._h4h_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._h4h_w]
|
||||
self._4hh_w = [self.transpose_impl(moe_w1.data) for moe_w1 in self._4hh_w]
|
||||
|
||||
def transpose_residual(self):
|
||||
self._res_h4h_w.data = self.transpose_impl(self._res_h4h_w.data)
|
||||
self._res_4hh_w.data = self.transpose_impl(self._res_4hh_w.data)
|
||||
self._res_coef.data = self.transpose_impl(self._res_coef.data)
|
||||
|
||||
def apply_tensor_parallelism(self, mp_replace):
|
||||
# setup the new Attention module
|
||||
self.attention_qkv_mp(mp_replace)
|
||||
self.attention_o_mp(mp_replace)
|
||||
|
||||
# quantize attention weights
|
||||
self.attention_quantization()
|
||||
|
||||
# setup the new MLP module
|
||||
self.mlp_mp()
|
||||
|
||||
def mlp_mp(self):
|
||||
gpu_index = dist.get_rank()
|
||||
for ep_index in range(self.local_ep_size):
|
||||
# mlp inter
|
||||
self.module.mlp[ep_index].inter_w.data = self._h4h_w[gpu_index * self.local_ep_size + ep_index].to(
|
||||
get_accelerator().current_device_name())
|
||||
self.module.mlp[ep_index].inter_b.data = self._h4h_b[gpu_index * self.local_ep_size + ep_index].to(
|
||||
get_accelerator().current_device_name())
|
||||
|
||||
# mlp output
|
||||
self.module.mlp[ep_index].output_w.data = self._4hh_w[gpu_index * self.local_ep_size + ep_index].to(
|
||||
get_accelerator().current_device_name())
|
||||
self.module.mlp[ep_index].output_b.data = self._4hh_b[gpu_index * self.local_ep_size + ep_index].to(
|
||||
get_accelerator().current_device_name())
|
||||
|
||||
def copy_data_to_new_module(self):
|
||||
self.module.attn_nw.data = self.attn_nw.to(get_accelerator().current_device_name())
|
||||
self.module.attn_nb.data = self.attn_nb.to(get_accelerator().current_device_name())
|
||||
|
||||
self.module.norm_w.data.copy_(self.input_nw.to(get_accelerator().current_device_name()))
|
||||
self.module.norm_b.data.copy_(self.input_nb.to(get_accelerator().current_device_name()))
|
||||
|
||||
if self.config.moe.type == 'residual':
|
||||
self.module.res_mlp.inter_w.data = self._res_h4h_w.to(get_accelerator().current_device_name())
|
||||
self.module.res_mlp.inter_b.data = self._res_h4h_b.to(get_accelerator().current_device_name())
|
||||
self.module.res_mlp.output_w.data = self._res_4hh_w.to(get_accelerator().current_device_name())
|
||||
self.module.res_mlp.output_b.data = self._res_4hh_b.to(get_accelerator().current_device_name())
|
||||
self.module.res_coef.data = self._res_coef.to(get_accelerator().current_device_name())
|
||||
@@ -0,0 +1,93 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
|
||||
|
||||
class DS_BERTContainer(BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
self.return_tuple = True
|
||||
self.triangular_masking = False
|
||||
self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
|
||||
class HFBertLayerPolicy(TransformerPolicy):
|
||||
|
||||
def __init__(self, client_module, inference=False):
|
||||
super().__init__(inference, pre_attn_norm=False)
|
||||
self.client_module = client_module
|
||||
self.cuda_graph_supported = True
|
||||
|
||||
if HFBertLayerPolicy._orig_layer_class is None:
|
||||
try:
|
||||
import transformers
|
||||
HFBertLayerPolicy._orig_layer_class = [
|
||||
transformers.models.bert.modeling_bert.BertLayer,
|
||||
transformers.models.roberta.modeling_roberta.RobertaLayer
|
||||
]
|
||||
except Exception:
|
||||
HFBertLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
if self.pre_attn_norm:
|
||||
attention_layernorm = self.client_module.PostAttentionLayerNorm
|
||||
else:
|
||||
attention_layernorm = self.client_module.attention.output.LayerNorm
|
||||
return self.client_module.attention.self.query.weight.shape[1], \
|
||||
self.client_module.attention.self.num_attention_heads, \
|
||||
attention_layernorm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.attention.self.query.weight
|
||||
qb = self.client_module.attention.self.query.bias
|
||||
kw = self.client_module.attention.self.key.weight
|
||||
kb = self.client_module.attention.self.key.bias
|
||||
vw = self.client_module.attention.self.value.weight
|
||||
vb = self.client_module.attention.self.value.bias
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
qkvb, \
|
||||
self.client_module.attention.output.dense.weight, \
|
||||
self.client_module.attention.output.dense.bias, \
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
if self.pre_attn_norm:
|
||||
intermediate_ff = self.client_module.intermediate.dense_act
|
||||
else:
|
||||
intermediate_ff = self.client_module.intermediate.dense
|
||||
|
||||
return intermediate_ff.weight, intermediate_ff.bias, \
|
||||
self.client_module.output.dense.weight, \
|
||||
self.client_module.output.dense.bias
|
||||
|
||||
def layernorm(self):
|
||||
if self.pre_attn_norm:
|
||||
attention_layernorm = self.client_module.PostAttentionLayerNorm
|
||||
transformer_layernorm = self.client_module.PreAttentionLayerNorm
|
||||
else:
|
||||
attention_layernorm = self.client_module.attention.output.LayerNorm
|
||||
transformer_layernorm = self.client_module.output.LayerNorm
|
||||
return attention_layernorm.weight, \
|
||||
attention_layernorm.bias, \
|
||||
transformer_layernorm.weight, \
|
||||
transformer_layernorm.bias
|
||||
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features.meta_tensor import MetaTensorContainer
|
||||
from .features.hybrid_engine import HybridEngineContainer
|
||||
from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference
|
||||
from ..policy import TransformerPolicy
|
||||
from ..policy import transformer_param_names
|
||||
from ..policy import maybe_copy
|
||||
|
||||
from ..policy import maybe_get_lora
|
||||
|
||||
supported_models = {None}
|
||||
|
||||
|
||||
class DS_BloomContainer(MetaTensorContainer, HybridEngineContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
# Check transformers version, error if > 4.43.4 (breaks at 4.44.0)
|
||||
from importlib.metadata import version
|
||||
v_transformers = version('transformers')
|
||||
vers = v_transformers.split('.')
|
||||
major = int(vers[0])
|
||||
minor = int(vers[1])
|
||||
if major > 4 or (major == 4 and minor > 43):
|
||||
raise RuntimeError(
|
||||
f"Transformers version {v_transformers} exceeds version 4.43.4! After transformers version 4.43.4, BLOOM inference with DeepSpeed is no longer supported."
|
||||
)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
self.bigscience_bloom = True
|
||||
self.triangular_masking = False
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
|
||||
self.module = DeepSpeedBloomInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
self.module.config.invert_mask = False
|
||||
return self.module
|
||||
|
||||
def attention_qkv_mp(self, mp_replace, reversed_dim=False):
|
||||
self.module.attention.attn_qkvw = mp_replace.copy(self.module.attention.attn_qkvw, self.qkvw)
|
||||
self.module.attention.attn_qkvb = mp_replace.copy(self.module.attention.attn_qkvb, self.qkvb)
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params()
|
||||
ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)]
|
||||
return ret
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, self.policy.
|
||||
client_module.self_attention.query_key_value, self.policy.client_module.self_attention.dense
|
||||
]
|
||||
]
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'self_attention.query_key_value.weight', \
|
||||
'self_attention.query_key_value.bias', \
|
||||
'self_attention.dense.weight', \
|
||||
'self_attention.dense.bias', \
|
||||
'mlp.dense_h_to_4h.weight', \
|
||||
'mlp.dense_h_to_4h.bias', \
|
||||
'mlp.dense_4h_to_h.weight', \
|
||||
'mlp.dense_4h_to_h.bias', \
|
||||
'post_attention_layernorm.weight', \
|
||||
'post_attention_layernorm.bias', \
|
||||
'input_layernorm.weight', \
|
||||
'input_layernorm.bias'
|
||||
)
|
||||
for i in range(0, 2):
|
||||
maybe_copy(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
transformer_param_names[i],
|
||||
prefix + param_names[i],
|
||||
qkv=True,
|
||||
megatron_v2=self.policy.is_megatron_v2,
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(2, 4):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i],
|
||||
prefix + param_names[i])
|
||||
for i in range(4, 10):
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i],
|
||||
prefix + param_names[i])
|
||||
for i in range(10, 12):
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i])
|
||||
|
||||
|
||||
class BLOOMLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self, client_module, inference=True, use_load_prefix=True, split_qkv=False):
|
||||
super().__init__(inference, linear_layer=True, use_load_prefix=use_load_prefix, split_qkv=split_qkv)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
BLOOMLayerPolicy._orig_layer_class = transformers.models.bloom.modeling_bloom.BloomBlock
|
||||
global supported_models
|
||||
supported_models.update({transformers.models.bloom.modeling_bloom.BloomModel})
|
||||
except Exception as e:
|
||||
print(f"WARNING! Setting BLOOMLayerPolicy._orig_layer_class to None due to Exception: {e}")
|
||||
BLOOMLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.self_attention.hidden_size, \
|
||||
self.client_module.self_attention.num_heads, \
|
||||
self.client_module.input_layernorm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
return self.client_module.self_attention.query_key_value.weight, \
|
||||
self.client_module.self_attention.query_key_value.bias, \
|
||||
self.client_module.self_attention.dense.weight, \
|
||||
self.client_module.self_attention.dense.bias,
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.dense_h_to_4h.weight, \
|
||||
self.client_module.mlp.dense_h_to_4h.bias, \
|
||||
self.client_module.mlp.dense_4h_to_h.weight, \
|
||||
self.client_module.mlp.dense_4h_to_h.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.post_attention_layernorm.weight, \
|
||||
self.client_module.post_attention_layernorm.bias, \
|
||||
self.client_module.input_layernorm.weight, \
|
||||
self.client_module.input_layernorm.bias
|
||||
@@ -0,0 +1,73 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
|
||||
|
||||
class DS_CLIPContainer(BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
|
||||
class HFCLIPLayerPolicy(TransformerPolicy):
|
||||
|
||||
def __init__(self, client_module, inference=False):
|
||||
super().__init__(inference, pre_attn_norm=True, scale_attention=True)
|
||||
self.client_module = client_module
|
||||
self.cuda_graph_supported = True
|
||||
|
||||
if HFCLIPLayerPolicy._orig_layer_class is None:
|
||||
try:
|
||||
import transformers
|
||||
HFCLIPLayerPolicy._orig_layer_class = transformers.models.clip.modeling_clip.CLIPEncoderLayer
|
||||
except Exception:
|
||||
HFCLIPLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.self_attn.q_proj.weight.shape[1], \
|
||||
self.client_module.self_attn.num_heads, \
|
||||
self.client_module.layer_norm1.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.self_attn.q_proj.weight
|
||||
qb = self.client_module.self_attn.q_proj.bias
|
||||
kw = self.client_module.self_attn.k_proj.weight
|
||||
kb = self.client_module.self_attn.k_proj.bias
|
||||
vw = self.client_module.self_attn.v_proj.weight
|
||||
vb = self.client_module.self_attn.v_proj.bias
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
qkvb, \
|
||||
self.client_module.self_attn.out_proj.weight, \
|
||||
self.client_module.self_attn.out_proj.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.fc1.weight, \
|
||||
self.client_module.mlp.fc1.bias, \
|
||||
self.client_module.mlp.fc2.weight, \
|
||||
self.client_module.mlp.fc2.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.layer_norm2.weight, \
|
||||
self.client_module.layer_norm2.bias, \
|
||||
self.client_module.layer_norm1.weight, \
|
||||
self.client_module.layer_norm1.bias
|
||||
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
|
||||
|
||||
class DS_DistilBERTContainer(BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
self.triangular_masking = False
|
||||
self.return_single_tuple = True
|
||||
self.use_triton = kwargs['config'].use_triton and deepspeed.HAS_TRITON
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedBERTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
|
||||
class HFDistilBertLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self, client_module, inference=False, preln=False):
|
||||
super().__init__(inference)
|
||||
self.client_module = client_module
|
||||
self.preln = preln
|
||||
self.cuda_graph_supported = True
|
||||
if HFDistilBertLayerPolicy._orig_layer_class is None:
|
||||
try:
|
||||
import transformers
|
||||
HFDistilBertLayerPolicy._orig_layer_class = [
|
||||
transformers.models.distilbert.modeling_distilbert.TransformerBlock,
|
||||
]
|
||||
except Exception:
|
||||
HFDistilBertLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.attention.q_lin.weight.shape[1], \
|
||||
self.client_module.attention.n_heads, \
|
||||
self.client_module.sa_layer_norm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.attention.q_lin.weight
|
||||
qb = self.client_module.attention.q_lin.bias
|
||||
kw = self.client_module.attention.k_lin.weight
|
||||
kb = self.client_module.attention.k_lin.bias
|
||||
vw = self.client_module.attention.v_lin.weight
|
||||
vb = self.client_module.attention.v_lin.bias
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
qkvb, \
|
||||
self.client_module.attention.out_lin.weight, \
|
||||
self.client_module.attention.out_lin.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
intermediate_ff = self.client_module.ffn.lin1
|
||||
|
||||
return intermediate_ff.weight, intermediate_ff.bias, \
|
||||
self.client_module.ffn.lin2.weight, \
|
||||
self.client_module.ffn.lin2.bias
|
||||
|
||||
def layernorm(self):
|
||||
attention_layernorm = self.client_module.sa_layer_norm
|
||||
transformer_layernorm = self.client_module.output_layer_norm
|
||||
return attention_layernorm.weight, \
|
||||
attention_layernorm.bias, \
|
||||
transformer_layernorm.weight, \
|
||||
transformer_layernorm.bias
|
||||
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .gated_mlp import HybridGatedMLPContainer
|
||||
from .megatron import MegatronContainer
|
||||
from .meta_tensor import MetaTensorContainer
|
||||
from .split_qkv import HybridSplitQKVContainer
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
from .hybrid_engine import HybridEngineContainer
|
||||
|
||||
|
||||
class HybridGatedMLPContainer(HybridEngineContainer):
|
||||
"""
|
||||
The HybridGatedMLPContainer supports models for which the first MLP layer
|
||||
is represented with two separate weights, one for the activation function
|
||||
and one for the gating function.
|
||||
"""
|
||||
|
||||
def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b):
|
||||
super().set_mlp(_h4h_w, _h4h_b, _4hh_w, _4hh_b)
|
||||
self.set_mlp_gate()
|
||||
|
||||
@abstractmethod
|
||||
def set_mlp_gate(self):
|
||||
"""
|
||||
In `set_mlp_gate`, it is necessary to populate the following variables (where appropriate)
|
||||
for the given model:
|
||||
self.inter_up_w: inter up weight
|
||||
self.inter_up_b: inter up bias
|
||||
self.inter_gate_w: inter gate weight
|
||||
self.inter_gate_b: inter gate bias
|
||||
If the parameter does not exist in the original model, set the attribute to None.
|
||||
"""
|
||||
raise NotImplementedError("A set_mlp_gate() function must be defined in the model container \
|
||||
in order to set the unfused inter up and gate tensors.")
|
||||
|
||||
def mlp_inter_mp(self, mp_replace, reversed_dim=False):
|
||||
# Only need to alter behavior if we can't do the normal destructive copy
|
||||
if self.module.mlp.inter_w is None:
|
||||
params = [
|
||||
(self.module.mlp.inter_up_w, self.inter_up_w),
|
||||
(self.module.mlp.inter_up_b, self.inter_up_b),
|
||||
(self.module.mlp.inter_gate_w, self.inter_gate_w),
|
||||
(self.module.mlp.inter_gate_b, self.inter_gate_b),
|
||||
]
|
||||
for dst, src in params:
|
||||
dst = mp_replace.copy(dst[:self.inter_up_w.shape[0] // mp_replace.mp_size],
|
||||
src,
|
||||
int8=reversed_dim,
|
||||
allocate_tensor=reversed_dim) if src is not None else None
|
||||
else:
|
||||
self.module.mlp.inter_w = mp_replace.strided_copy(self.module.mlp.inter_w,
|
||||
self._h4h_w,
|
||||
num_splits=2,
|
||||
int8=reversed_dim)
|
||||
self.module.mlp.inter_b = mp_replace.strided_copy(self.module.mlp.inter_b,
|
||||
self._h4h_b,
|
||||
num_splits=2,
|
||||
int8=reversed_dim)
|
||||
|
||||
def release_mlp(self):
|
||||
super().release_mlp()
|
||||
gated_mlp_params = [
|
||||
(self.module.mlp.inter_up_w, self.inter_up_w),
|
||||
(self.module.mlp.inter_up_b, self.inter_up_b),
|
||||
(self.module.mlp.inter_gate_w, self.inter_gate_w),
|
||||
(self.module.mlp.inter_gate_b, self.inter_gate_b),
|
||||
]
|
||||
|
||||
self._release_params(gated_mlp_params)
|
||||
|
||||
def reset_mlp(self):
|
||||
self._h4h_w.data[:self.inter_up_w.shape[0]] = self.inter_up_w.data
|
||||
self._h4h_w.data[self.inter_up_w.shape[0]:] = self.inter_gate_w.data
|
||||
|
||||
if self.inter_up_b is not None:
|
||||
self._h4h_b.data[:self.inter_up_b.shape[0]] = self.inter_up_b.data
|
||||
self._h4h_b.data[self.inter_up_b.shape[0]:] = self.inter_gate_b.data
|
||||
|
||||
inter_data = [self.inter_up_w.data, self.inter_gate_w.data]
|
||||
if self.inter_up_b is not None:
|
||||
inter_data.extend([self.inter_up_b.data, self.inter_gate_b.data])
|
||||
|
||||
self.inter_up_w.data = self._h4h_w.data[:self.inter_up_w.shape[0]]
|
||||
self.inter_gate_w.data = self._h4h_w.data[self.inter_up_w.shape[0]:]
|
||||
|
||||
if self.inter_up_b is not None:
|
||||
self.inter_up_b.data = self._h4h_b.data[:self.inter_up_b.shape[0]]
|
||||
self.inter_gate_b.data = self._h4h_b.data[self.inter_up_b.shape[0]:]
|
||||
|
||||
for data in inter_data:
|
||||
del data
|
||||
|
||||
def set_mlp_params_wo_copy(self, Z3_enabled=False):
|
||||
self.module.mlp.output_w = self._4hh_w
|
||||
self.module.mlp.output_b = self._4hh_b
|
||||
|
||||
if not Z3_enabled:
|
||||
# In initialize_tensors, we create a fused inter projection with the appropriate shape
|
||||
# and copy the up projection and gate projection into it
|
||||
self.module.mlp.inter_w = self._h4h_w
|
||||
self.module.mlp.inter_b = self._h4h_b
|
||||
|
||||
self.inter_up_w.data = self._h4h_w[:self.inter_up_w.shape[0], :]
|
||||
self.inter_gate_w.data = self._h4h_w[self.inter_up_w.shape[0]:, :]
|
||||
|
||||
if self.inter_up_b is not None:
|
||||
self.inter_up_b.data = self._h4h_b[:self.inter_up_w.shape[0]] if self._h4h_b is not None else None
|
||||
self.inter_gate_b.data = self._h4h_b[self.inter_up_w.shape[0]:] if self._h4h_b is not None else None
|
||||
else:
|
||||
self.module.mlp.inter_up_w = self.inter_up_w
|
||||
self.module.mlp.inter_up_b = self.inter_up_b
|
||||
self.module.mlp.inter_gate_w = self.inter_gate_w
|
||||
self.module.mlp.inter_gate_b = self.inter_gate_b
|
||||
|
||||
def get_mlp_params(self):
|
||||
params = super().get_mlp_params()
|
||||
params.extend([self.inter_up_w, self.inter_up_b, self.inter_gate_w, self.inter_gate_b])
|
||||
return params
|
||||
@@ -0,0 +1,212 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class HybridEngineContainer(ABC):
|
||||
"""
|
||||
This container identifies which methods need to be overridden in addition to
|
||||
the base container to enable use in the RLHF pipeline. These methods are not
|
||||
necessary for inference alone.
|
||||
|
||||
NOTE: If you are using this feature with a container that
|
||||
also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer`
|
||||
is inherited before `HybridEngineContainer` in the class definition.
|
||||
"""
|
||||
|
||||
def initialize_tensors(self, enable_training=False):
|
||||
"""
|
||||
Same purposes as the base container, but also grabs the hooks for any LoRA
|
||||
parameters. If it's necessary to override specific sub-components of the model,
|
||||
it's best to augment the specific `set_[component]` itself rather than modifying
|
||||
the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example.
|
||||
"""
|
||||
super().initialize_tensors(enable_training=enable_training)
|
||||
self.set_lora_params()
|
||||
|
||||
def transform_for_training(self):
|
||||
"""
|
||||
If the views on certain parameters are largely incompatible, it may be necessary to do
|
||||
more substantial transformations to the parameters. This method should be overridden to
|
||||
transform the inference format to what is necessary for training.
|
||||
"""
|
||||
pass
|
||||
|
||||
def transform_for_inference(self):
|
||||
"""
|
||||
If the views on certain parameters are largely incompatible, it may be necessary to do
|
||||
more substantial transformations to the parameters. This method should be overridden to
|
||||
transform the training format to what is necessary for inference.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
If available, set the LoRA parameters for the module. An implementation
|
||||
for this would iterate over all parameters of the model and use the `maybe_get_lora` helper
|
||||
method to check if the parameter does in fact have any LoRA params.
|
||||
"""
|
||||
raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.")
|
||||
|
||||
@abstractmethod
|
||||
def get_lora_matched_pair(self):
|
||||
"""Get the pair of lora params and its matched model parameters."""
|
||||
raise NotImplementedError("get_lora_matched_pair() must be defined for the relevant parameters.")
|
||||
|
||||
def fuse_lora(self):
|
||||
"""Fuse the LoRA parameters for the inference mode."""
|
||||
for maybe_lora_param, param in self.get_lora_matched_pair():
|
||||
if len(maybe_lora_param) == 3:
|
||||
lora_right_weight, \
|
||||
lora_left_weight, \
|
||||
lora_scaling = maybe_lora_param
|
||||
param.data += lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t())
|
||||
|
||||
def unfuse_lora(self):
|
||||
"""Unfuse the LoRA parameters for the training mode."""
|
||||
for maybe_lora_param, param in self.get_lora_matched_pair():
|
||||
if len(maybe_lora_param) == 3:
|
||||
lora_right_weight, \
|
||||
lora_left_weight, \
|
||||
lora_scaling = maybe_lora_param
|
||||
param.data -= lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t())
|
||||
|
||||
def apply_tensor_parallelism(self, mp_replace, reversed_dim=False):
|
||||
"""
|
||||
Add support for reversed dim in tensor parallelism. If necessary, override
|
||||
the called methods to handle partitioned weights (i.e. if qkv is split, override
|
||||
the `attention_qkv_mp` method). If the model component is not split, it should
|
||||
be safe to use the default implementation.
|
||||
"""
|
||||
# Setup the new Attention module
|
||||
self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim)
|
||||
self.attention_o_mp(mp_replace, reversed_dim=reversed_dim)
|
||||
|
||||
# Setup the new MLP module
|
||||
self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim)
|
||||
self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim)
|
||||
|
||||
# Apply weight quantization
|
||||
# TODO(cmikeh2): Re-enable this once verified
|
||||
#self.apply_weight_quantization()
|
||||
|
||||
def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]):
|
||||
"""
|
||||
Helper for `release_[component]` methods. Accepts a list of tuples where the first
|
||||
element is the module param that needs to be deleted, and the second is the reassignment
|
||||
from the container.
|
||||
"""
|
||||
for module_param, container_param in param_pairs:
|
||||
if module_param is not None:
|
||||
del module_param
|
||||
module_param = container_param
|
||||
|
||||
def release_memory(self):
|
||||
"""
|
||||
Delete module parameters if they exist and point them back to the container. The primary
|
||||
purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the
|
||||
parameters we've created for inference to free their memory.
|
||||
"""
|
||||
general_params = [
|
||||
(self.module.attention.attn_ow, self.dense_w),
|
||||
(self.module.attention.attn_ob, self.dense_b),
|
||||
(self.module.mlp.attn_nw, self.attn_nw),
|
||||
(self.module.mlp.attn_nb, self.attn_nb),
|
||||
(self.module.norm_w, self.input_nw),
|
||||
(self.module.norm_b, self.input_nb),
|
||||
]
|
||||
|
||||
self._release_params(general_params)
|
||||
|
||||
self.release_qkv()
|
||||
self.release_mlp()
|
||||
|
||||
def release_qkv(self):
|
||||
"""
|
||||
Release for QKV parameters (as well as any aliases).
|
||||
"""
|
||||
qkv_params = [
|
||||
(self.module.attention.attn_qkvw, self.qkvw),
|
||||
(self.module.attention.attn_qkvb, self.qkvb),
|
||||
]
|
||||
|
||||
self._release_params(qkv_params)
|
||||
|
||||
def release_mlp(self):
|
||||
"""
|
||||
Release for MLP parameters (as well as any aliases).
|
||||
"""
|
||||
mlp_params = [
|
||||
(self.module.mlp.inter_w, self._h4h_w),
|
||||
(self.module.mlp.inter_b, self._h4h_b),
|
||||
(self.module.mlp.output_w, self._4hh_w),
|
||||
(self.module.mlp.output_b, self._4hh_b),
|
||||
]
|
||||
|
||||
self._release_params(mlp_params)
|
||||
|
||||
def reset_params(self):
|
||||
"""
|
||||
The purpose of reset params is to get the weights from the FP16 training
|
||||
copy of the model and copy to them to contiguous inference view. This only needs
|
||||
to be performed when the container parameters cannot be used directly for inference.
|
||||
"""
|
||||
self.reset_qkv()
|
||||
self.reset_mlp()
|
||||
|
||||
def reset_qkv(self):
|
||||
"""
|
||||
Perform any necessary resets of the model parameters for the QKV components.
|
||||
"""
|
||||
pass
|
||||
|
||||
def reset_mlp(self):
|
||||
"""
|
||||
Perform any necessary resets of the model parameters for the MLP components.
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_lora_params(self):
|
||||
"""
|
||||
Return a list of all parameters that would have LoRA for the module.
|
||||
"""
|
||||
if not hasattr(self, "lora_params"):
|
||||
self.set_lora_params()
|
||||
return self.lora_params
|
||||
|
||||
def set_params_wo_copy(self, Z3_enabled=False):
|
||||
"""
|
||||
Rather than copying into, set the parameters directly. This is necessary to provide
|
||||
an inexpensive (low-memory-overhead) view onto the FP16 forward weights.
|
||||
"""
|
||||
self.module.mlp.attn_nw = self.attn_nw
|
||||
self.module.mlp.attn_nb = self.attn_nb
|
||||
self.module.norm_w = self.input_nw
|
||||
self.module.norm_b = self.input_nb
|
||||
self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled)
|
||||
self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled)
|
||||
|
||||
def set_attn_params_wo_copy(self, **kwargs):
|
||||
"""
|
||||
Narrower sub-method for finer grained overriding.
|
||||
"""
|
||||
self.module.attention.attn_ow = self.dense_w
|
||||
self.module.attention.attn_ob = self.dense_b
|
||||
self.module.attention.attn_qkvw = self.qkvw
|
||||
self.module.attention.attn_qkvb = self.qkvb
|
||||
|
||||
def set_mlp_params_wo_copy(self, **kwargs):
|
||||
"""
|
||||
Narrower sub-method for finer grained overriding.
|
||||
"""
|
||||
self.module.mlp.inter_w = self._h4h_w
|
||||
self.module.mlp.inter_b = self._h4h_b
|
||||
self.module.mlp.output_w = self._4hh_w
|
||||
self.module.mlp.output_b = self._4hh_b
|
||||
@@ -0,0 +1,87 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
from .hybrid_engine import HybridEngineContainer
|
||||
from .megatron import MegatronContainer
|
||||
|
||||
|
||||
class HybridMegatronContainer(MegatronContainer, HybridEngineContainer):
|
||||
|
||||
def _align_qkv(self, x: torch.Tensor):
|
||||
"""
|
||||
Internal helper for accepting the head-contiguous weight matrix and chunking
|
||||
the query, key, and value components.
|
||||
"""
|
||||
attention_head_size = x.shape[0] // self.num_attention_heads
|
||||
new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:]
|
||||
x_1 = x.view(*new_x_shape)
|
||||
div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1
|
||||
(q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim)
|
||||
if len(q.shape) > 2:
|
||||
x.data.copy_(
|
||||
torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])),
|
||||
dim=0).reshape(x.shape))
|
||||
else:
|
||||
x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape))
|
||||
|
||||
def transform_for_inference(self) -> None:
|
||||
"""
|
||||
Overrides the HybridEngineContainer implementation.
|
||||
|
||||
The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V
|
||||
are sequential in memory. This is different from the default layout in which all of the Qs
|
||||
are sequential, followed by all of the Ks, and then all of the Vs. Here, we take the default
|
||||
layout and transform it to the inference layout.
|
||||
"""
|
||||
if hasattr(self.qkvw, 'ds_id'):
|
||||
from deepspeed.runtime.zero import GatheredParameters
|
||||
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
|
||||
param_list = [self.qkvw, self.qkvb]
|
||||
non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \
|
||||
param.ds_status == ZeroParamStatus.NOT_AVAILABLE)]
|
||||
with GatheredParameters(non_active_params):
|
||||
self._align_qkv(self.qkvw)
|
||||
self._align_qkv(self.qkvb)
|
||||
else:
|
||||
self._align_qkv(self.qkvw)
|
||||
self._align_qkv(self.qkvb)
|
||||
|
||||
def _partition_qkv(self, x: torch.Tensor):
|
||||
"""
|
||||
Internal helper for taking contiguous QKV and partitioning it for contiguous
|
||||
heads.
|
||||
"""
|
||||
q_k_v = torch.split(x, (x.shape[0] // 3), dim=0)
|
||||
attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads
|
||||
new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:]
|
||||
q, k, v = [data.view(*new_x_shape) for data in q_k_v]
|
||||
if len(q.shape) > 2:
|
||||
x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1]))
|
||||
else:
|
||||
x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1))
|
||||
|
||||
def transform_for_training(self):
|
||||
"""
|
||||
Overrides the HybridEngineContainer implementation.
|
||||
|
||||
The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V
|
||||
are sequential in memory. This is different from the default layout in which all of the Qs
|
||||
are sequential, followed by all of the Ks, and then all of the Vs. This function takes the inference format and reverts it back to the default format.
|
||||
"""
|
||||
# If parameter is distributed, handle gathering it
|
||||
if hasattr(self.qkvw, 'ds_id'):
|
||||
from deepspeed.runtime.zero import GatheredParameters
|
||||
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
|
||||
param_list = [self.qkvw, self.qkvb]
|
||||
non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \
|
||||
param.ds_status == ZeroParamStatus.NOT_AVAILABLE)]
|
||||
with GatheredParameters(non_active_params):
|
||||
self._partition_qkv(self.qkvw)
|
||||
self._partition_qkv(self.qkvb)
|
||||
else:
|
||||
self._partition_qkv(self.qkvw)
|
||||
self._partition_qkv(self.qkvb)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from abc import ABC
|
||||
|
||||
|
||||
class MegatronContainer(ABC):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.megatron_v2 = self.policy.is_megatron_v2
|
||||
|
||||
def _align_qkv_transposed(self, x):
|
||||
attention_head_size = x.shape[-1] // self.num_attention_heads
|
||||
new_x_shape = x.size()[:-1] + (self.num_attention_heads, attention_head_size)
|
||||
x_1 = x.view(*new_x_shape)
|
||||
(q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1))
|
||||
if len(q.shape) > 2:
|
||||
return torch.cat((q.reshape(q.shape[0], -1), k.reshape(q.shape[0], -1), v.reshape(q.shape[0], -1)),
|
||||
dim=-1).reshape(x.shape)
|
||||
else:
|
||||
return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)
|
||||
|
||||
def transpose(self):
|
||||
super().transpose()
|
||||
if self.megatron_v2:
|
||||
self.qkvw = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvw).contiguous())
|
||||
self.qkvb = torch.nn.parameter.Parameter(self._align_qkv_transposed(self.qkvb).contiguous())
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from packaging import version as pkg_version
|
||||
import torch
|
||||
|
||||
|
||||
class MetaTensorContainer(ABC):
|
||||
"""
|
||||
NOTE: If you are using this feature with a container that
|
||||
also inherits from `HybridEngineContainer`, ensure that `MetaTensorContainer`
|
||||
is inherited before `HybridEngineContainer` in the class definition.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__):
|
||||
raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+")
|
||||
super().__init__(**kwargs)
|
||||
self.is_meta = False
|
||||
self.ckpt_load_enabled = True
|
||||
|
||||
def initialize_tensors(self, enable_training=False):
|
||||
super().initialize_tensors(enable_training=enable_training)
|
||||
self.is_meta = self.qkvw.is_meta
|
||||
|
||||
def apply_tensor_parallelism(self, mp_replace, **kwargs):
|
||||
if self.is_meta:
|
||||
if self.qkvb is None:
|
||||
self.module.attention.attn_qkvb = None
|
||||
if self.dense_b is None:
|
||||
self.module.attention.attn_ob = None
|
||||
else:
|
||||
super().apply_tensor_parallelism(mp_replace, **kwargs)
|
||||
|
||||
def copy_data_to_new_module(self):
|
||||
if self.is_meta:
|
||||
if self.attn_nw is None:
|
||||
self.module.mlp.attn_nw = self.attn_nw
|
||||
self.module.mlp.attn_nb = self.attn_nb
|
||||
else:
|
||||
super().copy_data_to_new_module()
|
||||
|
||||
def transpose(self):
|
||||
if not self.is_meta:
|
||||
super().transpose()
|
||||
|
||||
@abstractmethod
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
"""
|
||||
Load all the transformer parameter from the checkpoint file (sd).
|
||||
In addition to the parameter names, we require two
|
||||
more parameters to help read the data correctly
|
||||
from the checkpoint and split the qkv heads in the
|
||||
right order:
|
||||
1. `use_load_prefix` (Default: False): this specifies
|
||||
whether we need to use the name of first abstraction
|
||||
layer of the model for searching the parameter's name
|
||||
in a checkpoint file. For more information of how this
|
||||
is used please see
|
||||
https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/load_checkpoint.py
|
||||
2. `split_qkv` (Default: True): we use this flag when splitting
|
||||
the qkv parameter into heads. If it is False, it means the heads
|
||||
of q, k, and v are stored together and needs to split in the
|
||||
DeepSpeed-Inference API.
|
||||
"""
|
||||
raise NotImplementedError("A load_params() function must be defined in the model container \
|
||||
when inheriting the MetaTensorContainer feature")
|
||||
@@ -0,0 +1,159 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import abstractmethod
|
||||
import torch
|
||||
|
||||
from .hybrid_engine import HybridEngineContainer
|
||||
|
||||
|
||||
class HybridSplitQKVContainer(HybridEngineContainer):
|
||||
|
||||
def set_attention(self, qkvw, qkvb, dense_w, dense_b):
|
||||
super().set_attention(qkvw, qkvb, dense_w, dense_b)
|
||||
self.set_q_k_v()
|
||||
|
||||
@abstractmethod
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
In `set_q_k_v`, it is necessary to populate the following variables (where appropriate)
|
||||
for the given model:
|
||||
self.qw: q weight
|
||||
self.qb: q bias
|
||||
self.kw: k weight
|
||||
self.kb: k bias
|
||||
self.vw: v weight
|
||||
self.vb: v bias
|
||||
"""
|
||||
raise NotImplementedError("A set_q_k_v() function must be defined in the model container \
|
||||
in order to set the unfused q, k, and v tensors.")
|
||||
|
||||
def attention_qkv_mp(self, mp_replace, reversed_dim=False):
|
||||
# Only need to alter
|
||||
if self.module.attention.attn_qkvw is None:
|
||||
params = [
|
||||
(self.module.attention.attn_qw, self.qw),
|
||||
(self.module.attention.attn_qb, self.qb),
|
||||
(self.module.attention.attn_kw, self.kw),
|
||||
(self.module.attention.attn_kb, self.kb),
|
||||
(self.module.attention.attn_vw, self.vw),
|
||||
(self.module.attention.attn_vb, self.vb),
|
||||
]
|
||||
for dst, src in params:
|
||||
dst = mp_replace.copy(
|
||||
dst[:self.qw.shape[0] // mp_replace.mp_size], src, int8=reversed_dim,
|
||||
allocate_tensor=reversed_dim) if src is not None else None
|
||||
else:
|
||||
super().attention_qkv_mp(mp_replace)
|
||||
|
||||
def release_qkv(self):
|
||||
super().release_qkv()
|
||||
split_qkv_params = [
|
||||
(self.module.attention.attn_qw, self.qw),
|
||||
(self.module.attention.attn_qb, self.qb),
|
||||
(self.module.attention.attn_kw, self.kw),
|
||||
(self.module.attention.attn_kb, self.kb),
|
||||
(self.module.attention.attn_vw, self.vw),
|
||||
(self.module.attention.attn_vb, self.vb),
|
||||
]
|
||||
|
||||
self._release_params(split_qkv_params)
|
||||
|
||||
def reset_qkv(self):
|
||||
self.qkvw.data[:self.qw.shape[0]] = self.qw.data
|
||||
self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data
|
||||
self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data
|
||||
|
||||
qkv_data = [self.qw.data, self.kw.data, self.vw.data]
|
||||
|
||||
self.qw.data = self.qkvw.data[:self.qw.shape[0]]
|
||||
self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]]
|
||||
self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:]
|
||||
|
||||
if self.qkvb is not None:
|
||||
self.qkvb.data[:self.qw.shape[0]] = self.qb.data
|
||||
self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data
|
||||
self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data
|
||||
|
||||
qkv_data.extend([self.qb.data, self.kb.data, self.vb.data])
|
||||
|
||||
self.qb.data = self.qkvb.data[:self.qw.shape[0]]
|
||||
self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]]
|
||||
self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:]
|
||||
|
||||
for data in qkv_data:
|
||||
del data
|
||||
|
||||
def reset_qkv_experimental(self):
|
||||
"""
|
||||
WIP - experimental and likely to be changed/improved.
|
||||
Unused by keeping for now.
|
||||
"""
|
||||
if self.module.attention.attn_qkvw is None:
|
||||
self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3,
|
||||
self.qw.shape[0],
|
||||
dtype=self.qw.dtype,
|
||||
device=self.qw.device)
|
||||
self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3,
|
||||
dtype=self.qw.dtype,
|
||||
device=self.qw.device)
|
||||
self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data
|
||||
self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data
|
||||
self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data
|
||||
self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data
|
||||
self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data
|
||||
self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data
|
||||
|
||||
qkv_data = [self.qw.data, \
|
||||
self.qb.data, \
|
||||
self.kw.data, \
|
||||
self.kb.data, \
|
||||
self.vw.data, \
|
||||
self.vb.data]
|
||||
|
||||
self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]]
|
||||
self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]]
|
||||
self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]]
|
||||
self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]]
|
||||
self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:]
|
||||
self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:]
|
||||
|
||||
for data in qkv_data:
|
||||
del data
|
||||
|
||||
def set_attn_params_wo_copy(self, Z3_enabled=False):
|
||||
self.module.attention.attn_ow = self.dense_w
|
||||
self.module.attention.attn_ob = self.dense_b
|
||||
if not Z3_enabled:
|
||||
# In initialize_tensors, we create a fused qkvw with the appropriate shape
|
||||
# and copy the qw, qb, kw, kb, vw, vb into it
|
||||
self.module.attention.attn_qkvw = self.qkvw
|
||||
self.module.attention.attn_qkvb = self.qkvb
|
||||
|
||||
# We reset the data for qw (which is the original model parameter) to point
|
||||
# to the fused weight matrix we have created here
|
||||
self.qw.data = self.qkvw[:self.qw.shape[0], :]
|
||||
self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :]
|
||||
self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :]
|
||||
|
||||
# Assume if one of the biases is not None, then all of them are not None
|
||||
if self.qb is not None:
|
||||
self.qb.data = self.qkvb[:self.qw.shape[0]]
|
||||
self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]]
|
||||
self.vb.data = self.qkvb[self.qw.shape[0] * 2:]
|
||||
else:
|
||||
# In ZeRO-3 this will be managed by ZeRO and handled separately in the
|
||||
# forward of ds_attention
|
||||
self.module.attention.attn_qw = self.qw
|
||||
self.module.attention.attn_qb = self.qb
|
||||
self.module.attention.attn_kw = self.kw
|
||||
self.module.attention.attn_kb = self.kb
|
||||
self.module.attention.attn_vw = self.vw
|
||||
self.module.attention.attn_vb = self.vb
|
||||
|
||||
def get_attn_params(self):
|
||||
params = super().get_attn_params()
|
||||
params.extend([self.qw, self.qb, self.kw, self.kb, self.vw, self.vb])
|
||||
return params
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
from ..policy import TransformerPolicy
|
||||
|
||||
|
||||
class DS_GPT2Container(BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
|
||||
class HFGPT2LayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
# HuggingFace GPT2 uses convolutional layer instead of linear layer
|
||||
super().__init__(inference, linear_layer=False)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
HFGPT2LayerPolicy._orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block
|
||||
except Exception:
|
||||
HFGPT2LayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.attn.embed_dim, \
|
||||
self.client_module.attn.num_heads, \
|
||||
self.client_module.ln_1.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
return self.client_module.attn.c_attn.weight, \
|
||||
self.client_module.attn.c_attn.bias, \
|
||||
self.client_module.attn.c_proj.weight, \
|
||||
self.client_module.attn.c_proj.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.c_fc.weight, \
|
||||
self.client_module.mlp.c_fc.bias, \
|
||||
self.client_module.mlp.c_proj.weight, \
|
||||
self.client_module.mlp.c_proj.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.ln_2.weight, \
|
||||
self.client_module.ln_2.bias, \
|
||||
self.client_module.ln_1.weight, \
|
||||
self.client_module.ln_1.bias
|
||||
@@ -0,0 +1,132 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features.meta_tensor import MetaTensorContainer
|
||||
from .features.split_qkv import HybridSplitQKVContainer
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
from ..policy import transformer_param_names
|
||||
from ..policy import maybe_copy
|
||||
from ..policy import maybe_copy_qkv
|
||||
|
||||
from ..policy import maybe_get_lora
|
||||
|
||||
|
||||
class DS_GPTJContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.fc_in, self.policy.client_module.mlp.fc_out,
|
||||
self.policy.client_module.attn.q_proj, self.policy.client_module.attn.k_proj,
|
||||
self.policy.client_module.attn.v_proj, self.policy.client_module.attn.out_proj
|
||||
]
|
||||
]
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw),
|
||||
(k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.attn.q_proj.weight
|
||||
self.qb = None
|
||||
self.kw = self.policy.client_module.attn.k_proj.weight
|
||||
self.kb = None
|
||||
self.vw = self.policy.client_module.attn.v_proj.weight
|
||||
self.vb = None
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'attn.q_proj.weight', \
|
||||
'attn.k_proj.weight', \
|
||||
'attn.v_proj.weight', \
|
||||
'attn.out_proj.weight', \
|
||||
'mlp.fc_in.weight', \
|
||||
'mlp.fc_in.bias', \
|
||||
'mlp.fc_out.weight', \
|
||||
'mlp.fc_out.bias', \
|
||||
'ln_1.weight', \
|
||||
'ln_1.bias'
|
||||
)
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(3, 4):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
for i in range(4, 8):
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i],
|
||||
prefix + param_names[i])
|
||||
for i in range(8, 10):
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i + 2],
|
||||
prefix + param_names[i])
|
||||
|
||||
|
||||
class HFGPTJLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(inference, scale_attention=True)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
HFGPTJLayerPolicy._orig_layer_class = transformers.models.gptj.modeling_gptj.GPTJBlock
|
||||
except Exception:
|
||||
HFGPTJLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.attn.embed_dim, \
|
||||
self.client_module.attn.num_attention_heads, \
|
||||
self.client_module.ln_1.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.attn.q_proj.weight
|
||||
kw = self.client_module.attn.k_proj.weight
|
||||
vw = self.client_module.attn.v_proj.weight
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
None, \
|
||||
self.client_module.attn.out_proj.weight, \
|
||||
None,
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.fc_in.weight, \
|
||||
self.client_module.mlp.fc_in.bias, \
|
||||
self.client_module.mlp.fc_out.weight, \
|
||||
self.client_module.mlp.fc_out.bias
|
||||
|
||||
def layernorm(self):
|
||||
return None, \
|
||||
None, \
|
||||
self.client_module.ln_1.weight, \
|
||||
self.client_module.ln_1.bias
|
||||
@@ -0,0 +1,145 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features.meta_tensor import MetaTensorContainer
|
||||
from .features.split_qkv import HybridSplitQKVContainer
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
from ..policy import transformer_param_names
|
||||
from ..policy import maybe_copy
|
||||
from ..policy import maybe_copy_qkv
|
||||
|
||||
from ..policy import maybe_get_lora
|
||||
|
||||
|
||||
class DS_GPTNEOContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.c_fc, self.policy.client_module.mlp.c_proj,
|
||||
self.policy.client_module.attn.attention.q_proj, self.policy.client_module.attn.attention.k_proj,
|
||||
self.policy.client_module.attn.attention.v_proj, self.policy.client_module.attn.attention.out_proj
|
||||
]
|
||||
]
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.attn.attention.q_proj.weight
|
||||
self.qb = None
|
||||
self.kw = self.policy.client_module.attn.attention.k_proj.weight
|
||||
self.kb = None
|
||||
self.vw = self.policy.client_module.attn.attention.v_proj.weight
|
||||
self.vb = None
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw),
|
||||
(k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'attn.attention.q_proj.weight', \
|
||||
'attn.attention.k_proj.weight', \
|
||||
'attn.attention.v_proj.weight', \
|
||||
'attn.attention.out_proj.weight', \
|
||||
'attn.attention.out_proj.bias', \
|
||||
'mlp.c_fc.weight', \
|
||||
'mlp.c_fc.bias', \
|
||||
'mlp.c_proj.weight', \
|
||||
'mlp.c_proj.bias', \
|
||||
'ln_2.weight', \
|
||||
'ln_2.bias', \
|
||||
'ln_1.weight', \
|
||||
'ln_1.bias'
|
||||
)
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(3, 5):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
for i in range(5, 11):
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
for i in range(11, 13):
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
|
||||
|
||||
class HFGPTNEOLayerPolicy(TransformerPolicy):
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(inference, scale_attention=False)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
HFGPTNEOLayerPolicy._orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock
|
||||
except Exception:
|
||||
HFGPTNEOLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.attn.attention.embed_dim, \
|
||||
self.client_module.attn.attention.num_heads, \
|
||||
self.client_module.ln_1.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def get_q_k_v(self):
|
||||
return self.client_module.attn.attention.q_proj.weight, \
|
||||
None, \
|
||||
self.client_module.attn.attention.k_proj.weight, \
|
||||
None, \
|
||||
self.client_module.attn.attention.v_proj.weight, \
|
||||
None
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.attn.attention.q_proj.weight
|
||||
kw = self.client_module.attn.attention.k_proj.weight
|
||||
vw = self.client_module.attn.attention.v_proj.weight
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
None, \
|
||||
self.client_module.attn.attention.out_proj.weight, \
|
||||
self.client_module.attn.attention.out_proj.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.c_fc.weight, \
|
||||
self.client_module.mlp.c_fc.bias, \
|
||||
self.client_module.mlp.c_proj.weight, \
|
||||
self.client_module.mlp.c_proj.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.ln_2.weight, \
|
||||
self.client_module.ln_2.bias, \
|
||||
self.client_module.ln_1.weight, \
|
||||
self.client_module.ln_1.bias
|
||||
@@ -0,0 +1,146 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features.meta_tensor import MetaTensorContainer
|
||||
from .features.hybrid_megatron import HybridMegatronContainer
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
import torch
|
||||
from ..policy import TransformerPolicy
|
||||
from ..policy import transformer_param_names
|
||||
from ..policy import maybe_copy
|
||||
from packaging import version as pkg_version
|
||||
|
||||
from ..policy import maybe_get_lora
|
||||
|
||||
|
||||
class DS_GPTNEOXContainer(MetaTensorContainer, HybridMegatronContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
|
||||
if self.megatron_v2:
|
||||
self.module.config.rotate_half = True
|
||||
self.module.config.rotate_every_two = False
|
||||
|
||||
return self.module
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
fc1_lora, fc2_lora, qkv_lora, out_lora = self.get_lora_params()
|
||||
ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (qkv_lora, self.qkvw), (out_lora, self.dense_w)]
|
||||
return ret
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
if GPTNEOXLayerPolicy.version == 0:
|
||||
attention = self.policy.client_module.attention
|
||||
else:
|
||||
attention = self.policy.client_module.self_attention
|
||||
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h,
|
||||
attention.query_key_value, attention.dense
|
||||
]
|
||||
]
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'attention.query_key_value.weight', \
|
||||
'attention.query_key_value.bias', \
|
||||
'attention.dense.weight', \
|
||||
'attention.dense.bias', \
|
||||
'mlp.dense_h_to_4h.weight', \
|
||||
'mlp.dense_h_to_4h.bias', \
|
||||
'mlp.dense_4h_to_h.weight', \
|
||||
'mlp.dense_4h_to_h.bias', \
|
||||
'post_attention_layernorm.weight', \
|
||||
'post_attention_layernorm.bias', \
|
||||
'input_layernorm.weight', \
|
||||
'input_layernorm.bias'
|
||||
)
|
||||
for i in range(0, 2):
|
||||
maybe_copy(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
transformer_param_names[i],
|
||||
prefix + param_names[i],
|
||||
qkv=True,
|
||||
megatron_v2=self.policy.is_megatron_v2,
|
||||
split_qkv=self.policy.split_qkv,
|
||||
heads=self.policy.client_module.attention.num_attention_heads)
|
||||
for i in range(2, 4):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i],
|
||||
prefix + param_names[i])
|
||||
for i in range(4, 10):
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i],
|
||||
prefix + param_names[i])
|
||||
for i in range(10, 12):
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i], prefix + param_names[i])
|
||||
|
||||
|
||||
class GPTNEOXLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
version = 0
|
||||
|
||||
def __init__(self, client_module, inference=True, megatron_v2=True, split_qkv=False):
|
||||
super().__init__(inference, megatron_v2=megatron_v2, split_qkv=split_qkv)
|
||||
self.client_module = client_module
|
||||
if GPTNEOXLayerPolicy._orig_layer_class is None:
|
||||
if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"):
|
||||
GPTNEOXLayerPolicy._orig_layer_class = None
|
||||
else:
|
||||
try:
|
||||
from transformers import GPTNeoXLayer
|
||||
GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer
|
||||
except ImportError:
|
||||
GPTNEOXLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
if GPTNEOXLayerPolicy.version == 0:
|
||||
attention = self.client_module.attention
|
||||
else:
|
||||
attention = self.client_module.self_attention
|
||||
|
||||
return self.client_module.attention.hidden_size, \
|
||||
self.client_module.attention.num_attention_heads, \
|
||||
self.client_module.input_layernorm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
if GPTNEOXLayerPolicy.version == 0:
|
||||
attention = self.client_module.attention
|
||||
else:
|
||||
attention = self.client_module.self_attention
|
||||
|
||||
return attention.query_key_value.weight, \
|
||||
attention.query_key_value.bias, \
|
||||
attention.dense.weight, \
|
||||
attention.dense.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.mlp.dense_h_to_4h.weight, \
|
||||
self.client_module.mlp.dense_h_to_4h.bias, \
|
||||
self.client_module.mlp.dense_4h_to_h.weight, \
|
||||
self.client_module.mlp.dense_4h_to_h.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.post_attention_layernorm.weight, \
|
||||
self.client_module.post_attention_layernorm.bias, \
|
||||
self.client_module.input_layernorm.weight, \
|
||||
self.client_module.input_layernorm.bias
|
||||
@@ -0,0 +1,181 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import importlib
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
from deepspeed.utils.types import ActivationFuncType, NormType
|
||||
|
||||
from ..policy import (TransformerPolicy, maybe_copy, maybe_copy_geglu, maybe_copy_qkv, maybe_get_lora,
|
||||
transformer_param_names)
|
||||
from .base import *
|
||||
from .features import HybridGatedMLPContainer, HybridSplitQKVContainer
|
||||
|
||||
|
||||
class DS_InternLMContainer(HybridGatedMLPContainer, HybridSplitQKVContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
|
||||
_config.rotate_half = True
|
||||
_config.rotate_every_two = False
|
||||
_config.rotary_dim = self.hidden_size // self.num_attention_heads
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight,
|
||||
self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight,
|
||||
self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight,
|
||||
self.policy.client_module.self_attn.o_proj.weight
|
||||
]
|
||||
]
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w),
|
||||
(out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.self_attn.q_proj.weight
|
||||
self.qb = self.policy.client_module.self_attn.q_proj.bias
|
||||
self.kw = self.policy.client_module.self_attn.k_proj.weight
|
||||
self.kb = self.policy.client_module.self_attn.k_proj.bias
|
||||
self.vw = self.policy.client_module.self_attn.v_proj.weight
|
||||
self.vb = self.policy.client_module.self_attn.v_proj.bias
|
||||
|
||||
def set_mlp_gate(self):
|
||||
"""
|
||||
Necessary to implement for `HybridGatedMLPContainer`
|
||||
"""
|
||||
self.inter_up_w = self.policy.client_module.mlp.up_proj.weight
|
||||
self.inter_up_b = None
|
||||
self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight
|
||||
self.inter_gate_b = None
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'self_attn.q_proj.weight', \
|
||||
'self_attn.k_proj.weight', \
|
||||
'self_attn.v_proj.weight', \
|
||||
'self_attn.o_proj.weight', \
|
||||
'mlp.up_proj.weight', \
|
||||
'mlp.gate_proj.weight', \
|
||||
'mlp.down_proj.weight', \
|
||||
'input_layernorm.weight', \
|
||||
'post_attention_layernorm.weight'
|
||||
'self_attn.q_proj.bias', \
|
||||
'self_attn.k_proj.bias', \
|
||||
'self_attn.v_proj.bias', \
|
||||
'self_attn.o_proj.bias', \
|
||||
)
|
||||
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvb', [prefix + param_names[9], prefix + param_names[10], prefix + param_names[11]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[2],
|
||||
prefix + param_names[3])
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[3],
|
||||
prefix + param_names[12])
|
||||
maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w',
|
||||
[prefix + param_names[4], prefix + param_names[5]])
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6])
|
||||
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7])
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8])
|
||||
|
||||
|
||||
class InternLMLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = []
|
||||
_orig_layer_class_inited = False
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(
|
||||
inference,
|
||||
mlp_act_func_type=ActivationFuncType.GATED_SILU,
|
||||
norm_type=NormType.RMSNorm,
|
||||
)
|
||||
self.client_module = client_module
|
||||
|
||||
self._init_orig_layer_class_once()
|
||||
|
||||
def _init_orig_layer_class_once(self):
|
||||
if InternLMLayerPolicy._orig_layer_class_inited:
|
||||
return
|
||||
|
||||
for sub_pkg in ['', '.internlm-7b', '.internlm-chat-7b']:
|
||||
try:
|
||||
from transformers.utils import TRANSFORMERS_DYNAMIC_MODULE_NAME
|
||||
module = importlib.import_module(f"{TRANSFORMERS_DYNAMIC_MODULE_NAME}{sub_pkg}.modeling_internlm")
|
||||
if module.InternLMDecoderLayer not in InternLMLayerPolicy._orig_layer_class:
|
||||
InternLMLayerPolicy._orig_layer_class.append(module.InternLMDecoderLayer)
|
||||
except ImportError:
|
||||
continue
|
||||
|
||||
InternLMLayerPolicy._orig_layer_class_inited = True
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.self_attn.q_proj.weight.shape[1], \
|
||||
self.client_module.self_attn.num_heads, \
|
||||
self.client_module.input_layernorm.variance_epsilon, \
|
||||
self.client_module.mlp.gate_proj.weight.shape[0]
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.self_attn.q_proj.weight
|
||||
kw = self.client_module.self_attn.k_proj.weight
|
||||
vw = self.client_module.self_attn.v_proj.weight
|
||||
qb = self.client_module.self_attn.q_proj.bias
|
||||
kb = self.client_module.self_attn.k_proj.bias
|
||||
vb = self.client_module.self_attn.v_proj.bias
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
qkvb, \
|
||||
self.client_module.self_attn.o_proj.weight, \
|
||||
self.client_module.self_attn.o_proj.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
mlp1_up = self.client_module.mlp.up_proj.weight
|
||||
mlp1_gate = self.client_module.mlp.gate_proj.weight
|
||||
mlp2 = self.client_module.mlp.down_proj.weight
|
||||
|
||||
mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training)
|
||||
|
||||
return mlp1, None, mlp2, None
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.post_attention_layernorm.weight, \
|
||||
None, \
|
||||
self.client_module.input_layernorm.weight, \
|
||||
None
|
||||
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer
|
||||
from deepspeed.utils.types import ActivationFuncType, NormType
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from ..policy import (
|
||||
TransformerPolicy,
|
||||
transformer_param_names,
|
||||
maybe_copy,
|
||||
maybe_copy_qkv,
|
||||
maybe_copy_geglu,
|
||||
maybe_get_lora,
|
||||
)
|
||||
|
||||
|
||||
class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer,
|
||||
BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
|
||||
_config.rotate_half = True
|
||||
_config.rotate_every_two = False
|
||||
_config.rotary_dim = self.hidden_size // self.num_attention_heads
|
||||
if hasattr(self.policy.client_module.self_attn, 'config'):
|
||||
_config.rope_theta = self.policy.client_module.self_attn.config.rope_theta
|
||||
else:
|
||||
_config.rope_theta = self.policy.client_module.self_attn.rope_theta
|
||||
self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
|
||||
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight,
|
||||
self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight,
|
||||
self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight,
|
||||
self.policy.client_module.self_attn.o_proj.weight
|
||||
]
|
||||
]
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w),
|
||||
(out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.self_attn.q_proj.weight
|
||||
self.qb = None
|
||||
self.kw = self.policy.client_module.self_attn.k_proj.weight
|
||||
self.kb = None
|
||||
self.vw = self.policy.client_module.self_attn.v_proj.weight
|
||||
self.vb = None
|
||||
|
||||
def set_mlp_gate(self):
|
||||
"""
|
||||
Necessary to implement for `HybridGatedMLPContainer`
|
||||
"""
|
||||
self.inter_up_w = self.policy.client_module.mlp.up_proj.weight
|
||||
self.inter_up_b = None
|
||||
self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight
|
||||
self.inter_gate_b = None
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'self_attn.q_proj.weight', \
|
||||
'self_attn.k_proj.weight', \
|
||||
'self_attn.v_proj.weight', \
|
||||
'self_attn.o_proj.weight', \
|
||||
'mlp.up_proj.weight', \
|
||||
'mlp.gate_proj.weight', \
|
||||
'mlp.down_proj.weight', \
|
||||
'post_attention_layernorm.weight', \
|
||||
'input_layernorm.weight',
|
||||
)
|
||||
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(3, 4):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w',
|
||||
[prefix + param_names[4], prefix + param_names[5]])
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6])
|
||||
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7])
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8])
|
||||
|
||||
# This line is necessary for proper output when kernels + meta tensors are used in Llama models
|
||||
# TODO: Investigate root-cause and fix meta tensor loading
|
||||
module.mlp.output_b = None
|
||||
|
||||
|
||||
class LLAMALayerPolicy(TransformerPolicy):
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(
|
||||
inference,
|
||||
mlp_act_func_type=ActivationFuncType.GATED_SILU,
|
||||
norm_type=NormType.RMSNorm,
|
||||
)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
LLAMALayerPolicy._orig_layer_class = transformers.models.llama.modeling_llama.LlamaDecoderLayer # type: ignore
|
||||
except Exception:
|
||||
LLAMALayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
if hasattr(self.client_module.self_attn, 'config'):
|
||||
num_heads = self.client_module.self_attn.config.num_attention_heads
|
||||
else:
|
||||
num_heads = self.client_module.self_attn.num_heads
|
||||
hidden_heads = (
|
||||
self.client_module.self_attn.q_proj.in_features,
|
||||
num_heads,
|
||||
self.client_module.input_layernorm.variance_epsilon,
|
||||
self.client_module.mlp.gate_proj.out_features,
|
||||
)
|
||||
return hidden_heads
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.self_attn.q_proj.weight
|
||||
kw = self.client_module.self_attn.k_proj.weight
|
||||
vw = self.client_module.self_attn.v_proj.weight
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
None, \
|
||||
self.client_module.self_attn.o_proj.weight, \
|
||||
None
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
mlp1_up = self.client_module.mlp.up_proj.weight
|
||||
mlp1_gate = self.client_module.mlp.gate_proj.weight
|
||||
mlp2 = self.client_module.mlp.down_proj.weight
|
||||
|
||||
mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training)
|
||||
|
||||
return mlp1, None, mlp2, None
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.post_attention_layernorm.weight, \
|
||||
None, \
|
||||
self.client_module.input_layernorm.weight, \
|
||||
None
|
||||
@@ -0,0 +1,158 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features import HybridSplitQKVContainer, HybridGatedMLPContainer, MetaTensorContainer
|
||||
from deepspeed.utils.types import ActivationFuncType, NormType
|
||||
from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from ..policy import (
|
||||
TransformerPolicy,
|
||||
transformer_param_names,
|
||||
maybe_copy,
|
||||
maybe_copy_qkv,
|
||||
maybe_copy_geglu,
|
||||
maybe_get_lora,
|
||||
)
|
||||
|
||||
|
||||
class DS_LLAMA2Container(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer,
|
||||
BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
|
||||
_config.rotate_half = False
|
||||
_config.rotate_every_two = True
|
||||
_config.rotary_dim = self.hidden_size // self.num_attention_heads
|
||||
_config.num_kv = self.policy.client_module.attention.n_kv_heads
|
||||
self.module = DeepSpeedLlama2Inference(_config, mp_group=self.mp_group)
|
||||
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.feed_forward.w3.weight, self.policy.client_module.feed_forward.w1.weight,
|
||||
self.policy.client_module.feed_forward.w2.weight, self.policy.client_module.attention.wq.weight,
|
||||
self.policy.client_module.attention.wk.weight, self.policy.client_module.attention.wv.weight,
|
||||
self.policy.client_module.attention.wo.weight
|
||||
]
|
||||
]
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
up_proj_lora, gate_proj_lora, down_proj_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(up_proj_lora, self.inter_up_w), (gate_proj_lora, self.inter_gate_w), (down_proj_lora, self._4hh_w),
|
||||
(out_lora, self.dense_w), (q_lora, self.qw), (k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.attention.wq.weight
|
||||
self.qb = None
|
||||
self.kw = self.policy.client_module.attention.wk.weight
|
||||
self.kb = None
|
||||
self.vw = self.policy.client_module.attention.wv.weight
|
||||
self.vb = None
|
||||
|
||||
def set_mlp_gate(self):
|
||||
"""
|
||||
Necessary to implement for `HybridGatedMLPContainer`
|
||||
"""
|
||||
self.inter_up_w = self.policy.client_module.feed_forward.w2.weight
|
||||
self.inter_up_b = None
|
||||
self.inter_gate_w = self.policy.client_module.feed_forward.w1.weight
|
||||
self.inter_gate_b = None
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'attention.wq.weight', \
|
||||
'attention.wk.weight', \
|
||||
'attention.wv.weight', \
|
||||
'attention.wo.weight', \
|
||||
'feed_forward.w3.weight', \
|
||||
'feed_forward.w1.weight', \
|
||||
'feed_forward.w2.weight', \
|
||||
'ffn_norm.weight', \
|
||||
'attention_norm.weight'
|
||||
)
|
||||
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(3, 4):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1],
|
||||
prefix + param_names[i])
|
||||
maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w',
|
||||
[prefix + param_names[4], prefix + param_names[5]])
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6])
|
||||
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7])
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8])
|
||||
|
||||
|
||||
class LLAMA2LayerPolicy(TransformerPolicy):
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(
|
||||
inference,
|
||||
mlp_act_func_type=ActivationFuncType.GATED_SILU,
|
||||
norm_type=NormType.RMSNorm,
|
||||
)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import llama
|
||||
LLAMA2LayerPolicy._orig_layer_class = llama.model.TransformerBlock # type: ignore
|
||||
except Exception:
|
||||
LLAMA2LayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.attention.wq.weight.shape[1], \
|
||||
self.client_module.n_heads, \
|
||||
self.client_module.ffn_norm.eps, \
|
||||
(self.client_module.feed_forward.w1.weight.shape[0] * \
|
||||
deepspeed.comm.get_world_size() if deepspeed.comm.is_initialized() else 1) # this is a hack to inject when model is already partitioned!
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.attention.wq.weight
|
||||
kw = self.client_module.attention.wk.weight
|
||||
vw = self.client_module.attention.wv.weight
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
|
||||
return qkvw, \
|
||||
None, \
|
||||
self.client_module.attention.wo.weight, \
|
||||
None
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
mlp1_up = self.client_module.feed_forward.w3.weight
|
||||
mlp1_gate = self.client_module.feed_forward.w1.weight
|
||||
mlp2 = self.client_module.feed_forward.w2.weight
|
||||
|
||||
mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training)
|
||||
|
||||
return mlp1, None, mlp2, None
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.ffn_norm.weight, \
|
||||
None, \
|
||||
self.client_module.attention_norm.weight, \
|
||||
None
|
||||
@@ -0,0 +1,119 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features.megatron import MegatronContainer
|
||||
from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference
|
||||
import torch
|
||||
from ..policy import TransformerPolicy
|
||||
from packaging import version as pkg_version
|
||||
|
||||
|
||||
class DS_MegatronGPTContainer(MegatronContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
|
||||
if self.megatron_v2:
|
||||
self.module.config.rotate_half = True
|
||||
self.module.config.rotate_every_two = False
|
||||
|
||||
return self.module
|
||||
|
||||
|
||||
# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp
|
||||
# TODO: Generalize MoE overall goal, expand beyond Megatron
|
||||
class MegatronLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
version = 0
|
||||
moe_type = 'standard'
|
||||
megatron_v2 = True
|
||||
use_mup = False
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(inference, megatron_v2=MegatronLayerPolicy.megatron_v2, use_mup=MegatronLayerPolicy.use_mup)
|
||||
self.client_module = client_module
|
||||
# we use megatron version to differentiate between the old and new
|
||||
# megatron-lm source code
|
||||
if MegatronLayerPolicy._orig_layer_class is None:
|
||||
if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"):
|
||||
MegatronLayerPolicy._orig_layer_class = None
|
||||
else:
|
||||
try:
|
||||
from megatron.model.transformer import ParallelTransformerLayer
|
||||
MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer
|
||||
MegatronLayerPolicy.version = 1
|
||||
except ImportError:
|
||||
MegatronLayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_hidden_heads(self):
|
||||
if MegatronLayerPolicy.version == 0:
|
||||
return self.client_module.attention.query_key_value.weight.shape[1], \
|
||||
self.client_module.attention.num_attention_heads, \
|
||||
self.client_module.input_layernorm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
else:
|
||||
return self.client_module.self_attention.query_key_value.weight.shape[1], \
|
||||
self.client_module.self_attention.num_attention_heads, \
|
||||
self.client_module.input_layernorm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
if self.inference:
|
||||
if MegatronLayerPolicy.version == 0:
|
||||
attention = self.client_module.attention
|
||||
else:
|
||||
attention = self.client_module.self_attention
|
||||
else:
|
||||
return None
|
||||
|
||||
return attention.query_key_value.weight, \
|
||||
attention.query_key_value.bias, \
|
||||
attention.dense.weight, \
|
||||
attention.dense.bias
|
||||
|
||||
def mlp(self, moe_type='standard', enable_training=False):
|
||||
from deepspeed.moe.utils import has_moe_layers
|
||||
moe, _ = has_moe_layers(self.client_module)
|
||||
|
||||
if moe:
|
||||
moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \
|
||||
self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts
|
||||
num_experts = len(moe_experts)
|
||||
if moe_type == 'standard':
|
||||
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)]
|
||||
else:
|
||||
|
||||
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \
|
||||
self.client_module.mlp.mlp.dense_h_to_4h.weight, \
|
||||
self.client_module.mlp.mlp.dense_h_to_4h.bias, \
|
||||
self.client_module.mlp.mlp.dense_4h_to_h.weight, \
|
||||
self.client_module.mlp.mlp.dense_4h_to_h.bias, \
|
||||
self.client_module.mlp.coefficient.weight
|
||||
|
||||
else:
|
||||
return self.client_module.mlp.dense_h_to_4h.weight, \
|
||||
self.client_module.mlp.dense_h_to_4h.bias, \
|
||||
self.client_module.mlp.dense_4h_to_h.weight, \
|
||||
self.client_module.mlp.dense_4h_to_h.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.post_attention_layernorm.weight, \
|
||||
self.client_module.post_attention_layernorm.bias, \
|
||||
self.client_module.input_layernorm.weight, \
|
||||
self.client_module.input_layernorm.bias
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .base_moe import *
|
||||
from .features.megatron import MegatronContainer
|
||||
from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference
|
||||
import torch
|
||||
from .megatron_gpt import MegatronLayerPolicy
|
||||
from packaging import version as pkg_version
|
||||
|
||||
|
||||
class DS_MegatronGPTMoEContainer(MegatronContainer, BaseTransformerMoEContainer):
|
||||
|
||||
def __init__(self, policy, config, model_config, layer_id):
|
||||
super().__init__(policy, config, model_config, layer_id)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedMegatronGPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
|
||||
if self.megatron_v2:
|
||||
self.module.config.rotate_half = True
|
||||
self.module.config.rotate_every_two = False
|
||||
|
||||
return self.module
|
||||
|
||||
|
||||
# TODO: Megatron GPT MoE inherits from Megatron policy and replaces mlp
|
||||
# TODO: Generalize MoE overall goal, expand beyond Megatron
|
||||
class MegatronMoELayerPolicy(MegatronLayerPolicy):
|
||||
_orig_layer_class = None
|
||||
version = 0
|
||||
moe_type = 'standard'
|
||||
num_experts = 1
|
||||
|
||||
def __init__(self, client_module, inference=True):
|
||||
super().__init__(inference)
|
||||
self.client_module = client_module
|
||||
# we use megatron version to differentiate between the old and new
|
||||
# megatron-lm source code
|
||||
if MegatronMoELayerPolicy._orig_layer_class is None:
|
||||
if pkg_version.parse(torch.__version__) <= pkg_version.parse("1.2"):
|
||||
MegatronMoELayerPolicy._orig_layer_class = None
|
||||
else:
|
||||
try:
|
||||
from megatron.model.transformer import ParallelTransformerLayer
|
||||
MegatronMoELayerPolicy._orig_layer_class = ParallelTransformerLayer
|
||||
except ImportError:
|
||||
MegatronMoELayerPolicy._orig_layer_class = None
|
||||
|
||||
def get_num_experts(self):
|
||||
return self.num_experts
|
||||
|
||||
def mlp(self, moe_type='standard', enable_training=False):
|
||||
# for now, all of this is tightly coupled to megatron-deepspeed moe implementation
|
||||
# todo: think and refactor this to be more general
|
||||
|
||||
#from deepspeed.moe.utils import has_moe_layers
|
||||
#moe, _ = has_moe_layers(self.client_module)
|
||||
|
||||
moe_experts = self.client_module.mlp.deepspeed_moe.experts.deepspeed_experts if moe_type == 'standard' else \
|
||||
self.client_module.mlp.moe.deepspeed_moe.experts.deepspeed_experts
|
||||
num_experts = len(moe_experts)
|
||||
self.num_experts = num_experts
|
||||
|
||||
if moe_type == 'standard':
|
||||
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)]
|
||||
else:
|
||||
return [moe_experts[i].dense_h_to_4h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_h_to_4h.bias for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.weight for i in range(num_experts)], \
|
||||
[moe_experts[i].dense_4h_to_h.bias for i in range(num_experts)], \
|
||||
self.client_module.mlp.mlp.dense_h_to_4h.weight, \
|
||||
self.client_module.mlp.mlp.dense_h_to_4h.bias, \
|
||||
self.client_module.mlp.mlp.dense_4h_to_h.weight, \
|
||||
self.client_module.mlp.mlp.dense_4h_to_h.bias, \
|
||||
self.client_module.mlp.coefficient.weight
|
||||
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .base import *
|
||||
from .features import MetaTensorContainer, HybridSplitQKVContainer
|
||||
from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
from ..policy import TransformerPolicy
|
||||
from ..policy import transformer_param_names
|
||||
from ..policy import maybe_copy
|
||||
from ..policy import maybe_copy_qkv
|
||||
from ..policy import maybe_get_lora
|
||||
from deepspeed.utils.types import ActivationFuncType
|
||||
|
||||
|
||||
class DS_OPTContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# All model specific things should be defined here instead of the base class.
|
||||
|
||||
def create_module(self, config=None):
|
||||
_config = config if config is not None else self.ds_model_config
|
||||
self.module = DeepSpeedOPTInference(_config, mp_group=self.mp_group)
|
||||
self.module.config.scale_attention = self.scale_attention
|
||||
return self.module
|
||||
|
||||
def set_lora_params(self):
|
||||
"""
|
||||
Necessary to implement for `HybridEngineContainer`
|
||||
"""
|
||||
self.lora_params = [
|
||||
maybe_get_lora(p) for p in [
|
||||
self.policy.client_module.fc1,
|
||||
self.policy.client_module.fc2,
|
||||
self.policy.client_module.self_attn.q_proj,
|
||||
self.policy.client_module.self_attn.k_proj,
|
||||
self.policy.client_module.self_attn.v_proj,
|
||||
self.policy.client_module.self_attn.out_proj,
|
||||
]
|
||||
]
|
||||
|
||||
def set_q_k_v(self):
|
||||
"""
|
||||
Necessary to implement for `HybridSplitQKVContainer`
|
||||
"""
|
||||
self.qw = self.policy.client_module.self_attn.q_proj.weight
|
||||
self.qb = self.policy.client_module.self_attn.q_proj.bias
|
||||
self.kw = self.policy.client_module.self_attn.k_proj.weight
|
||||
self.kb = self.policy.client_module.self_attn.k_proj.bias
|
||||
self.vw = self.policy.client_module.self_attn.v_proj.weight
|
||||
self.vb = self.policy.client_module.self_attn.v_proj.bias
|
||||
|
||||
def get_lora_matched_pair(self):
|
||||
fc1_lora, fc2_lora, q_lora, k_lora, v_lora, out_lora = self.get_lora_params()
|
||||
ret = [(fc1_lora, self._h4h_w), (fc2_lora, self._4hh_w), (out_lora, self.dense_w), (q_lora, self.qw),
|
||||
(k_lora, self.kw), (v_lora, self.vw)]
|
||||
return ret
|
||||
|
||||
def load_params(self, module, sd, weight_quantizer, mp_replace, prefix):
|
||||
param_names = (
|
||||
'self_attn.q_proj.weight', \
|
||||
'self_attn.k_proj.weight', \
|
||||
'self_attn.v_proj.weight', \
|
||||
'self_attn.q_proj.bias', \
|
||||
'self_attn.k_proj.bias', \
|
||||
'self_attn.v_proj.bias', \
|
||||
'self_attn.out_proj.weight', \
|
||||
'self_attn.out_proj.bias', \
|
||||
'fc1.weight', \
|
||||
'fc1.bias', \
|
||||
'fc2.weight', \
|
||||
'fc2.bias', \
|
||||
'final_layer_norm.weight', \
|
||||
'final_layer_norm.bias', \
|
||||
'self_attn_layer_norm.weight', \
|
||||
'self_attn_layer_norm.bias'
|
||||
)
|
||||
|
||||
for i in range(0, 6, 3):
|
||||
maybe_copy_qkv(module.attention,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
transformer_param_names[i // 3],
|
||||
[prefix + param_names[i], prefix + param_names[i + 1], prefix + param_names[i + 2]],
|
||||
split_qkv=self.policy.split_qkv)
|
||||
for i in range(6, 8):
|
||||
maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4],
|
||||
prefix + param_names[i])
|
||||
for i in range(8, 14):
|
||||
maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4],
|
||||
prefix + param_names[i])
|
||||
for i in range(14, 16):
|
||||
maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[i - 4],
|
||||
prefix + param_names[i])
|
||||
|
||||
|
||||
class HFOPTLayerPolicy(TransformerPolicy):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self, client_module, inference=True, use_load_prefix=True):
|
||||
super().__init__(inference, linear_layer=True, pre_attn_norm=True, use_load_prefix=use_load_prefix)
|
||||
self.client_module = client_module
|
||||
try:
|
||||
import transformers
|
||||
HFOPTLayerPolicy._orig_layer_class = transformers.models.opt.modeling_opt.OPTDecoderLayer
|
||||
except Exception:
|
||||
HFOPTLayerPolicy._orig_layer_class = None
|
||||
|
||||
if hasattr(TransformerPolicy, "hf_model_config") and hasattr(TransformerPolicy.hf_model_config,
|
||||
"activation_function"):
|
||||
if TransformerPolicy.hf_model_config.activation_function == "relu":
|
||||
self.mlp_act_func_type = ActivationFuncType.ReLU
|
||||
elif TransformerPolicy.hf_model_config.activation_function in ["gelu", "gelu_new"]:
|
||||
self.mlp_act_func_type = ActivationFuncType.GELU
|
||||
else:
|
||||
raise ValueError("Unsupported activation function: {}".format(
|
||||
TransformerPolicy.hf_model_config.activation_function))
|
||||
else:
|
||||
self.mlp_act_func_type = ActivationFuncType.ReLU # default
|
||||
|
||||
def get_hidden_heads(self):
|
||||
return self.client_module.self_attn.embed_dim, \
|
||||
self.client_module.self_attn.num_heads, \
|
||||
self.client_module.self_attn_layer_norm.eps, \
|
||||
DEFAULT_INTERMEDIATE_SIZE
|
||||
|
||||
def attention(self, enable_training=False):
|
||||
qw = self.client_module.self_attn.q_proj.weight
|
||||
qb = self.client_module.self_attn.q_proj.bias
|
||||
|
||||
kw = self.client_module.self_attn.k_proj.weight
|
||||
kb = self.client_module.self_attn.k_proj.bias
|
||||
|
||||
vw = self.client_module.self_attn.v_proj.weight
|
||||
vb = self.client_module.self_attn.v_proj.bias
|
||||
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training)
|
||||
qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training)
|
||||
return qkvw, \
|
||||
qkvb, \
|
||||
self.client_module.self_attn.out_proj.weight, \
|
||||
self.client_module.self_attn.out_proj.bias
|
||||
|
||||
def mlp(self, enable_training=False):
|
||||
return self.client_module.fc1.weight, \
|
||||
self.client_module.fc1.bias, \
|
||||
self.client_module.fc2.weight, \
|
||||
self.client_module.fc2.bias
|
||||
|
||||
def layernorm(self):
|
||||
return self.client_module.final_layer_norm.weight, \
|
||||
self.client_module.final_layer_norm.bias, \
|
||||
self.client_module.self_attn_layer_norm.weight, \
|
||||
self.client_module.self_attn_layer_norm.bias
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from ..policy import DSPolicy
|
||||
from ...model_implementations.diffusers.unet import DSUNet
|
||||
|
||||
|
||||
class UNetPolicy(DSPolicy):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
try:
|
||||
import diffusers
|
||||
self._orig_layer_class = diffusers.models.unet_2d_condition.UNet2DConditionModel
|
||||
except AttributeError:
|
||||
self._orig_layer_class = diffusers.models.unets.unet_2d_condition.UNet2DConditionModel
|
||||
except ImportError:
|
||||
self._orig_layer_class = None
|
||||
|
||||
def match(self, module):
|
||||
return isinstance(module, self._orig_layer_class)
|
||||
|
||||
def match_replaced(self, module):
|
||||
return isinstance(module, DSUNet)
|
||||
|
||||
def apply(self, module, enable_cuda_graph=True):
|
||||
# TODO(cmikeh2): Enable cuda graph should be an inference configuration
|
||||
return DSUNet(module, enable_cuda_graph=enable_cuda_graph)
|
||||
|
||||
def attention(self, client_module):
|
||||
qw = client_module.to_q.weight
|
||||
kw = client_module.to_k.weight
|
||||
vw = client_module.to_v.weight
|
||||
|
||||
if qw.shape[1] == kw.shape[1]:
|
||||
qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=False)
|
||||
|
||||
return qkvw, \
|
||||
client_module.to_out[0].weight, \
|
||||
client_module.to_out[0].bias, \
|
||||
qw.shape[-1], \
|
||||
client_module.heads
|
||||
else:
|
||||
#return None
|
||||
#kvw = Parameter(torch.cat((kw, vw), dim=0), requires_grad=False)
|
||||
return qw, \
|
||||
kw, vw, \
|
||||
client_module.to_out[0].weight, \
|
||||
client_module.to_out[0].bias, \
|
||||
qw.shape[-1], \
|
||||
client_module.heads
|
||||
@@ -0,0 +1,42 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from ..policy import DSPolicy
|
||||
from ...model_implementations.diffusers.vae import DSVAE
|
||||
|
||||
|
||||
class VAEPolicy(DSPolicy):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
try:
|
||||
import diffusers
|
||||
if hasattr(diffusers.models, "autoencoders"):
|
||||
# Diffusers >= 0.25.0
|
||||
# Changes location to 'autoencoders' directory
|
||||
self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL
|
||||
elif hasattr(diffusers.models.vae, "AutoencoderKL"):
|
||||
# Diffusers < 0.12.0
|
||||
self._orig_layer_class = diffusers.models.vae.AutoencoderKL
|
||||
else:
|
||||
# Diffusers >= 0.12.0 & < 0.25.0
|
||||
# Changes location of AutoencoderKL
|
||||
self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL
|
||||
except ImportError:
|
||||
self._orig_layer_class = None
|
||||
|
||||
def match(self, module):
|
||||
return isinstance(module, self._orig_layer_class)
|
||||
|
||||
def match_replaced(self, module):
|
||||
return isinstance(module, DSVAE)
|
||||
|
||||
def apply(self, module, enable_cuda_graph=True):
|
||||
# TODO(cmikeh2): Enable cuda graph should be an inference configuration
|
||||
return DSVAE(module, enable_cuda_graph=enable_cuda_graph)
|
||||
|
||||
# NOTE (lekurile): Should we have a diffusers policy class?
|
||||
def attention(self, client_module):
|
||||
pass
|
||||
@@ -0,0 +1,231 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
import torch
|
||||
from deepspeed.utils.logging import warning_once
|
||||
from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list, get_num_kv_heads, get_n_embd, get_num_attention_heads
|
||||
|
||||
|
||||
def split_by_qkvlist_and_refuse(qkv_list, split_size, split_dim=0, cat_dim=0):
|
||||
qkv_split_list = [torch.split(mat, split_size, dim=split_dim) for mat in qkv_list]
|
||||
tp_fusedqkv_list = [
|
||||
torch.cat([qkv_s[i] for qkv_s in qkv_split_list], dim=cat_dim) for i in range(len(qkv_split_list[0]))
|
||||
]
|
||||
return tp_fusedqkv_list
|
||||
|
||||
|
||||
def require_tp_fused_qkvw(name, mp_size):
|
||||
fused_qkvw_name_list = ['qkv_proj', 'query_key_value', 'attn.Wqkv', 'self_attn.W_pack', 'c_attn']
|
||||
|
||||
if mp_size == 1:
|
||||
return False
|
||||
for fused_name in fused_qkvw_name_list:
|
||||
if fused_name in name:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def prepare_tp_fused_qkvw(module, src, mp_size, gpu_index):
|
||||
|
||||
module_str = str(module).strip()
|
||||
if src is None:
|
||||
return
|
||||
fused_type_dict = {
|
||||
'CodeGenBlock': 'codegentype',
|
||||
'BloomBlock': 'bloomtype',
|
||||
'GLMBlock': 'glmtype',
|
||||
"MPTBlock": 'glmtype',
|
||||
"MptBlock": 'glmtype',
|
||||
"BaichuanLayer": 'glmtype',
|
||||
"QWenBlock": 'qwentype',
|
||||
"FalconDecoderLayer": 'bloomtype',
|
||||
"GPTBigCodeBlock": 'bigcodetype',
|
||||
"DecoderLayer": 'glmtype',
|
||||
"Phi3DecoderLayer": "phi3type"
|
||||
}
|
||||
|
||||
def _codegen_type_transpose(input, mp_size, codegen_mp_num=4):
|
||||
# codegen_mp_num defined in https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py
|
||||
assert get_num_kv_heads() % (
|
||||
mp_size * codegen_mp_num) == 0, "codgen autoTP requires num_kv_heads % (mp_size*codegen_mp_num) == 0"
|
||||
#input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
|
||||
|
||||
shape = input.shape
|
||||
dst_shape = get_shard_size(shape[0], mp_size)
|
||||
num_mp_blocks = input.reshape(codegen_mp_num, shape[0] // codegen_mp_num, shape[1])
|
||||
|
||||
#num_mp_blocks : [codegen_mp_num, 3*hidden_dim/codegen_mp_num, :]
|
||||
src_split = list(torch.split(num_mp_blocks, num_mp_blocks.shape[1] // 3, dim=1))
|
||||
src_split = [x.reshape(codegen_mp_num * mp_size, -1, shape[1]) for x in src_split]
|
||||
|
||||
split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size(shape[0] // 3, mp_size), 0, 1)
|
||||
tp_fuseqkv_weight = torch.cat(split_fusedqkv, dim=0).reshape(shape[0], -1)
|
||||
|
||||
return tp_fuseqkv_weight[gpu_index * dst_shape:(gpu_index + 1) * dst_shape]
|
||||
|
||||
def _glm_type_transpose(input, mp_size):
|
||||
#input : [3*hidden_dim, hidden_dim](weight) or [3*hidden_dim](bias)
|
||||
|
||||
# For chatglm2 & chatglm3(kv_heads=2), need to special handle.
|
||||
if get_num_kv_heads() == 2:
|
||||
shape = input.shape
|
||||
hidden_dim = get_n_embd()
|
||||
kv_dim = (shape[0] - hidden_dim) // get_num_kv_heads()
|
||||
q = input[:hidden_dim]
|
||||
k = input[hidden_dim:hidden_dim + kv_dim]
|
||||
v = input[hidden_dim + kv_dim:]
|
||||
q_split = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
|
||||
k_split = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
|
||||
v_split = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
|
||||
return torch.cat((q_split[gpu_index], k_split[gpu_index], v_split[gpu_index]), dim=0)
|
||||
else:
|
||||
shape = input.shape
|
||||
src_split = torch.split(input, shape[0] // 3, dim=0)
|
||||
|
||||
split_fusedqkv = split_by_qkvlist_and_refuse(src_split, get_shard_size_list(shape[0] // 3, mp_size))
|
||||
return split_fusedqkv[gpu_index]
|
||||
|
||||
def _bloom_type_transpose(input, mp_size):
|
||||
shape = input.shape
|
||||
|
||||
split_fusedqkv = input.split(get_shard_size_list(shape[0], mp_size), dim=0)
|
||||
return split_fusedqkv[gpu_index]
|
||||
|
||||
def _qwen_type_transpose(input, mp_size, module):
|
||||
if not hasattr(module, "_ds_fusedqkv_entered"):
|
||||
# Adjust splitting absolute value variables
|
||||
setattr(module, "_ds_fusedqkv_entered", True)
|
||||
module.attn.split_size = get_shard_size(module.attn.split_size, mp_size)
|
||||
return _glm_type_transpose(input, mp_size)
|
||||
|
||||
def _bigcode_type_transpose(input, mp_size):
|
||||
n_embd = get_n_embd()
|
||||
q = input[:n_embd]
|
||||
kv = input[n_embd:]
|
||||
shape = q.shape
|
||||
split_q = q.split(get_shard_size_list(shape[0], mp_size), dim=0)
|
||||
return torch.cat((split_q[gpu_index], kv), dim=0)
|
||||
|
||||
def _phi3_type_transpose(input, mp_size):
|
||||
num_kv_heads = get_num_kv_heads()
|
||||
num_heads = get_num_attention_heads()
|
||||
hidden_size = input.shape[1]
|
||||
head_dim = hidden_size // num_heads
|
||||
q_pos = input.shape[0] - 2 * num_kv_heads * head_dim
|
||||
q = input[:q_pos]
|
||||
k = input[q_pos:q_pos + num_kv_heads * head_dim]
|
||||
v = input[q_pos + num_kv_heads * head_dim:]
|
||||
split_q = q.split(get_shard_size_list(q.shape[0], mp_size), dim=0)
|
||||
split_k = k.split(get_shard_size_list(k.shape[0], mp_size), dim=0)
|
||||
split_v = v.split(get_shard_size_list(v.shape[0], mp_size), dim=0)
|
||||
return torch.cat((split_q[gpu_index], split_k[gpu_index], split_v[gpu_index]), dim=0)
|
||||
|
||||
def _transpose_fused_qkvw(src, mp_size, fused_qkv_type=None, module=None):
|
||||
|
||||
# suppose num_heads=n, q(n)_w means the n-th q head linear weight, the weight format are as following
|
||||
# bloomtype: [q(1)_w,k(1)_w,v(1)_w,q(2)_w,k(2)_w,v(2)_w,...,q(n)_w,k(n)_w,v(n)_w]
|
||||
# glmtype: [q(1)_w, q(2)_w,...,q(n)_w,k(1)_w,k(2)_w,...,k(n)_w,v(1)_w,v(2)_w,...,v(n)_w]
|
||||
# codegentype: [q(1)_w,q(2)_w,...,q(n/t)_w,k(1)_w,k(2)_w,...,k(n/t)_w,v(1)_2,v(2)_w,...v(n/t)_w,q(n/t+1)_w,...], where t is a const defined in model file.
|
||||
|
||||
if fused_qkv_type == 'bloomtype':
|
||||
return _bloom_type_transpose(src, mp_size)
|
||||
elif fused_qkv_type == 'codegentype':
|
||||
return _codegen_type_transpose(src, mp_size)
|
||||
elif fused_qkv_type == 'glmtype':
|
||||
return _glm_type_transpose(src, mp_size)
|
||||
elif fused_qkv_type == 'qwentype':
|
||||
return _qwen_type_transpose(src, mp_size, module)
|
||||
elif fused_qkv_type == 'bigcodetype':
|
||||
return _bigcode_type_transpose(src, mp_size)
|
||||
elif fused_qkv_type == 'phi3type':
|
||||
return _phi3_type_transpose(src, mp_size)
|
||||
|
||||
raise ValueError("unknown fused_qkv_type")
|
||||
|
||||
module_name_matches = [k for k in fused_type_dict.keys() if k in module_str]
|
||||
if module_name_matches:
|
||||
# There can be overlap with matches (e.g., "DecoderLayer" and "FalconDecoderLayer").
|
||||
# We take the longest matching module_name
|
||||
module_name = max(module_name_matches, key=len)
|
||||
fused_type = fused_type_dict[module_name]
|
||||
return _transpose_fused_qkvw(src, mp_size, fused_type, module)
|
||||
warning_once("Unrecognized fusedkqv weight type, default to using bloom type,"
|
||||
"please check in prepare_tp_fused_qkvw() to avoid potential calculation errors")
|
||||
return _bloom_type_transpose(src, mp_size)
|
||||
|
||||
|
||||
# For share qk type:
|
||||
# q = [q1,...,q_{n/4}, q_{n/2+1},...,q_{3n/4}, k1,...,k_{n/4}, k_{n/2+1},...,k_{3n/4}]
|
||||
# k = [q_{n/4+1},...,q_{n/2}, q_{3n/4+1},...,qn, k_{n/4+1},...,k_{n/2}, k{3n/4+1},...,kn]
|
||||
# Avoid modifying the modeling code. We adjust the value and oproj weight to fit this qk type.
|
||||
def shard_value_with_share_qk(
|
||||
weight,
|
||||
bias,
|
||||
rank,
|
||||
world_size,
|
||||
shard_value=True # True -> shard_value; False -> shard_oproj
|
||||
):
|
||||
if shard_value:
|
||||
total_size = weight.shape[0]
|
||||
weight_cat_dim = 0
|
||||
else:
|
||||
total_size = weight.shape[1]
|
||||
weight_cat_dim = 1
|
||||
num_heads = get_num_kv_heads()
|
||||
head_dim = total_size // num_heads
|
||||
assert (num_heads % world_size == 0)
|
||||
if world_size > num_heads // 2:
|
||||
RuntimeError(f"world_size {world_size} is larger than half of num_heads {num_heads}")
|
||||
head_per_rank = num_heads // world_size
|
||||
q_head_start = rank * head_per_rank
|
||||
# mapping q_head to v_head
|
||||
v_head_ids = []
|
||||
i = 0
|
||||
# mapping neighbor q_head to v_head
|
||||
while i < head_per_rank:
|
||||
v_head_ids.append(q_head_start // 2)
|
||||
q_head_start += 2
|
||||
i = i + 2
|
||||
|
||||
# mapping neighbor k_head to v_head
|
||||
v_head_ids.extend([i + num_heads // 2 for i in v_head_ids])
|
||||
sharded_weight = []
|
||||
sharded_bias = []
|
||||
for head_id in v_head_ids:
|
||||
if shard_value:
|
||||
sharded_weight.append(weight[head_id * head_dim:(head_id + 1) * head_dim])
|
||||
if bias is not None:
|
||||
sharded_bias.append(bias.data[head_id * head_dim:(head_id + 1) * head_dim])
|
||||
else:
|
||||
sharded_weight.append(weight[:, head_id * head_dim:(head_id + 1) * head_dim])
|
||||
sharded_weight = torch.cat(sharded_weight, dim=weight_cat_dim)
|
||||
if bias is not None:
|
||||
if shard_value:
|
||||
sharded_bias = torch.cat(sharded_bias, dim=0)
|
||||
else:
|
||||
bias = bias / float(world_size)
|
||||
return torch.nn.Parameter(sharded_weight), torch.nn.Parameter(sharded_bias)
|
||||
else:
|
||||
return torch.nn.Parameter(sharded_weight), None
|
||||
|
||||
|
||||
# For phi3 with chunk mlp, adjust the weight order.
|
||||
def shard_chunk_mlp(
|
||||
weight,
|
||||
bias,
|
||||
rank,
|
||||
world_size,
|
||||
):
|
||||
weight_gate, weight_states = weight.chunk(2, dim=0)
|
||||
total_size = weight_gate.shape[0]
|
||||
split_weight_gate = weight_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0)
|
||||
split_weight_states = weight_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0)
|
||||
shard_weight = torch.cat((split_weight_gate[rank], split_weight_states[rank]), dim=0)
|
||||
if bias is not None:
|
||||
bias_gate, bias_states = bias.chunk(2, dim=0)
|
||||
split_bias_gate = bias_gate.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0)
|
||||
split_bias_states = bias_states.split(get_shard_size_list(total_size, world_size, "mlp"), dim=0)
|
||||
return shard_weight, torch.cat((split_bias_gate[rank], split_bias_states[rank]), dim=0)
|
||||
|
||||
return shard_weight, None
|
||||
Executable
+112
@@ -0,0 +1,112 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import copy
|
||||
import torch
|
||||
from deepspeed.ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig
|
||||
|
||||
|
||||
def module_inject(layer_obj, model, config, micro_batch_size, max_seq_length, seed, preln, fp16=True):
|
||||
for name, child in model.named_children():
|
||||
if isinstance(child, layer_obj):
|
||||
print('REPLACING BertLayer')
|
||||
|
||||
cuda_config = DeepSpeedTransformerConfig(batch_size=micro_batch_size,
|
||||
max_seq_length=max_seq_length,
|
||||
hidden_size=config.hidden_size,
|
||||
heads=config.num_attention_heads,
|
||||
attn_dropout_ratio=config.attention_probs_dropout_prob,
|
||||
hidden_dropout_ratio=config.hidden_dropout_prob,
|
||||
num_hidden_layers=config.num_hidden_layers,
|
||||
initializer_range=config.initializer_range,
|
||||
seed=seed,
|
||||
fp16=fp16,
|
||||
pre_layer_norm=preln)
|
||||
|
||||
new_module = DeepSpeedTransformerLayer(cuda_config)
|
||||
|
||||
# copy relevant state from child -> new module
|
||||
qw = child.attention.self.query.weight
|
||||
qb = child.attention.self.query.bias
|
||||
kw = child.attention.self.key.weight
|
||||
kb = child.attention.self.key.bias
|
||||
vw = child.attention.self.value.weight
|
||||
vb = child.attention.self.value.bias
|
||||
|
||||
qkvw = torch.cat((qw, kw, vw), 0)
|
||||
qkvb = torch.cat((qb, kb, vb), 0)
|
||||
|
||||
new_module.attn_qkvw.data = qkvw
|
||||
new_module.attn_qkvb.data = qkvb
|
||||
new_module.attn_ow.data = child.attention.output.dense.weight
|
||||
new_module.attn_ob.data = child.attention.output.dense.bias
|
||||
if preln:
|
||||
attention_layerNorm = child.PostAttentionLayerNorm
|
||||
else:
|
||||
attention_layerNorm = child.attention.output.LayerNorm
|
||||
new_module.attn_nw.data = attention_layerNorm.weight
|
||||
new_module.attn_nb.data = attention_layerNorm.bias
|
||||
if preln:
|
||||
intermediate_FF = child.intermediate.dense_act
|
||||
else:
|
||||
intermediate_FF = child.intermediate.dense
|
||||
new_module.inter_w.data = intermediate_FF.weight
|
||||
new_module.inter_b.data = intermediate_FF.bias
|
||||
new_module.output_w.data = child.output.dense.weight
|
||||
new_module.output_b.data = child.output.dense.bias
|
||||
if preln:
|
||||
transformer_LayerNorm = child.PreAttentionLayerNorm
|
||||
else:
|
||||
transformer_LayerNorm = child.output.LayerNorm
|
||||
new_module.norm_w.data = transformer_LayerNorm.weight
|
||||
new_module.norm_b.data = transformer_LayerNorm.bias
|
||||
|
||||
setattr(model, name, copy.deepcopy(new_module))
|
||||
|
||||
else:
|
||||
module_inject(layer_obj, child, config, micro_batch_size, max_seq_length, seed, preln, fp16)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def test_hi():
|
||||
from turing.nvidia_modelingpreln import BertConfig as BertConfigPreLN
|
||||
from turing.nvidia_modelingpreln import BertForQuestionAnswering as BertForQuestionAnsweringPreLN
|
||||
from turing.nvidia_modelingpreln import BertLayer
|
||||
bert_model_config = {
|
||||
"vocab_size_or_config_json_file": 119547,
|
||||
"hidden_size": 1024,
|
||||
"num_hidden_layers": 1,
|
||||
"num_attention_heads": 16,
|
||||
"intermediate_size": 4096,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.1,
|
||||
"attention_probs_dropout_prob": 0.1,
|
||||
"hidden_dropout_prob": 0.1,
|
||||
"attention_probs_dropout_prob": 0.1,
|
||||
"max_position_embeddings": 512,
|
||||
"type_vocab_size": 2,
|
||||
"initializer_range": 0.02
|
||||
}
|
||||
bert_config = BertConfigPreLN(**bert_model_config)
|
||||
base_model = BertForQuestionAnsweringPreLN(bert_config, args=None)
|
||||
|
||||
#base_model = LinearStack()
|
||||
|
||||
test_model = copy.deepcopy(base_model)
|
||||
test_model = module_inject(BertLayer, test_model, bert_config, 4, 384, 1234)
|
||||
|
||||
print('BASE', base_model)
|
||||
print('TEST', test_model)
|
||||
|
||||
#base_model.eval()
|
||||
#test_model.eval()
|
||||
|
||||
#test_input = torch.rand(1, base_model.input_dim)
|
||||
|
||||
#base_output = base_model(test_input)
|
||||
#test_output = test_model(test_input)
|
||||
#
|
||||
#assert torch.allclose(base_output, test_output, atol=3e-8)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,278 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from torch import nn
|
||||
from deepspeed.model_implementations.transformers.ds_bloom import DeepSpeedBloomInference
|
||||
from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference
|
||||
from deepspeed.model_implementations.transformers.ds_bert import DeepSpeedBERTInference
|
||||
from deepspeed.model_implementations.transformers.ds_megatron_gpt import DeepSpeedMegatronGPTInference
|
||||
from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference
|
||||
from deepspeed.model_implementations.transformers.ds_llama2 import DeepSpeedLlama2Inference
|
||||
|
||||
import deepspeed.ops.transformer as transformer_inference
|
||||
from .layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding, RMSNormalize
|
||||
import torch
|
||||
import gc
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
import re
|
||||
from .utils import transpose
|
||||
|
||||
|
||||
def load_model_with_checkpoint(r_module,
|
||||
sd,
|
||||
mp_replace,
|
||||
ckpt_type,
|
||||
ckpt_mp_size,
|
||||
weight_quantizer=None,
|
||||
rank=0,
|
||||
container=None):
|
||||
error_msgs = []
|
||||
|
||||
def prefix_check():
|
||||
# if keys start with 'model.' or 'transformer.', don't skip level 0 prefix
|
||||
for key in sd[0].keys():
|
||||
# OPT models
|
||||
if re.match("^model[.]", key):
|
||||
return False
|
||||
# BLOOM models
|
||||
if re.match("^transformer[.]", key):
|
||||
return False
|
||||
return True
|
||||
|
||||
skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix
|
||||
|
||||
def load(module, prefix):
|
||||
args = (sd[0], prefix, {}, True, [], [], error_msgs)
|
||||
|
||||
if hasattr(module, 'weight'):
|
||||
module.weight = mp_replace.copy(module.weight.data, sd[0][prefix + 'weight'])
|
||||
if prefix + 'bias' in sd[0].keys():
|
||||
if module.bias.data.is_meta:
|
||||
# meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
|
||||
module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"),
|
||||
requires_grad=module.bias.data.requires_grad)
|
||||
module.bias = mp_replace.copy(module.bias.data, sd[0][prefix + 'bias'])
|
||||
args = None
|
||||
gc.collect()
|
||||
|
||||
def load_transformer_layer(module, prefix):
|
||||
if ckpt_type == "tp":
|
||||
|
||||
def load_parameters(module, prefix):
|
||||
for n, p in module.named_parameters():
|
||||
if prefix + n in sd[0] and len(n.split('.')) == 1:
|
||||
if type(sd[0][prefix + n]) is list:
|
||||
tmp_data, scale = sd[0][prefix + n]
|
||||
tmp_data = tmp_data
|
||||
scale = scale.to(get_accelerator().current_device_name())
|
||||
# set the quantizer number of groups using the checkpoint scale shape
|
||||
weight_quantizer.num_groups = scale.shape[0]
|
||||
else:
|
||||
tmp_data = sd[0][prefix + n].to(get_accelerator().current_device_name())
|
||||
scale = None
|
||||
src_shape = tmp_data.shape
|
||||
dst_shape = p.shape
|
||||
inner_dim = 1 if tmp_data.dtype == torch.int8 else 0
|
||||
outer_dim = 0 if tmp_data.dtype == torch.int8 else 1
|
||||
if (len(src_shape) == 2 and len(dst_shape) == 2):
|
||||
if (src_shape[inner_dim] == dst_shape[0] and src_shape[outer_dim] == dst_shape[1]):
|
||||
if tmp_data.dtype != torch.int8:
|
||||
p = weight_quantizer.quantize(
|
||||
transpose(tmp_data) if weight_quantizer.q_int8 else tmp_data)
|
||||
else:
|
||||
p = torch.nn.parameter.Parameter(tmp_data, requires_grad=False)
|
||||
p.scale = scale
|
||||
setattr(module, n, p)
|
||||
else:
|
||||
dim = inner_dim if src_shape[inner_dim] != dst_shape[0] else outer_dim
|
||||
dim1 = 0 if src_shape[inner_dim] != dst_shape[0] else 1
|
||||
if src_shape[dim] > dst_shape[dim1]:
|
||||
weight_partition = torch.split(tmp_data, dst_shape[dim1], dim=dim)[rank].to(
|
||||
get_accelerator().current_device_name())
|
||||
assert tmp_data.dtype != torch.int8 or scale.numel() > weight_quantizer.num_groups * (rank+1), \
|
||||
'''ERROR: We require the quantization scales for larger TP-size when loading INT8 checkpoint!\
|
||||
Please use the FP16 checkpoint to generate INT8 checkpoint with the sharding parameters!'''
|
||||
scale = scale.view(-1)[weight_quantizer.num_groups * (rank + 1):].reshape(
|
||||
weight_quantizer.num_groups, -1).contiguous()
|
||||
else:
|
||||
assert tmp_data.dtype != torch.int8, \
|
||||
'''Merging of the checkpoints are not supported when using INT8 checkpoint! \
|
||||
Please use a as many GPUs as TP-size for the checkpoint'''
|
||||
all_data = [
|
||||
sd[j][prefix + n] if type(sd[j][prefix + n]) is list else sd[j][prefix + n].to(
|
||||
get_accelerator().current_device_name()) for j in range(len(sd))
|
||||
]
|
||||
# Check if the weight tensor is for the QKV parameter
|
||||
if src_shape[1] == (3 * src_shape[0]) // ckpt_mp_size:
|
||||
qkv_size = src_shape[outer_dim] // 3
|
||||
src_split = [
|
||||
torch.split(src[0].data, qkv_size, dim=outer_dim) for src in all_data
|
||||
]
|
||||
|
||||
weight_partition = torch.cat([
|
||||
torch.cat([qkv_s[i] for qkv_s in src_split], axis=outer_dim)
|
||||
for i in range(len(src_split[0]))
|
||||
],
|
||||
dim=dim)
|
||||
else:
|
||||
weight_partition = torch.cat([
|
||||
ad[0].to(get_accelerator().current_device_name())
|
||||
if type(ad) is list else ad for ad in all_data
|
||||
],
|
||||
dim=dim)
|
||||
if tmp_data.dtype == torch.int8:
|
||||
scale = torch.cat(
|
||||
[ad[1].to(get_accelerator().current_device_name()) for ad in all_data],
|
||||
dim=dim)
|
||||
|
||||
if tmp_data.dtype != torch.int8:
|
||||
weight_partition = weight_quantizer.quantize(
|
||||
transpose(weight_partition), \
|
||||
parallel_dim=(0 if dim == 1 else 1)) if weight_quantizer.q_int8 else \
|
||||
weight_quantizer.quantize(weight_partition)
|
||||
else:
|
||||
weight_partition = torch.nn.parameter.Parameter(weight_partition,
|
||||
requires_grad=False)
|
||||
weight_partition.scale = scale
|
||||
setattr(module, n, weight_partition)
|
||||
else:
|
||||
if src_shape[0] == dst_shape[0]:
|
||||
p.data.copy_(tmp_data)
|
||||
else:
|
||||
if src_shape[0] > dst_shape[0]:
|
||||
bias_split = torch.split(tmp_data, dst_shape[-1])[rank].to(
|
||||
get_accelerator().current_device_name()).contiguous()
|
||||
p.data.copy_(bias_split)
|
||||
else:
|
||||
# Check if the weight tensor is for the QKV parameter
|
||||
if src_shape[0] == (3 * r_module.config.hidden_size) // ckpt_mp_size:
|
||||
qkv_size = src_shape[0] // 3
|
||||
src_split = [
|
||||
torch.split(sd[j][prefix + n], qkv_size, dim=0) for j in range(len(sd))
|
||||
]
|
||||
|
||||
p.data.copy_(
|
||||
torch.cat([
|
||||
torch.cat([qkv_s[i] for qkv_s in src_split], axis=0)
|
||||
for i in range(len(src_split[0]))
|
||||
],
|
||||
dim=0).to(get_accelerator().current_device_name()).contiguous())
|
||||
else:
|
||||
p.data.copy_(
|
||||
torch.cat([sd[j][prefix + n] for j in range(len(sd))],
|
||||
dim=0).to(get_accelerator().current_device_name()).contiguous())
|
||||
|
||||
load_parameters(module, prefix)
|
||||
for n, child in module.named_children():
|
||||
load_parameters(child, prefix + n + '.')
|
||||
else:
|
||||
container.load_params(module, sd[0], weight_quantizer, mp_replace, prefix)
|
||||
|
||||
try:
|
||||
import transformers
|
||||
OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding
|
||||
if hasattr(transformers.models, "llama"):
|
||||
LlamaRMSNorm = transformers.models.llama.modeling_llama.LlamaRMSNorm
|
||||
else:
|
||||
LlamaRMSNorm = None
|
||||
except Exception:
|
||||
OPTLearnedPositionalEmbedding = None
|
||||
try:
|
||||
from fairscale.nn.model_parallel.layers import (
|
||||
ColumnParallelLinear,
|
||||
ParallelEmbedding,
|
||||
RowParallelLinear,
|
||||
)
|
||||
except Exception:
|
||||
ColumnParallelLinear = None
|
||||
ParallelEmbedding = None
|
||||
RowParallelLinear = None
|
||||
try:
|
||||
from llama.model import RMSNorm
|
||||
except Exception:
|
||||
RMSNorm = None
|
||||
layer_policies = {
|
||||
nn.Linear: load,
|
||||
nn.Embedding: load,
|
||||
nn.LayerNorm: load,
|
||||
EmbeddingLayer: load,
|
||||
LinearLayer: load,
|
||||
Normalize: load,
|
||||
transformer_inference.DeepSpeedTransformerInference: load_transformer_layer,
|
||||
DeepSpeedBloomInference: load_transformer_layer,
|
||||
DeepSpeedGPTInference: load_transformer_layer,
|
||||
DeepSpeedBERTInference: load_transformer_layer,
|
||||
DeepSpeedMegatronGPTInference: load_transformer_layer,
|
||||
DeepSpeedOPTInference: load_transformer_layer,
|
||||
DeepSpeedLlama2Inference: load_transformer_layer,
|
||||
OPTLearnedPositionalEmbedding: load,
|
||||
OPTEmbedding: load,
|
||||
LlamaRMSNorm: load,
|
||||
RMSNormalize: load,
|
||||
ColumnParallelLinear: load,
|
||||
ParallelEmbedding: load,
|
||||
RowParallelLinear: load,
|
||||
RMSNorm: load
|
||||
}
|
||||
|
||||
all_ds_ids = {}
|
||||
|
||||
def load_module_recursive(module, prefix='', level=0):
|
||||
for name, child in module.named_children():
|
||||
if child.__class__ in layer_policies:
|
||||
checking_key = prefix + name + '.'
|
||||
if not any(checking_key in item for item in sd[0].keys()):
|
||||
if hasattr(child, 'weight') and \
|
||||
(hasattr(child.weight, 'ds_id') and \
|
||||
child.weight.ds_id in all_ds_ids):
|
||||
prefix1 = all_ds_ids[child.weight.ds_id]
|
||||
if child.__class__ is nn.Linear:
|
||||
child = LinearLayer.from_weights(weight=all_ds_ids[child.weight.ds_id])
|
||||
setattr(module, name, child)
|
||||
continue
|
||||
child_params = list(child.parameters())
|
||||
if len(child_params) > 0 and (child_params[0].numel() == 0 or child_params[0].is_meta):
|
||||
if child.weight.is_meta:
|
||||
ds_shape = child.weight.shape
|
||||
else:
|
||||
ds_shape = child.weight.ds_shape
|
||||
if child.__class__ is nn.LayerNorm:
|
||||
child = Normalize(dim=ds_shape[-1], dtype=child.weight.dtype, eps=child.eps)
|
||||
setattr(module, name, child)
|
||||
elif child.__class__ in [nn.Linear, ColumnParallelLinear, RowParallelLinear]:
|
||||
child = LinearLayer.from_weights(weight_shape=child.weight.shape,
|
||||
dtype=child.weight.dtype,
|
||||
bias=child.bias)
|
||||
setattr(module, name, child)
|
||||
elif child.__class__ is OPTLearnedPositionalEmbedding:
|
||||
child = OPTEmbedding(weight_shape=ds_shape)
|
||||
setattr(module, name, child)
|
||||
elif child.__class__ in [LlamaRMSNorm, RMSNorm]:
|
||||
child = RMSNormalize(dim=ds_shape[-1],
|
||||
dtype=child.weight.dtype,
|
||||
eps=child.eps if hasattr(child, 'eps') else child.variance_epsilon)
|
||||
setattr(module, name, child)
|
||||
else:
|
||||
ds_id = None
|
||||
if hasattr(child.weight, 'ds_id'):
|
||||
ds_id = child.weight.ds_id
|
||||
child = EmbeddingLayer(weight_shape=ds_shape, dtype=child.weight.dtype)
|
||||
if ds_id is not None:
|
||||
all_ds_ids[ds_id] = child.weight
|
||||
setattr(module, name, child)
|
||||
layer_policies[child.__class__](child, prefix + name + '.')
|
||||
else:
|
||||
load_module_recursive(
|
||||
child,
|
||||
prefix if (level == 0 and ckpt_type == 'pp') and skip_level_0_prefix else \
|
||||
prefix + name + '.',
|
||||
level + 1)
|
||||
|
||||
load_module_recursive(r_module)
|
||||
|
||||
for sd_ in sd:
|
||||
del sd_
|
||||
sd = None
|
||||
gc.collect()
|
||||
Executable
+72
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def quantize_transformer_layer(orig_layer_impl, model, megatron=False, preln=False):
|
||||
""" Quantize bert-style transformer layers with DeepSpeed's transformer layer
|
||||
Arguments:
|
||||
orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for,
|
||||
e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer
|
||||
model (torch.nn.Module): user's nn.module representing their model
|
||||
|
||||
megatron (bool): megatron model-parallel implementation (this is supported for inference only)
|
||||
preln (bool): does the original layer implementation do pre or post layer norm?
|
||||
|
||||
Note: For Bert kind of models, we inject based on the DeepSpeed-Example models, if not setting huggingface flag.
|
||||
|
||||
Returns:
|
||||
Updated nn.module with quantized transformer layers
|
||||
"""
|
||||
|
||||
def quantize_weight(weight):
|
||||
return weight.to(torch.int8)
|
||||
|
||||
def megatron_layer_quantize(layer):
|
||||
layer.attention.query_key_value.weight.data = quantize_weight(layer.attention.query_key_value.weight.data)
|
||||
layer.attention.dense.weight.data = quantize_weight(layer.attention.dense.weight.data)
|
||||
layer.mlp.dense_h_to_4h.weight.data = quantize_weight(layer.mlp.dense_h_to_4h.weight.data)
|
||||
layer.mlp.dense_4h_to_h.weight.data = quantize_weight(layer.mlp.dense_4h_to_h.weight.data)
|
||||
|
||||
def bert_layer_quantize(layer):
|
||||
layer.attention.self.query.weight.data = quantize_weight(layer.attention.self.query.weight.data)
|
||||
layer.attention.self.key.weight.data = quantize_weight(layer.attention.self.key.weight.data)
|
||||
layer.attention.self.value.weight.data = quantize_weight(layer.attention.self.value.weight.data)
|
||||
layer.attention.output.dense.weight.data = quantize_weight(layer.attention.output.dense.weight.data)
|
||||
if preln:
|
||||
layer.intermediate.dense_act.weight.data = quantize_weight(layer.intermediate.dense_act.weight.data)
|
||||
else:
|
||||
layer.intermediate.dense.weight.data = quantize_weight(layer.intermediate.dense.weight.data)
|
||||
layer.output.dense.weight.data = quantize_weight(layer.output.dense.weight.data)
|
||||
|
||||
def quantize_fn(child):
|
||||
if megatron:
|
||||
# Quantize megatron GPT2 / GPT3 trained model
|
||||
megatron_layer_quantize(child)
|
||||
else:
|
||||
# Quantize either DeepSpeed or HuggingFace trained model
|
||||
bert_layer_quantize(child)
|
||||
|
||||
return child
|
||||
|
||||
return quantize_module(model=model, orig_class=orig_layer_impl, quantize_fn=quantize_fn)
|
||||
|
||||
|
||||
def quantize_module(model, orig_class, quantize_fn):
|
||||
policy = {orig_class: quantize_fn}
|
||||
return _quantize_module(model, policy)
|
||||
|
||||
|
||||
def _quantize_module(model, policies):
|
||||
for name, child in model.named_children():
|
||||
if child.__class__ in policies:
|
||||
orig = repr(child)
|
||||
setattr(model, name, policies[child.__class__](child))
|
||||
new = getattr(model, name)
|
||||
else:
|
||||
_quantize_module(child, policies)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,215 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from deepspeed.utils.types import ActivationFuncType, NormType
|
||||
import torch
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from .utils import transpose
|
||||
|
||||
transformer_param_names = (
|
||||
'attn_qkvw', \
|
||||
'attn_qkvb', \
|
||||
'attn_ow' , \
|
||||
'attn_ob', \
|
||||
'inter_w', \
|
||||
'inter_b', \
|
||||
'output_w', \
|
||||
'output_b', \
|
||||
'attn_nw', \
|
||||
'attn_nb', \
|
||||
'norm_w', \
|
||||
'norm_b')
|
||||
|
||||
|
||||
class DSPolicy(ABC):
|
||||
_orig_layer_class = None
|
||||
|
||||
def __init__(self):
|
||||
self.cuda_graph_supported = False
|
||||
|
||||
@abstractmethod
|
||||
def attention(self):
|
||||
"""
|
||||
Returns attention qkv and dense parameters
|
||||
weight: (3*hidden, hidden) and (hidden, hidden)
|
||||
bias: (3*hidden) and (hidden)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class TransformerPolicy(DSPolicy):
|
||||
# a static class variable containing the HuggingFace model configuration.
|
||||
# see e.g., transformers.models.opt.configuration_opt.OPTConfig
|
||||
hf_model_config = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inference=True,
|
||||
linear_layer=True,
|
||||
scale_attention=True,
|
||||
megatron_v2=False,
|
||||
use_mup=False,
|
||||
# the type of activation function used in MLP
|
||||
mlp_act_func_type=ActivationFuncType.GELU,
|
||||
# applies layer norm before attention if `pre_attn_norm` is set to True
|
||||
pre_attn_norm=True,
|
||||
# this flag shows whether or not using prefix in loading the checkpoint
|
||||
use_load_prefix=False,
|
||||
# whether or not the qkv is stored in the split-format
|
||||
split_qkv=True,
|
||||
# Type of normalization to perform
|
||||
norm_type=NormType.LayerNorm):
|
||||
super().__init__()
|
||||
self.cuda_graph_supported = False
|
||||
self.inference = inference
|
||||
self.linear_layer = linear_layer
|
||||
self.scale_attention = scale_attention
|
||||
self.is_megatron_v2 = megatron_v2
|
||||
self.use_mup = use_mup
|
||||
self.mlp_act_func_type = mlp_act_func_type
|
||||
self.pre_attn_norm = pre_attn_norm
|
||||
self.use_load_prefix = use_load_prefix
|
||||
self.split_qkv = split_qkv
|
||||
self.norm_type = norm_type
|
||||
|
||||
@abstractmethod
|
||||
def attention(self):
|
||||
"""
|
||||
Returns attention qkv and dense parameters
|
||||
weight: (3*hidden, hidden) and (hidden, hidden)
|
||||
bias: (3*hidden) and (hidden)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_hidden_heads(self):
|
||||
"""
|
||||
return hidden_size and number of heads
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def mlp(self):
|
||||
"""
|
||||
Returns mlp intermediate and output
|
||||
weight: (intermediate, hidden) and (hidden, intermediate)
|
||||
bias: (intermediate) and (hidden)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def layernorm(self):
|
||||
"""
|
||||
Returns LayerNorms used in transformer layer
|
||||
Post-Attention and pre/post layer norm
|
||||
gamma and beta with shape: (hidden)
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
# TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point
|
||||
def _transpose(x, heads=1, mp_replace=None):
|
||||
heads = heads // mp_replace.mp_size # type: ignore
|
||||
outer_dim = -1
|
||||
attention_head_size = x.shape[outer_dim] // heads
|
||||
new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size)
|
||||
x_1 = x.view(*new_x_shape)
|
||||
(q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=-1)
|
||||
if len(q.shape) > 2:
|
||||
new_shape = (q.shape[0], ) + (-1, )
|
||||
return torch.cat((q.reshape(new_shape), k.reshape(new_shape), v.reshape(new_shape)),
|
||||
dim=outer_dim).reshape(x.shape)
|
||||
else:
|
||||
return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)
|
||||
|
||||
|
||||
# This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor.
|
||||
# Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist!
|
||||
def maybe_copy(module,
|
||||
sd,
|
||||
weight_quantizer,
|
||||
mp_replace,
|
||||
dst_name,
|
||||
src_name,
|
||||
qkv=False,
|
||||
megatron_v2=False,
|
||||
split_qkv=False,
|
||||
heads=1):
|
||||
if src_name in sd:
|
||||
dst = getattr(module, dst_name)
|
||||
tmp = sd[src_name]
|
||||
if len(dst.shape) == 1:
|
||||
if split_qkv:
|
||||
dst = mp_replace.strided_copy(dst, tmp, num_splits=3)
|
||||
else:
|
||||
dst = mp_replace.copy(dst, tmp)
|
||||
if qkv and megatron_v2:
|
||||
dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous())
|
||||
else:
|
||||
if split_qkv:
|
||||
dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \
|
||||
(transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8)
|
||||
else:
|
||||
if qkv and megatron_v2:
|
||||
tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous()
|
||||
if weight_quantizer.q_int8:
|
||||
tmp = transpose(tmp)
|
||||
dst = mp_replace.copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \
|
||||
transpose(tmp)), int8=weight_quantizer.q_int8)
|
||||
setattr(module, dst_name, dst)
|
||||
|
||||
|
||||
# Extending the maybe_copy function for when the q, k, and v are in separate parameters!
|
||||
def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names, split_qkv=False):
|
||||
if src_names[0] in sd:
|
||||
q = sd[src_names[0]]
|
||||
k = sd[src_names[1]]
|
||||
v = sd[src_names[2]]
|
||||
qkv_data = torch.cat((q, k, v), dim=0)
|
||||
dst = getattr(module, dst_name)
|
||||
if len(dst.shape) == 1:
|
||||
if split_qkv:
|
||||
dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3)
|
||||
else:
|
||||
dst = mp_replace.copy(dst, qkv_data)
|
||||
else:
|
||||
if split_qkv:
|
||||
dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
|
||||
((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8)
|
||||
else:
|
||||
dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
|
||||
transpose(qkv_data)), int8=weight_quantizer.q_int8)
|
||||
setattr(module, dst_name, dst)
|
||||
|
||||
|
||||
# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU
|
||||
def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names):
|
||||
if src_names[0] in sd:
|
||||
reg_proj = sd[src_names[0]]
|
||||
gate_proj = sd[src_names[1]]
|
||||
|
||||
mlp1_data = torch.cat((reg_proj, gate_proj), dim=0)
|
||||
dst = getattr(module, dst_name)
|
||||
|
||||
dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
|
||||
transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8)
|
||||
setattr(module, dst_name, dst)
|
||||
|
||||
|
||||
def pack_lora_weights(p):
|
||||
return [
|
||||
p.lora_right_weight, \
|
||||
p.lora_left_weight, \
|
||||
p.lora_scaling
|
||||
]
|
||||
|
||||
|
||||
def maybe_get_lora(p):
|
||||
if hasattr(p, 'lora_right_weight'):
|
||||
lora_param = pack_lora_weights(p)
|
||||
else:
|
||||
lora_param = []
|
||||
return lora_param
|
||||
@@ -0,0 +1,734 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import os
|
||||
import torch
|
||||
import tqdm
|
||||
import deepspeed
|
||||
import deepspeed.ops.transformer as transformer_inference
|
||||
from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention
|
||||
from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
|
||||
from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from .replace_policy import replace_policies, generic_policies
|
||||
from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading
|
||||
from .layers import TensorParallelOcShardConv2d, TensorParallelIcShardConv2d
|
||||
from deepspeed.module_inject.layers import is_autotp_training_mode
|
||||
from deepspeed import comm as dist
|
||||
from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads, set_tp_grain_size
|
||||
|
||||
from .load_checkpoint import load_model_with_checkpoint
|
||||
import time
|
||||
|
||||
from .utils import policy_to_ds_container
|
||||
import gc
|
||||
|
||||
|
||||
def get_transformer_name(replaced_module):
|
||||
from .containers import supported_models
|
||||
from torch.nn import ModuleList
|
||||
transformer_name = ''
|
||||
for n, c in replaced_module.named_children():
|
||||
if c.__class__ in supported_models:
|
||||
transformer_name += n + '.'
|
||||
for name, child in c.named_children():
|
||||
if child.__class__ is ModuleList:
|
||||
transformer_name += name
|
||||
break
|
||||
break
|
||||
return transformer_name
|
||||
|
||||
|
||||
class GroupQuantizer:
|
||||
|
||||
def __init__(self, q_int8=True, group_size=1, num_bits=8, num_groups=0):
|
||||
self.group_size = group_size
|
||||
self.num_bits = num_bits
|
||||
self.q_int8 = q_int8
|
||||
|
||||
self.num_groups = num_groups
|
||||
|
||||
def quantize(self, inputs, qkv=True, count=1, parallel_dim=0):
|
||||
if not self.q_int8 or not qkv:
|
||||
inputs = torch.nn.Parameter(inputs, requires_grad=False)
|
||||
inputs.scale = torch.empty(1)
|
||||
return inputs
|
||||
q_range = 2**self.num_bits
|
||||
num_groups = self.num_groups if self.num_groups > 0 else inputs.shape[0] // self.group_size
|
||||
inputs = inputs.to(get_accelerator().current_device_name())
|
||||
input_flat = inputs.reshape(num_groups, -1).contiguous()
|
||||
input_min = torch.min(input_flat, dim=1, keepdim=True)[0].float()
|
||||
input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float()
|
||||
scale = torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range)
|
||||
input_flat = (input_flat / scale).round().clamp(-q_range // 2, q_range // 2 - 1)
|
||||
inputs_q = input_flat.reshape(inputs.shape).to(torch.int8).contiguous()
|
||||
out = torch.nn.Parameter(inputs_q, requires_grad=False)
|
||||
inputs_split = inputs.split(inputs.shape[parallel_dim] // 2, dim=parallel_dim)
|
||||
input_flat = [inputs_split[i].reshape(num_groups, -1).contiguous() for i in range(2)]
|
||||
input_min = [torch.min(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)]
|
||||
input_max = [torch.max(input_flat[i], dim=1, keepdim=True)[0].float() for i in range(2)]
|
||||
scale1 = [(torch.max(input_min[i].abs(), input_max[i].abs()) * 2.0 / (q_range)).squeeze().unsqueeze(0)
|
||||
for i in range(2)]
|
||||
|
||||
out.scale = torch.cat([scale.squeeze().unsqueeze(0), scale1[0], scale1[1]], dim=0).reshape(num_groups,
|
||||
-1).contiguous()
|
||||
return out
|
||||
|
||||
|
||||
def _module_match(module):
|
||||
for policy in generic_policies:
|
||||
policy = policy()
|
||||
if policy.match(module):
|
||||
return policy
|
||||
return None
|
||||
|
||||
|
||||
def generic_injection(module, dtype=None, enable_cuda_graph=True):
|
||||
|
||||
def replace_attn(child, policy):
|
||||
policy_attn = policy.attention(child)
|
||||
if policy_attn is None:
|
||||
return child
|
||||
if len(policy_attn) == 5:
|
||||
qkvw, attn_ow, attn_ob, hidden_size, heads = policy_attn
|
||||
qw, kw, vw = torch.empty(0), torch.empty(0), torch.empty(0)
|
||||
else:
|
||||
qw, kw, vw, attn_ow, attn_ob, hidden_size, heads = policy_attn
|
||||
qkvw = torch.empty(0)
|
||||
|
||||
config = transformer_inference.DeepSpeedInferenceConfig(
|
||||
hidden_size=hidden_size,
|
||||
heads=heads,
|
||||
dtype=dtype,
|
||||
triangular_masking=False,
|
||||
max_out_tokens=4096,
|
||||
)
|
||||
attn_module = DeepSpeedDiffusersAttention(config)
|
||||
|
||||
def transpose(data):
|
||||
data = data.contiguous()
|
||||
data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1))
|
||||
data = data.reshape(data.shape[-1], data.shape[-2])
|
||||
data.to(get_accelerator().current_device_name())
|
||||
return data
|
||||
|
||||
if len(policy_attn) == 5:
|
||||
assert qkvw is not None and qkvw.data is not None, "qkvw can't be None"
|
||||
attn_module.attn_qkvw.data = transpose(qkvw.data)
|
||||
else:
|
||||
attn_module.attn_qkvw = None
|
||||
assert qw is not None and qw.data is not None, "qw can't be None"
|
||||
attn_module.attn_qw.data = transpose(qw.data)
|
||||
assert kw is not None and kw.data is not None, "kw can't be None"
|
||||
attn_module.attn_kw.data = transpose(kw.data)
|
||||
assert vw is not None and vw.data is not None, "vw can't be None"
|
||||
attn_module.attn_vw.data = transpose(vw.data)
|
||||
|
||||
attn_module.attn_qkvb = None
|
||||
attn_module.attn_ow.data = transpose(attn_ow.data)
|
||||
attn_module.attn_ob.data.copy_(attn_ob.data.to(get_accelerator().current_device_name()))
|
||||
return attn_module
|
||||
|
||||
def replace_attn_block(child, policy):
|
||||
config = Diffusers2DTransformerConfig()
|
||||
return DeepSpeedDiffusersTransformerBlock(child, config)
|
||||
|
||||
if isinstance(module, torch.nn.Module):
|
||||
pass
|
||||
else:
|
||||
if dtype not in [torch.float16, torch.half]:
|
||||
raise ValueError("Generic injection only supported with FP16")
|
||||
|
||||
try:
|
||||
import diffusers
|
||||
if hasattr(diffusers.models.attention, 'CrossAttention'):
|
||||
cross_attention = diffusers.models.attention.CrossAttention
|
||||
else:
|
||||
cross_attention = diffusers.models.attention_processor.Attention
|
||||
attention_block = diffusers.models.attention.BasicTransformerBlock
|
||||
new_policies = {
|
||||
cross_attention: replace_attn,
|
||||
attention_block: replace_attn_block,
|
||||
}
|
||||
except ImportError:
|
||||
new_policies = {}
|
||||
|
||||
#replace_transformer_layer(None,
|
||||
# module.text_encoder,
|
||||
# training=False,
|
||||
# replace_with_kernel_inject=True,
|
||||
# triangular_masking=True,
|
||||
# max_out_tokens=8192)
|
||||
from ..model_implementations.transformers.clip_encoder import DSClipEncoder
|
||||
cg_encoder = DSClipEncoder(module.text_encoder, enable_cuda_graph=enable_cuda_graph)
|
||||
setattr(module, 'text_encoder', cg_encoder)
|
||||
for name in module.__dict__.keys():
|
||||
sub_module = getattr(module, name)
|
||||
policy = _module_match(sub_module)
|
||||
|
||||
if policy is not None:
|
||||
|
||||
def _replace_module(module, policy):
|
||||
for name, child in module.named_children():
|
||||
_replace_module(child, policy)
|
||||
if child.__class__ in new_policies:
|
||||
replaced_module = new_policies[child.__class__](child, policy)
|
||||
setattr(module, name, replaced_module)
|
||||
|
||||
_replace_module(sub_module, policy)
|
||||
new_module = policy.apply(sub_module, enable_cuda_graph=enable_cuda_graph)
|
||||
print(f"**** found and replaced {name} w. {type(new_module)}")
|
||||
setattr(module, name, new_module)
|
||||
|
||||
|
||||
container_g = None
|
||||
|
||||
|
||||
def replace_transformer_layer(orig_layer_impl, model, checkpoint_dict, config, model_config):
|
||||
""" Replace bert-style transformer layers with DeepSpeed's transformer layer
|
||||
Arguments:
|
||||
orig_layer_impl (torch.nn.Module): the original transformer layer implementation to look for,
|
||||
e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer
|
||||
model (torch.nn.Module): user's nn.module representing their model
|
||||
checkpoint_dict: Dictionary for checkpoint passed from the Inference Engine
|
||||
config: top-level DS Inference config defined in inference/config.py
|
||||
model_config: HuggingFace model config passed from the inference/engine.py
|
||||
Returns:
|
||||
Updated nn.module with replaced transformer layers
|
||||
"""
|
||||
# defining globals as internally defined functions inherit these everywhere
|
||||
quantize = (config.dtype == torch.int8)
|
||||
# todo: Refactor later. In future, let's minimize the style used above and use config.** instead
|
||||
|
||||
linear_layer_setting = None
|
||||
'''
|
||||
linear_layer_setting (tuple of modules) [Optional]: shows which two classes are used for linear layers and embedding layers
|
||||
'''
|
||||
micro_batch_size = -1
|
||||
seed = -1
|
||||
local_rank = -1
|
||||
|
||||
mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group,
|
||||
mp_size=config.tensor_parallel.tp_size) #, out_dim=0, in_dim=1)
|
||||
|
||||
def replace_with_policy(child, policy_cls, triangular_masking, inference=False, layer_id=0):
|
||||
policy = policy_cls(child, inference=inference)
|
||||
if not policy.cuda_graph_supported:
|
||||
# policy says cuda graph is not supported raise an error if set
|
||||
assert not config.enable_cuda_graph, "cuda graph is not supported with this model, please disable"
|
||||
|
||||
from deepspeed.moe.layer import MoE
|
||||
moe = False
|
||||
if hasattr(child, 'mlp') and isinstance(child.mlp, MoE):
|
||||
num_experts = child.mlp.num_experts
|
||||
moe = True
|
||||
|
||||
# 1. Create a model-specific container object using the policy object.
|
||||
_container = policy_to_ds_container(policy=policy,
|
||||
config=config,
|
||||
model_config=model_config,
|
||||
layer_id=layer_id,
|
||||
child=child)
|
||||
_container.set_moe(moe)
|
||||
|
||||
# 2. Set the tensor parallelism config
|
||||
_container.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group)
|
||||
|
||||
# 3. Initialize tensors
|
||||
_container.initialize_tensors()
|
||||
|
||||
# 4. deal with data types -- needs refactor to use dtype instead of fp16
|
||||
if config.dtype in [torch.float16, torch.bfloat16, torch.int8]:
|
||||
_container.convert_to_required_dtype()
|
||||
|
||||
# 5. Set the quantization config
|
||||
quantizer = GroupQuantizer(q_int8=quantize)
|
||||
_container.set_quantization_config(quantizer)
|
||||
|
||||
# 6. create a DS Inference config object
|
||||
_container.create_ds_model_config()
|
||||
|
||||
# 7. use the config and create the module
|
||||
_container.create_module()
|
||||
|
||||
# 8. transpose the weights and bias if needed
|
||||
_container.transpose()
|
||||
|
||||
# 9. deal with tensor parallelism.
|
||||
_container.apply_tensor_parallelism(mp_replace)
|
||||
|
||||
# 10. copy the tensors from the model-specific container to the new module
|
||||
_container.copy_data_to_new_module()
|
||||
|
||||
# 11. set global for generic checkpoint loading
|
||||
global container_g
|
||||
|
||||
if container_g is None:
|
||||
container_g = _container
|
||||
|
||||
return _container.module
|
||||
|
||||
def replace_wo_policy(module, all_reduce_linears, prefix="", state_dict=None):
|
||||
#mp_replace = ReplaceWithTensorSlicing(mp_group=config.tensor_parallel.tp_group)
|
||||
|
||||
# Get the configurable partition_config if available
|
||||
partition_config = None
|
||||
if hasattr(config, 'get_partition_config_object'):
|
||||
partition_config = config.get_partition_config_object()
|
||||
|
||||
# 1. Create AutoTP object
|
||||
_autotp = AutoTP(module,
|
||||
all_reduce_linears,
|
||||
prefix,
|
||||
state_dict,
|
||||
linear_layer_setting,
|
||||
orig_layer_impl,
|
||||
config.keep_module_on_host,
|
||||
partition_config=partition_config)
|
||||
|
||||
# 2. Set the tensor parallelism config
|
||||
_autotp.set_tensor_parallel_config(config.tensor_parallel.tp_size, config.tensor_parallel.tp_group)
|
||||
|
||||
# 3. Try to get num_key_heads from model_config.num_key_value_heads
|
||||
if hasattr(model_config, "vision_config"):
|
||||
if "MllamaVisionEncoderLayer" in str(module):
|
||||
num_kv_heads = _autotp.get_model_num_kv_heads(model_config.vision_config)
|
||||
elif hasattr(model_config, "text_config"):
|
||||
num_kv_heads = _autotp.get_model_num_kv_heads(model_config.text_config)
|
||||
else:
|
||||
num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
|
||||
else:
|
||||
num_kv_heads = _autotp.get_model_num_kv_heads(model_config)
|
||||
|
||||
# 4. When we have num_kv_heads defined, uneven division is possible, otherwise enforce even division
|
||||
set_num_kv_heads(num_kv_heads)
|
||||
|
||||
# 4.1 Get n_embd
|
||||
n_embd = None
|
||||
multi_query_n_embd_names = ['n_embd', 'hidden_size']
|
||||
for name in multi_query_n_embd_names:
|
||||
if hasattr(model_config, name):
|
||||
n_embd = getattr(model_config, name)
|
||||
if n_embd != None:
|
||||
break
|
||||
|
||||
# 4.2 set n_embd
|
||||
set_n_embd(n_embd)
|
||||
|
||||
# 4.3 set attention_heads
|
||||
if hasattr(model_config, 'num_attention_heads'):
|
||||
set_num_attention_heads(getattr(model_config, 'num_attention_heads'))
|
||||
|
||||
# 4.4 set tp_grain_size
|
||||
set_tp_grain_size(config.tensor_parallel.tp_grain_size)
|
||||
|
||||
# 5. Set linear policies
|
||||
_autotp.update_linear_policies()
|
||||
|
||||
# 6. Replace modules
|
||||
if "lm_head" in all_reduce_linears or "embed_out" in all_reduce_linears:
|
||||
return _autotp._replace_last_linear_module(module)
|
||||
return _autotp._replace_module(module)
|
||||
|
||||
def replace_fn(child, _policy, layer_id=0, prefix="", state_dict=None):
|
||||
# copy relevant state from child -> new module
|
||||
if not is_autotp_training_mode() and config.replace_with_kernel_inject:
|
||||
new_module = replace_with_policy(child,
|
||||
_policy,
|
||||
config.triangular_masking,
|
||||
inference=True,
|
||||
layer_id=layer_id)
|
||||
else:
|
||||
new_module = replace_wo_policy(child, _policy, prefix=prefix, state_dict=state_dict)
|
||||
|
||||
return new_module
|
||||
|
||||
def set_lm_head(module):
|
||||
if is_autotp_training_mode():
|
||||
# we need to handle autoTP training mode separately.
|
||||
return
|
||||
|
||||
embedding_weight = None
|
||||
for n, p in module.named_parameters():
|
||||
if "word_embeddings." in n or "embed_tokens." in n or "wte." in n:
|
||||
embedding_weight = p
|
||||
if embedding_weight is not None and hasattr(module, "lm_head") and hasattr(
|
||||
module.lm_head, "weight") and module.lm_head.weight.is_meta:
|
||||
module.lm_head.weight = embedding_weight
|
||||
# enable tensor parallel for the last linear
|
||||
if hasattr(module, "lm_head") and hasattr(module.lm_head, "weight") and isinstance(
|
||||
module.lm_head, torch.nn.Linear):
|
||||
module = replace_wo_policy(module, ("lm_head", ), 0, "lm_head")
|
||||
elif hasattr(module, "embed_out") and hasattr(module.embed_out, "weight") and isinstance(
|
||||
module.embed_out, torch.nn.Linear):
|
||||
module = replace_wo_policy(module, ("embed_out", ), 0, "embed_out")
|
||||
elif hasattr(module, "language_model") and hasattr(module.language_model, "lm_head"):
|
||||
module = replace_wo_policy(module.language_model, ("lm_head", ), 0, "lm_head")
|
||||
return module
|
||||
|
||||
def conv2d_parallel_shard_weights(model, rank, world_size):
|
||||
# add conv policy
|
||||
shard_oc_name = ["conv1"]
|
||||
shard_ic_name = ["conv2"]
|
||||
for name, sub_m in model.named_children():
|
||||
for l_name, l_sub_m in sub_m.named_children():
|
||||
if l_name in shard_oc_name:
|
||||
TPConv2d = TensorParallelOcShardConv2d(
|
||||
l_sub_m,
|
||||
rank,
|
||||
world_size,
|
||||
)
|
||||
setattr(sub_m, l_name, TPConv2d)
|
||||
if l_name in shard_ic_name:
|
||||
TPConv2d = TensorParallelIcShardConv2d(
|
||||
l_sub_m,
|
||||
rank,
|
||||
world_size,
|
||||
)
|
||||
setattr(sub_m, l_name, TPConv2d)
|
||||
conv2d_parallel_shard_weights(sub_m, rank, world_size)
|
||||
|
||||
if checkpoint_dict is not None and not config.replace_with_kernel_inject:
|
||||
# AutoTP shard loading
|
||||
checkpoint = checkpoint_dict["checkpoints"]
|
||||
pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards")
|
||||
for i in range(len(checkpoint)):
|
||||
checkpoint_file = os.path.join(config.base_dir, checkpoint[i])
|
||||
replaced_module = replace_module(model=model,
|
||||
orig_class=orig_layer_impl,
|
||||
replace_fn=replace_fn,
|
||||
_replace_policy=config.injection_policy_tuple,
|
||||
checkpoint=checkpoint_file)
|
||||
pbar.update(1)
|
||||
gc.collect()
|
||||
# conv2d tp module replace
|
||||
# Now is for yuan model. Add model list and conv policy to decide whether to replace conv.
|
||||
if 'Yuan' in str(replaced_module):
|
||||
conv2d_parallel_shard_weights(replaced_module, dist.get_rank(), dist.get_world_size())
|
||||
else:
|
||||
replaced_module = replace_module(model=model,
|
||||
orig_class=orig_layer_impl,
|
||||
replace_fn=replace_fn,
|
||||
_replace_policy=config.injection_policy_tuple)
|
||||
# AutoTP default set lm_head tp
|
||||
if not config.replace_with_kernel_inject:
|
||||
replaced_module = set_lm_head(replaced_module)
|
||||
|
||||
quantizer = GroupQuantizer(q_int8=quantize)
|
||||
world_size = dist.get_world_size() if dist.is_initialized() else 1
|
||||
rank = dist.get_rank() if dist.is_initialized() else 0
|
||||
if checkpoint_dict is not None and config.replace_with_kernel_inject:
|
||||
assert container_g.ckpt_load_enabled, \
|
||||
f"Meta Tensor checkpoint loading not supported in {container_g.__class__.__name__} container"
|
||||
start_time = time.time()
|
||||
checkpoint = checkpoint_dict['checkpoints']
|
||||
ckpt_list = checkpoint["tp"] if type(checkpoint) is dict else checkpoint
|
||||
ckpt_type = checkpoint_dict.get('parallelization', 'pp')
|
||||
ckpt_mp_size = checkpoint_dict.get('tp_size', len(ckpt_list))
|
||||
ckpt_mp_size = checkpoint_dict.get('mp_size', ckpt_mp_size)
|
||||
base_dir1 = checkpoint_dict.get('base_dir', config.base_dir)
|
||||
|
||||
if ckpt_type == 'pp' and type(checkpoint) is list:
|
||||
pbar = tqdm.tqdm(total=len(checkpoint), desc=f"Loading {len(checkpoint)} checkpoint shards")
|
||||
|
||||
for i in range(len(checkpoint)):
|
||||
sd = [torch.load(os.path.join(base_dir1, checkpoint[i]), map_location='cpu', weights_only=False)]
|
||||
load_model_with_checkpoint(replaced_module,
|
||||
sd,
|
||||
mp_replace,
|
||||
ckpt_type,
|
||||
ckpt_mp_size,
|
||||
quantizer,
|
||||
container=container_g)
|
||||
pbar.update(1)
|
||||
else:
|
||||
num_checkpoints = len(ckpt_list) // ckpt_mp_size
|
||||
tp_split_size = (world_size / ckpt_mp_size)
|
||||
sd_offset = int(rank / tp_split_size)
|
||||
sd_count = int((rank + max(1, tp_split_size)) / tp_split_size) - sd_offset
|
||||
pbar = tqdm.tqdm(total=num_checkpoints, desc=f"Loading {num_checkpoints} checkpoint shards")
|
||||
for i in range(num_checkpoints):
|
||||
pbar.update(1)
|
||||
ckpt_index = i * ckpt_mp_size + sd_offset
|
||||
ckpt_files = [
|
||||
os.path.join(base_dir1, ckpt_list[ckpt_index + j]) if base_dir1 else ckpt_list[ckpt_index + j]
|
||||
for j in range(sd_count)
|
||||
]
|
||||
sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False) for ckpt_file in ckpt_files]
|
||||
load_model_with_checkpoint(replaced_module,
|
||||
sds,
|
||||
mp_replace,
|
||||
ckpt_type,
|
||||
ckpt_mp_size,
|
||||
quantizer,
|
||||
int(rank % tp_split_size),
|
||||
container=container_g)
|
||||
sds = [None for _ in sds]
|
||||
gc.collect()
|
||||
|
||||
if "non_tp" in checkpoint:
|
||||
pbar = tqdm.tqdm(total=len(checkpoint["non_tp"]),
|
||||
desc=f"Loading {len(checkpoint['non_tp'])} checkpoint shards")
|
||||
|
||||
for i in range(len(checkpoint["non_tp"])):
|
||||
pbar.update(1)
|
||||
ckpt_file = os.path.join(base_dir1,
|
||||
checkpoint["non_tp"][i]) if base_dir1 else checkpoint["non_tp"][i]
|
||||
sds = [torch.load(ckpt_file, map_location='cpu', weights_only=False)]
|
||||
load_model_with_checkpoint(replaced_module,
|
||||
sds,
|
||||
mp_replace,
|
||||
ckpt_type,
|
||||
ckpt_mp_size,
|
||||
quantizer,
|
||||
int(rank % tp_split_size),
|
||||
container=container_g)
|
||||
sds = [None for _ in sds]
|
||||
gc.collect()
|
||||
set_lm_head(replaced_module)
|
||||
print(f"checkpoint loading time at rank {rank}: {time.time()-start_time} sec")
|
||||
|
||||
if not is_autotp_training_mode() and config.save_mp_checkpoint_path is not None:
|
||||
from collections import OrderedDict
|
||||
import json
|
||||
num_partitions = 8
|
||||
|
||||
if checkpoint_dict is None:
|
||||
ckpt_name = "ds_model"
|
||||
try:
|
||||
from transformers.models.bloom.modeling_bloom import BloomForCausalLM
|
||||
if isinstance(model, BloomForCausalLM):
|
||||
ckpt_name = "bloom"
|
||||
except ImportError:
|
||||
ckpt_name = "ds_model"
|
||||
else:
|
||||
ckpt_name = checkpoint_dict['type']
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
transformer_name = get_transformer_name(replaced_module)
|
||||
non_tp_ckpt_name = 'non-tp.pt'
|
||||
ckpt_files = [non_tp_ckpt_name]
|
||||
os.makedirs(config.save_mp_checkpoint_path, exist_ok=True)
|
||||
|
||||
if not dist.is_initialized() or dist.get_rank() == 0:
|
||||
print("Saving tp-sharded checkpoints")
|
||||
torch.save(
|
||||
OrderedDict({
|
||||
k: v
|
||||
for k, v in dict(replaced_module.state_dict()).items() if transformer_name not in k
|
||||
}), f'{config.save_mp_checkpoint_path}/{non_tp_ckpt_name}')
|
||||
|
||||
dtype_reprs = {
|
||||
torch.float32: 'float32',
|
||||
torch.float16: 'float16',
|
||||
torch.int8: 'int8',
|
||||
torch.bfloat16: 'bfloat16'
|
||||
}
|
||||
|
||||
ckpt_config = json.dumps({
|
||||
'type': ckpt_name,
|
||||
'base_dir': f'{config.save_mp_checkpoint_path}',
|
||||
'checkpoints': {
|
||||
"non_tp": ckpt_files,
|
||||
"tp": [f'tp_{r:0>2d}_{m:0>2d}.pt' for m in range(num_partitions) for r in range(world_size)]
|
||||
},
|
||||
'version': 1.0,
|
||||
'parallelization': 'tp',
|
||||
'tp_size': world_size,
|
||||
'dtype': dtype_reprs[config.dtype]
|
||||
})
|
||||
with open(f"{config.save_mp_checkpoint_path}/ds_inference_config.json", "w") as cfg:
|
||||
cfg.write(ckpt_config)
|
||||
|
||||
rep_sd = replaced_module.state_dict()
|
||||
for n, p in replaced_module.named_parameters():
|
||||
if hasattr(p, 'scale'):
|
||||
rep_sd[n] = [p, p.scale]
|
||||
keys = list(rep_sd.keys())
|
||||
partition_size = (len(keys) // num_partitions + 1)
|
||||
for m in range(num_partitions):
|
||||
torch.save(
|
||||
OrderedDict({
|
||||
k: [rep_sd[k], rep_sd[k].scale] if hasattr(rep_sd[k], 'scale') else rep_sd[k]
|
||||
for k in keys[m * partition_size:(m + 1) * partition_size] if transformer_name in k
|
||||
}), f'{config.save_mp_checkpoint_path}/tp_{rank:0>2d}_{m:0>2d}.pt')
|
||||
|
||||
return replaced_module
|
||||
|
||||
|
||||
def revert_transformer_layer(orig_layer_impl, model, config, preln=False):
|
||||
""" Revert DeepSpeed's transformer layer back to original bert-style transformer layer
|
||||
Arguments:
|
||||
orig_layer_impl (torch.nn.Module): the original transformer layer implementation that was replaced,
|
||||
e.g., transformers.models.bert.modeling_bert.BertLayer or transformers.BertLayer
|
||||
model (torch.nn.Module): user's nn.module representing their model
|
||||
config (dict): model config containing hidden size, attention heads, etc.
|
||||
Returns:
|
||||
Updated nn.module with original bert-style transformer layers
|
||||
"""
|
||||
|
||||
def replace_fn(child, _replace_policy, layer_id):
|
||||
#from turing.nvidia_modelingpreln import BertLayer
|
||||
orig_module = orig_layer_impl(config)
|
||||
|
||||
# copy relevant state from child -> original module
|
||||
qkvw = child.attn_qkvw.data
|
||||
qkvb = child.attn_qkvb.data
|
||||
|
||||
qw, kw, vw = torch.chunk(qkvw, 3, axis=0)
|
||||
qb, kb, vb = torch.chunk(qkvb, 3, axis=0)
|
||||
|
||||
orig_module.attention.self.query.weight.data = qw
|
||||
orig_module.attention.self.query.bias.data = qb
|
||||
orig_module.attention.self.key.weight.data = kw
|
||||
orig_module.attention.self.key.bias.data = kb
|
||||
orig_module.attention.self.value.weight.data = vw
|
||||
orig_module.attention.self.value.bias.data = vb
|
||||
|
||||
orig_module.attention.output.dense.weight.data = child.attn_ow.data
|
||||
orig_module.attention.output.dense.bias.data = child.attn_ob.data
|
||||
|
||||
attn_ln_w = child.attn_nw.data
|
||||
attn_ln_b = child.attn_nb.data
|
||||
if preln:
|
||||
orig_module.PostAttentionLayerNorm.weight.data = attn_ln_w
|
||||
orig_module.PostAttentionLayerNorm.bias.data = attn_ln_b
|
||||
else:
|
||||
orig_module.attention.output.LayerNorm.weight.data = attn_ln_w
|
||||
orig_module.attention.output.LayerNorm.bias.data = attn_ln_b
|
||||
|
||||
inter_ff_w = child.inter_w.data
|
||||
inter_ff_b = child.inter_b.data
|
||||
if preln:
|
||||
orig_module.intermediate.dense_act.weight.data = inter_ff_w
|
||||
orig_module.intermediate.dense_act.bias.data = inter_ff_b
|
||||
else:
|
||||
orig_module.intermediate.dense.weight.data = inter_ff_w
|
||||
orig_module.intermediate.dense.bias.data = inter_ff_b
|
||||
|
||||
orig_module.output.dense.weight.data = child.output_w.data
|
||||
orig_module.output.dense.bias.data = child.output_b.data
|
||||
|
||||
transformer_ln_w = child.norm_w.data
|
||||
transformer_ln_b = child.norm_b.data
|
||||
if preln:
|
||||
orig_module.PreAttentionLayerNorm.weight.data = transformer_ln_w
|
||||
orig_module.PreAttentionLayerNorm.bias.data = transformer_ln_b
|
||||
else:
|
||||
orig_module.output.LayerNorm.weight.data = transformer_ln_w
|
||||
orig_module.output.LayerNorm.bias.data = transformer_ln_b
|
||||
return orig_module
|
||||
|
||||
return replace_module(model=model,
|
||||
orig_class=deepspeed.DeepSpeedTransformerLayer,
|
||||
replace_fn=replace_fn,
|
||||
_replace_policy=None)
|
||||
|
||||
|
||||
def replace_module(model, orig_class, replace_fn, _replace_policy, checkpoint=None):
|
||||
""" Scan the model for instances of ``orig_clas:`` to replace using ``replace_fn``.
|
||||
Arguments:
|
||||
model (torch.nn.Module): the model to augment
|
||||
orig_class (torch.nn.Module): the module to search for
|
||||
replace_fn (method): a method to convert instances of ``orig_class`` to the
|
||||
desired type and return a new instance.
|
||||
Returns:
|
||||
A modified ``model``.
|
||||
"""
|
||||
sd = None
|
||||
if checkpoint is not None:
|
||||
if checkpoint.endswith(".safetensors"):
|
||||
from safetensors.torch import load_file
|
||||
sd = load_file(checkpoint)
|
||||
else:
|
||||
sd = torch.load(checkpoint, map_location='cpu', weights_only=False)
|
||||
|
||||
policy = {}
|
||||
if orig_class is not None:
|
||||
policy.update({orig_class: (replace_fn, _replace_policy)})
|
||||
else:
|
||||
for plcy in replace_policies:
|
||||
# instantiate a throw-away policy in order to populate the _orig_layer_class
|
||||
_ = plcy(None)
|
||||
if isinstance(plcy._orig_layer_class, list):
|
||||
for orig_layer_class in plcy._orig_layer_class:
|
||||
policy.update({orig_layer_class: (replace_fn, plcy)})
|
||||
elif plcy._orig_layer_class is not None:
|
||||
policy.update({plcy._orig_layer_class: (replace_fn, plcy)})
|
||||
assert len(policy.items()) > 0,\
|
||||
"No default policy found! Please specify your policy injection_policy (like {BertLayer:HFBEertLayerPolicy})." +\
|
||||
"You can find some samples here: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/module_inject/replace_policy.py"
|
||||
|
||||
replaced_module, _ = _replace_module(model, policy, state_dict=sd)
|
||||
return replaced_module
|
||||
|
||||
|
||||
from ..pipe import PipelineModule
|
||||
|
||||
import re
|
||||
|
||||
|
||||
def skip_level_0_prefix(model, state_dict):
|
||||
model = str(model)
|
||||
key = re.search(r": (.*?)Model", model)
|
||||
if key is None:
|
||||
key = re.search(r": (.*?)Stack", model)
|
||||
if key is None:
|
||||
key = re.match(r"(.*?)Model", model)
|
||||
# if keys start with 'model.', don't skip level 0 prefix
|
||||
if state_dict is not None:
|
||||
for item in state_dict.keys():
|
||||
if re.match("^model[.]", item):
|
||||
return False
|
||||
if key is not None and key.group(1).lower() in ["bloom", "opt"]:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _replace_module(model, policies, prefix='', layer_id=0, level_id=0, state_dict=None):
|
||||
""" Traverse model's children recursively and apply any transformations in ``policies``.
|
||||
Arguments:
|
||||
model (torch.nn.Module): model to augment
|
||||
policies (dict): Mapping of source class to replacement function.
|
||||
Returns:
|
||||
Modified ``model``.
|
||||
"""
|
||||
for name, child in model.named_children():
|
||||
if child.__class__ in policies:
|
||||
replaced_module = policies[child.__class__][0](child,
|
||||
policies[child.__class__][-1],
|
||||
layer_id,
|
||||
prefix=prefix + name,
|
||||
state_dict=state_dict)
|
||||
setattr(model, name, replaced_module)
|
||||
if isinstance(model, PipelineModule):
|
||||
assert hasattr(model, 'forward_funcs'),\
|
||||
"we require pipe-module to have the list of fwd_functions"
|
||||
model.forward_funcs[model.fwd_map[name]] = replaced_module
|
||||
layer_id += 1
|
||||
else:
|
||||
checking_key = prefix + name + '.'
|
||||
if Loading.is_load_module(child) and state_dict is not None:
|
||||
if any(checking_key in item for item in state_dict):
|
||||
Loading.load(
|
||||
child,
|
||||
state_dict,
|
||||
checking_key,
|
||||
)
|
||||
else:
|
||||
continue
|
||||
if len(child._buffers) != 0 and state_dict is not None:
|
||||
Loading.load_buffer(child, state_dict, checking_key)
|
||||
_, layer_id = _replace_module(child,
|
||||
policies,
|
||||
prefix if level_id == 0 and skip_level_0_prefix(model, state_dict) else \
|
||||
prefix + name + '.',
|
||||
layer_id=layer_id,
|
||||
level_id=level_id + 1,
|
||||
state_dict=state_dict)
|
||||
|
||||
# Add the reset_cache func to the model, so that it can be called in the beginning of text-generation.
|
||||
model.reset_cache = transformer_inference.DeepSpeedTransformerInference.reset_cache
|
||||
return model, layer_id
|
||||
Executable
+30
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from .containers import HFGPT2LayerPolicy
|
||||
from .containers import HFBertLayerPolicy
|
||||
from .containers import BLOOMLayerPolicy
|
||||
from .containers import HFGPTJLayerPolicy
|
||||
from .containers import HFGPTNEOLayerPolicy
|
||||
from .containers import GPTNEOXLayerPolicy
|
||||
from .containers import HFOPTLayerPolicy
|
||||
from .containers import MegatronLayerPolicy
|
||||
from .containers import HFDistilBertLayerPolicy
|
||||
from .containers import HFCLIPLayerPolicy
|
||||
from .containers import LLAMALayerPolicy
|
||||
from .containers import UNetPolicy
|
||||
from .containers import VAEPolicy
|
||||
from .containers import LLAMA2LayerPolicy
|
||||
from .containers import InternLMLayerPolicy
|
||||
|
||||
# transformer-based policies
|
||||
replace_policies = [
|
||||
HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy,
|
||||
HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy,
|
||||
LLAMALayerPolicy, LLAMA2LayerPolicy, InternLMLayerPolicy
|
||||
]
|
||||
|
||||
# non-transformer-based policies
|
||||
generic_policies = [UNetPolicy, VAEPolicy]
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import logging
|
||||
from typing import List, Dict, Optional
|
||||
from .autotp_config import TPLayerSpec, PartitionType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED_STYLES = {"colwise", "rowwise"}
|
||||
|
||||
|
||||
class TPPlanConverter:
|
||||
"""Convert HuggingFace tp_plan format to DeepSpeed TPLayerSpec format."""
|
||||
|
||||
@staticmethod
|
||||
def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
|
||||
"""Convert HF tp_plan to DeepSpeed layer specs.
|
||||
|
||||
Returns None if the plan contains any unsupported partition styles,
|
||||
allowing the caller to fall back to the existing AutoTP path.
|
||||
"""
|
||||
unsupported = {style for style in hf_tp_plan.values() if style.lower() not in SUPPORTED_STYLES}
|
||||
if unsupported:
|
||||
logger.warning(
|
||||
"HuggingFace tp_plan contains unsupported partition style(s): %s. "
|
||||
"Falling back to AutoTP preset-based partitioning.", sorted(unsupported))
|
||||
return None
|
||||
|
||||
layer_specs = []
|
||||
|
||||
for pattern, partition in hf_tp_plan.items():
|
||||
regex_pattern = TPPlanConverter._wildcard_to_regex(pattern)
|
||||
|
||||
if partition.lower() == "colwise":
|
||||
partition_type = PartitionType.COLUMN
|
||||
elif partition.lower() == "rowwise":
|
||||
partition_type = PartitionType.ROW
|
||||
|
||||
# Only add .weight suffix if not already present
|
||||
if not regex_pattern.endswith(r"\.weight"):
|
||||
regex_pattern += r"\.weight$"
|
||||
else:
|
||||
regex_pattern += r"$"
|
||||
|
||||
layer_specs.append(TPLayerSpec(
|
||||
patterns=[regex_pattern],
|
||||
partition_type=partition_type,
|
||||
))
|
||||
|
||||
return layer_specs
|
||||
|
||||
@staticmethod
|
||||
def _wildcard_to_regex(pattern: str) -> str:
|
||||
regex = pattern.replace('.', r'\.')
|
||||
regex = regex.replace('*', r'.*')
|
||||
return ".*" + regex
|
||||
@@ -0,0 +1,79 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed import comm as dist
|
||||
|
||||
# Defaults for optional TP globals. These can be overridden by setters.
|
||||
num_kv_heads = None
|
||||
num_attention_heads = None
|
||||
n_embd = None
|
||||
tp_grain_size = 1
|
||||
|
||||
|
||||
def set_num_kv_heads(num):
|
||||
global num_kv_heads
|
||||
num_kv_heads = num
|
||||
|
||||
|
||||
def set_num_attention_heads(num):
|
||||
global num_attention_heads
|
||||
num_attention_heads = num
|
||||
|
||||
|
||||
def set_n_embd(num):
|
||||
global n_embd
|
||||
n_embd = num
|
||||
|
||||
|
||||
def set_tp_grain_size(num):
|
||||
global tp_grain_size
|
||||
tp_grain_size = num
|
||||
|
||||
|
||||
def get_num_kv_heads():
|
||||
global num_kv_heads
|
||||
if 'num_kv_heads' in globals():
|
||||
return num_kv_heads
|
||||
return None
|
||||
|
||||
|
||||
def get_num_attention_heads():
|
||||
global num_attention_heads
|
||||
return num_attention_heads
|
||||
|
||||
|
||||
def get_shard_size(total_size, mp_size, name=None, rank=None):
|
||||
global num_kv_heads
|
||||
last_linear = ["lm_head", "embed_out"]
|
||||
# MoE MLP layer use near even division will get better perf.
|
||||
moe_mlp_layer = ["gate_proj", "up_proj", "down_proj", "w1", "w2", "w3"]
|
||||
not_moe_mlp_layer = True
|
||||
if name != None and any(s in str(name) for s in moe_mlp_layer):
|
||||
not_moe_mlp_layer = False
|
||||
# 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 and "mlp" not in str(name) and str(
|
||||
name) not in last_linear and not_moe_mlp_layer:
|
||||
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 = total_size // tp_grain_size
|
||||
return (grain_size // mp_size + (1 if rank < (grain_size % mp_size) else 0)) * tp_grain_size
|
||||
else:
|
||||
return total_size // mp_size + (1 if rank < (total_size % mp_size) else 0)
|
||||
|
||||
|
||||
def get_n_embd():
|
||||
global n_embd
|
||||
return n_embd
|
||||
|
||||
|
||||
def get_shard_size_list(total_size, mp_size, name=None):
|
||||
shard_sizes = []
|
||||
for i in range(mp_size):
|
||||
shard_sizes.append(get_shard_size(total_size, mp_size, name, i))
|
||||
return shard_sizes
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.utils import log_dist
|
||||
|
||||
|
||||
def transpose(data):
|
||||
return data.transpose(-1, -2).contiguous()
|
||||
|
||||
|
||||
# helper function to map between DS policies and DS containers
|
||||
def policy_to_ds_container(**kwargs):
|
||||
from .containers import HFGPT2LayerPolicy, DS_GPT2Container
|
||||
from .containers import HFBertLayerPolicy, DS_BERTContainer
|
||||
from .containers import BLOOMLayerPolicy, DS_BloomContainer
|
||||
from .containers import HFGPTJLayerPolicy, DS_GPTJContainer
|
||||
from .containers import HFGPTNEOLayerPolicy, DS_GPTNEOContainer
|
||||
from .containers import GPTNEOXLayerPolicy, DS_GPTNEOXContainer
|
||||
from .containers import HFOPTLayerPolicy, DS_OPTContainer
|
||||
from .containers import MegatronLayerPolicy, DS_MegatronGPTContainer
|
||||
from .containers import HFDistilBertLayerPolicy, DS_DistilBERTContainer
|
||||
from .containers import LLAMALayerPolicy, DS_LLAMAContainer
|
||||
from .containers import LLAMA2LayerPolicy, DS_LLAMA2Container
|
||||
from .containers import InternLMLayerPolicy, DS_InternLMContainer
|
||||
|
||||
policy_to_container = {
|
||||
HFGPT2LayerPolicy: DS_GPT2Container,
|
||||
HFBertLayerPolicy: DS_BERTContainer,
|
||||
BLOOMLayerPolicy: DS_BloomContainer,
|
||||
HFGPTJLayerPolicy: DS_GPTJContainer,
|
||||
HFGPTNEOLayerPolicy: DS_GPTNEOContainer,
|
||||
GPTNEOXLayerPolicy: DS_GPTNEOXContainer,
|
||||
HFOPTLayerPolicy: DS_OPTContainer,
|
||||
MegatronLayerPolicy: DS_MegatronGPTContainer,
|
||||
HFDistilBertLayerPolicy: DS_DistilBERTContainer,
|
||||
LLAMALayerPolicy: DS_LLAMAContainer,
|
||||
LLAMA2LayerPolicy: DS_LLAMA2Container,
|
||||
InternLMLayerPolicy: DS_InternLMContainer
|
||||
}
|
||||
|
||||
container = None
|
||||
policy = kwargs['policy']
|
||||
assert policy is not None, "Policy cannot be None"
|
||||
policy_type = type(policy)
|
||||
|
||||
if policy_type not in policy_to_container:
|
||||
log_dist(f"Policy type {policy_type} not supported", [0])
|
||||
else:
|
||||
container = policy_to_container[policy_type](**kwargs)
|
||||
|
||||
return container
|
||||
Reference in New Issue
Block a user