# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team """Expert weight repacking for AutoEP. Converts HuggingFace expert weight formats into TorchTitan-compatible grouped tensors [E_local, hidden_dim, dim] for grouped GEMM. """ from __future__ import annotations from contextlib import contextmanager import torch import torch.nn as nn from deepspeed.module_inject.auto_ep_config import MoELayerSpec from deepspeed.moe.fused_expert_layout import classify_fused_gate_up_layout from deepspeed.runtime.zero import GatheredParameters from deepspeed.runtime.zero.utils import is_zero_param @contextmanager def _gather_source_zero_params(params): """Gather source ZeRO params while AutoEP reads full tensor values.""" zero_params = [param for param in params if is_zero_param(param)] if not zero_params: yield return with GatheredParameters(zero_params, modifier_rank=None, enabled=True): yield def _source_data(param: torch.Tensor | nn.Parameter) -> torch.Tensor: return param.data if torch.is_tensor(param) else param def repack_expert_weights( experts_source: nn.Module, spec: MoELayerSpec, ep_rank: int, ep_size: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Repack expert weights from HF format to TorchTitan grouped format. Returns (w1, w2, w3) where: w1: [E_local, ffn_hidden_size, hidden_size] w2: [E_local, hidden_size, ffn_hidden_size] w3: [E_local, ffn_hidden_size, hidden_size] For fused_3d storage where expert_w3 is None (gate+up fused): Standard HF layout: Source gate_up_proj: [E, 2*ffn_hidden, hidden] Source down_proj: [E, hidden, ffn_hidden] Alternate Transformers fused layout: Source gate_up_proj: [E, hidden, 2*ffn_hidden] Source down_proj: [E, ffn_hidden, hidden] In both cases, the returned grouped-expert tensors are normalized to: w1 = gate_proj: [E_local, ffn_hidden, hidden] w3 = up_proj: [E_local, ffn_hidden, hidden] w2 = down_proj: [E_local, hidden, ffn_hidden] """ num_local_experts = spec.num_experts // ep_size expert_start = ep_rank * num_local_experts expert_end = expert_start + num_local_experts if spec.expert_storage == "fused_3d": return _repack_fused_3d(experts_source, spec, expert_start, expert_end) elif spec.expert_storage == "module_list": return _repack_module_list(experts_source, spec, expert_start, expert_end) else: raise ValueError(f"Unknown expert_storage type: {spec.expert_storage}") def repack_expert_requires_grad_flags( experts_source: nn.Module, spec: MoELayerSpec, ep_rank: int, ep_size: int, ) -> tuple[bool, bool, bool]: """Return the requires_grad flags for repacked (w1, w2, w3) tensors.""" num_local_experts = spec.num_experts // ep_size expert_start = ep_rank * num_local_experts expert_end = expert_start + num_local_experts if spec.expert_storage == "fused_3d": return _repack_fused_3d_requires_grad_flags(experts_source, spec) elif spec.expert_storage == "module_list": return _repack_module_list_requires_grad_flags(experts_source, spec, expert_start, expert_end) else: raise ValueError(f"Unknown expert_storage type: {spec.expert_storage}") def _repack_fused_3d( experts_source: nn.Module, spec: MoELayerSpec, expert_start: int, expert_end: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Repack from fused 3D parameter tensors (transformers 5.0.0+).""" w1_full = getattr(experts_source, spec.expert_w1_name) w2_full = getattr(experts_source, spec.expert_w2_name) source_params = [w1_full, w2_full] if spec.expert_w3_name is not None: source_params.append(getattr(experts_source, spec.expert_w3_name)) with _gather_source_zero_params(source_params): w1_full_data = _source_data(w1_full) w2_full_data = _source_data(w2_full) # Slice to local experts w1_local = w1_full_data[expert_start:expert_end].clone() w2_local = w2_full_data[expert_start:expert_end].clone() if spec.expert_w3_name is None: layout = classify_fused_gate_up_layout(tuple(w1_local.shape), tuple(w2_local.shape)) if layout is None: raise ValueError("Unsupported fused expert weight layout for AutoEP repacking: " f"{spec.expert_w1_name}={tuple(w1_local.shape)}, " f"{spec.expert_w2_name}={tuple(w2_local.shape)}") ffn_hidden = layout.ffn_hidden_size if layout.layout == "gate_up_first": w1 = w1_local[:, :ffn_hidden, :].contiguous() # [E_local, ffn, hidden] w3 = w1_local[:, ffn_hidden:, :].contiguous() # [E_local, ffn, hidden] w2 = w2_local.contiguous() # [E_local, hidden, ffn] else: w1 = w1_local[:, :, :ffn_hidden].transpose(1, 2).contiguous() # [E_local, ffn, hidden] w3 = w1_local[:, :, ffn_hidden:].transpose(1, 2).contiguous() # [E_local, ffn, hidden] w2 = w2_local.transpose(1, 2).contiguous() # [E_local, hidden, ffn] else: # Separate w1 (gate), w3 (up) w3_full = getattr(experts_source, spec.expert_w3_name) w3_local = _source_data(w3_full)[expert_start:expert_end].clone() w1 = w1_local.contiguous() # [E_local, ffn, hidden] w2 = w2_local.contiguous() # [E_local, hidden, ffn] w3 = w3_local.contiguous() # [E_local, ffn, hidden] return w1, w2, w3 def _repack_fused_3d_requires_grad_flags( experts_source: nn.Module, spec: MoELayerSpec, ) -> tuple[bool, bool, bool]: w1_param = getattr(experts_source, spec.expert_w1_name) w2_param = getattr(experts_source, spec.expert_w2_name) w1_requires_grad = _requires_grad(w1_param) w2_requires_grad = _requires_grad(w2_param) if spec.expert_w3_name is None: w3_requires_grad = w1_requires_grad else: w3_requires_grad = _requires_grad(getattr(experts_source, spec.expert_w3_name)) return w1_requires_grad, w2_requires_grad, w3_requires_grad def _repack_module_list( experts_source: nn.Module, spec: MoELayerSpec, expert_start: int, expert_end: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Repack from nn.ModuleList of individual expert modules (legacy transformers).""" assert isinstance(experts_source, nn.ModuleList), \ f"Expected nn.ModuleList for module_list storage, got {type(experts_source)}" w1_list = [] w2_list = [] w3_list = [] for expert_idx in range(len(experts_source)): expert = experts_source[expert_idx] # Get weight tensors - handle both nn.Linear children and direct attributes w1_param = _get_expert_weight(expert, spec.expert_w1_name) w2_param = _get_expert_weight(expert, spec.expert_w2_name) w3_param = _get_expert_weight(expert, spec.expert_w3_name) if spec.expert_w3_name is not None else None with _gather_source_zero_params([w1_param, w2_param, w3_param]): if expert_start <= expert_idx < expert_end: # nn.Linear stores weight as [out_features, in_features]. # TorchTitan expects [ffn_hidden, hidden] for w1/w3 and [hidden, ffn_hidden] for w2. # nn.Linear.weight is already [out, in], which matches TorchTitan's [ffn, hidden] for w1. w1_list.append(w1_param.data.clone()) w2_list.append(w2_param.data.clone()) if w3_param is not None: w3_list.append(w3_param.data.clone()) _require_consistent_dtype_device(w1_list, spec.expert_w1_name, expert_start, expert_end) _require_consistent_dtype_device(w2_list, spec.expert_w2_name, expert_start, expert_end) if spec.expert_w3_name is not None: _require_consistent_dtype_device(w3_list, spec.expert_w3_name, expert_start, expert_end) w1 = torch.stack(w1_list) # [E_local, ffn_hidden, hidden] w2 = torch.stack(w2_list) # [E_local, hidden, ffn_hidden] if spec.expert_w3_name is not None: w3 = torch.stack(w3_list) # [E_local, ffn_hidden, hidden] else: # If no w3, this is fused gate+up - split w1 ffn_hidden = w1.shape[1] // 2 w3 = w1[:, ffn_hidden:, :].contiguous() w1 = w1[:, :ffn_hidden, :].contiguous() return w1, w2, w3 def _repack_module_list_requires_grad_flags( experts_source: nn.Module, spec: MoELayerSpec, expert_start: int, expert_end: int, ) -> tuple[bool, bool, bool]: assert isinstance(experts_source, nn.ModuleList), \ f"Expected nn.ModuleList for module_list storage, got {type(experts_source)}" w1_flags = [] w2_flags = [] w3_flags = [] for expert_idx in range(expert_start, expert_end): expert = experts_source[expert_idx] w1_flags.append(_get_expert_weight(expert, spec.expert_w1_name).requires_grad) w2_flags.append(_get_expert_weight(expert, spec.expert_w2_name).requires_grad) if spec.expert_w3_name is not None: w3_flags.append(_get_expert_weight(expert, spec.expert_w3_name).requires_grad) w1_requires_grad = _require_consistent_requires_grad( w1_flags, spec.expert_w1_name, expert_start, expert_end, ) w2_requires_grad = _require_consistent_requires_grad( w2_flags, spec.expert_w2_name, expert_start, expert_end, ) if spec.expert_w3_name is None: w3_requires_grad = w1_requires_grad else: w3_requires_grad = _require_consistent_requires_grad( w3_flags, spec.expert_w3_name, expert_start, expert_end, ) return w1_requires_grad, w2_requires_grad, w3_requires_grad def _get_expert_weight(expert_module: nn.Module, weight_name: str) -> torch.Tensor: """Get expert weight tensor by name, handling both attribute and child module patterns.""" # Direct attribute param = getattr(expert_module, weight_name, None) if param is not None: if isinstance(param, nn.Linear): return param.weight if isinstance(param, (nn.Parameter, torch.Tensor)): return param # Try as child module name for name, child in expert_module.named_children(): if name == weight_name: if isinstance(child, nn.Linear): return child.weight if hasattr(child, 'weight'): return child.weight raise ValueError(f"Could not find weight '{weight_name}' in expert module " f"{type(expert_module).__name__}. Available attributes: " f"{[n for n, _ in expert_module.named_parameters(recurse=False)]}") def _requires_grad(weight: torch.Tensor | nn.Parameter) -> bool: return bool(getattr(weight, "requires_grad", False)) def _require_consistent_dtype_device( tensors: list[torch.Tensor], weight_name: str, expert_start: int, expert_end: int, ) -> None: expected_dtype = tensors[0].dtype expected_device = tensors[0].device for expert_idx, tensor in enumerate(tensors, start=expert_start): if tensor.dtype != expected_dtype or tensor.device != expected_device: raise ValueError("AutoEP cannot preserve mixed dtype/device for " f"module_list experts {expert_start}:{expert_end} weight '{weight_name}' " "because they are packed into one grouped expert tensor. " f"Expected {expected_dtype} on {expected_device}, " f"found {tensor.dtype} on {tensor.device} at expert {expert_idx}.") def _require_consistent_requires_grad( flags: list[bool], weight_name: str, expert_start: int, expert_end: int, ) -> bool: if len(set(flags)) != 1: raise ValueError("AutoEP cannot preserve mixed requires_grad flags for " f"module_list experts {expert_start}:{expert_end} weight '{weight_name}' " "because they are packed into one grouped expert tensor.") return flags[0]