This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,49 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from transformers.utils import strtobool
|
||||
|
||||
from .fsdp import NPUCastError
|
||||
from .mindspeed import patch_mindspeed_te_cp_implementation
|
||||
|
||||
_APPLIED = False
|
||||
_ENABLE_NPU_MODEL_PATCH_ARGS = ('--enable_npu_model_patch', '--enable-npu-model-patch')
|
||||
|
||||
|
||||
def _parse_model_patch_enabled(value: str) -> bool:
|
||||
try:
|
||||
return bool(strtobool(value))
|
||||
except ValueError as exc:
|
||||
raise ValueError('--enable_npu_model_patch must be true or false.') from exc
|
||||
|
||||
|
||||
def _is_model_patch_enabled_from_argv() -> bool:
|
||||
for i, arg in enumerate(sys.argv):
|
||||
if arg in _ENABLE_NPU_MODEL_PATCH_ARGS:
|
||||
if i + 1 >= len(sys.argv) or sys.argv[i + 1].startswith('--'):
|
||||
raise ValueError('--enable_npu_model_patch requires a value: true or false.')
|
||||
return _parse_model_patch_enabled(sys.argv[i + 1])
|
||||
if any(arg.startswith(f'{name}=') for name in _ENABLE_NPU_MODEL_PATCH_ARGS):
|
||||
value = arg.split('=', 1)[1]
|
||||
return _parse_model_patch_enabled(value)
|
||||
return True
|
||||
|
||||
|
||||
def apply_all_patches() -> None:
|
||||
global _APPLIED
|
||||
if _APPLIED:
|
||||
return
|
||||
|
||||
from . import env, fsdp
|
||||
|
||||
env.apply_patch()
|
||||
fsdp.apply_patch()
|
||||
# The model patch switch is checked only on the first import; monkey patches are not reversible.
|
||||
if _is_model_patch_enabled_from_argv():
|
||||
from . import model
|
||||
model.apply_patch()
|
||||
_APPLIED = True
|
||||
|
||||
|
||||
__all__ = ['NPUCastError', 'apply_all_patches', 'patch_mindspeed_te_cp_implementation']
|
||||
@@ -0,0 +1,51 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
_DEFAULT_NPU_HCCL_CONNECT_TIMEOUT = '600'
|
||||
_TORCH_NPU_GETENV_MODULE = 'torch_npu.utils.patch_getenv'
|
||||
|
||||
|
||||
def _patch_torch_npu_getenv() -> None:
|
||||
try:
|
||||
from torch_npu.utils import patch_getenv
|
||||
except Exception:
|
||||
return
|
||||
|
||||
orig_environ_get = getattr(patch_getenv, '_orig_environ_get', None)
|
||||
current_get = os.environ.get
|
||||
current_getenv = os.getenv
|
||||
getenv_module = getattr(current_getenv, '__module__', None)
|
||||
environ_get_module = getattr(current_get, '__module__', None)
|
||||
if not (getenv_module == _TORCH_NPU_GETENV_MODULE or environ_get_module == _TORCH_NPU_GETENV_MODULE):
|
||||
return
|
||||
if getattr(orig_environ_get, '__self__', None) is None:
|
||||
return
|
||||
|
||||
log_once = getattr(patch_getenv, '_log_once', None)
|
||||
|
||||
def _get_from_current_environ(key, default=None):
|
||||
hit = key in os.environ
|
||||
value = os.environ[key] if hit else default
|
||||
if hit and isinstance(value, str) and value != '' and log_once is not None:
|
||||
log_once(key, value)
|
||||
return value
|
||||
|
||||
os.getenv = _get_from_current_environ
|
||||
os.environ.get = _get_from_current_environ
|
||||
logger.info('Patched torch_npu getenv to read from current os.environ.')
|
||||
|
||||
|
||||
def apply_patch() -> None:
|
||||
_patch_torch_npu_getenv()
|
||||
|
||||
if 'HCCL_CONNECT_TIMEOUT' in os.environ:
|
||||
return
|
||||
|
||||
os.environ['HCCL_CONNECT_TIMEOUT'] = _DEFAULT_NPU_HCCL_CONNECT_TIMEOUT
|
||||
logger.info(f'Set HCCL_CONNECT_TIMEOUT={_DEFAULT_NPU_HCCL_CONNECT_TIMEOUT} by default for NPU.')
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import accelerate.utils.fsdp_utils as fsdp_utils
|
||||
import torch
|
||||
from accelerate.accelerator import Accelerator
|
||||
from functools import wraps
|
||||
|
||||
|
||||
class NPUCastError(RuntimeError):
|
||||
"""Raised when fp32 casting fails during NPU FSDP2 preparation."""
|
||||
|
||||
|
||||
def _cast_module_to_fp32_for_npu_if_needed(module: torch.nn.Module, accelerator: Accelerator) -> torch.nn.Module:
|
||||
if accelerator.device.type != 'npu':
|
||||
return module
|
||||
|
||||
param = next(module.parameters(recurse=True), None)
|
||||
if param is None:
|
||||
return module
|
||||
|
||||
if not param.is_floating_point() or param.dtype == torch.float32:
|
||||
return module
|
||||
|
||||
# Accelerate FSDP2 flattens and shards parameters during prepare. On NPU,
|
||||
# entering that path with bf16/fp16 parameters can fail before mixed
|
||||
# precision policy has a chance to manage runtime compute dtype. Cast early
|
||||
# while parameters are still on CPU or meta, so only dtype changes here.
|
||||
|
||||
# GRPO with vLLM colocate mode may preload the model onto NPU before
|
||||
# Accelerator.prepare() is called. In that case, casting fp32 on NPU
|
||||
# would temporarily duplicate the full model (bf16 + fp32), causing OOM.
|
||||
# We move the model back to CPU first to free NPU memory, then cast.
|
||||
try:
|
||||
if param.device.type == 'npu':
|
||||
import torch_npu
|
||||
module = module.cpu()
|
||||
torch_npu.npu.synchronize()
|
||||
torch_npu.npu.empty_cache()
|
||||
return module.to(torch.float32)
|
||||
except Exception as exc:
|
||||
raise NPUCastError(f'Failed to cast {module.__class__.__name__} to fp32.') from exc
|
||||
|
||||
|
||||
_original_fsdp2_prepare_model = fsdp_utils.fsdp2_prepare_model
|
||||
|
||||
|
||||
@wraps(_original_fsdp2_prepare_model)
|
||||
def wrapped_fsdp2_prepare_model(
|
||||
accelerator: Accelerator,
|
||||
model: torch.nn.Module,
|
||||
):
|
||||
# Public utility entry used by some code paths before Accelerator.prepare.
|
||||
model = _cast_module_to_fp32_for_npu_if_needed(model, accelerator)
|
||||
return _original_fsdp2_prepare_model(accelerator, model)
|
||||
|
||||
|
||||
_original_prepare_fsdp2 = Accelerator._prepare_fsdp2
|
||||
|
||||
|
||||
@wraps(_original_prepare_fsdp2)
|
||||
def wrapped_prepare_fsdp2(
|
||||
self: Accelerator,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
# Accelerator.prepare may receive one or more modules directly; patch this
|
||||
# private entry too so all FSDP2 NPU preparation paths get the same fp32 cast.
|
||||
patched_args = [
|
||||
_cast_module_to_fp32_for_npu_if_needed(obj, self) if isinstance(obj, torch.nn.Module) else obj for obj in args
|
||||
]
|
||||
|
||||
return _original_prepare_fsdp2(self, *patched_args, **kwargs)
|
||||
|
||||
|
||||
_APPLIED = False
|
||||
|
||||
|
||||
def apply_patch() -> None:
|
||||
global _APPLIED
|
||||
if _APPLIED:
|
||||
return
|
||||
|
||||
fsdp_utils.fsdp2_prepare_model = wrapped_fsdp2_prepare_model
|
||||
Accelerator._prepare_fsdp2 = wrapped_prepare_fsdp2
|
||||
_APPLIED = True
|
||||
@@ -0,0 +1,209 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""NPU-only Megatron checkpoint compatibility helpers.
|
||||
|
||||
MindSpeed patches Megatron's distributed optimizer on NPU, but some Megatron-Core
|
||||
checkpoint formats still need the native Megatron param_state loaders.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from contextlib import contextmanager
|
||||
|
||||
from swift.utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _iter_optimizer_param_groups(optimizer):
|
||||
visited = set()
|
||||
|
||||
def visit(obj):
|
||||
if obj is None or id(obj) in visited:
|
||||
return
|
||||
visited.add(id(obj))
|
||||
|
||||
param_groups = getattr(obj, 'param_groups', None)
|
||||
if param_groups is not None:
|
||||
yield param_groups
|
||||
|
||||
inner_optimizer = getattr(obj, 'optimizer', None)
|
||||
if inner_optimizer is not obj:
|
||||
yield from visit(inner_optimizer)
|
||||
|
||||
for child in getattr(obj, 'chained_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
for child in getattr(obj, 'sub_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
|
||||
yield from visit(optimizer)
|
||||
|
||||
|
||||
def _step_to_int(step):
|
||||
if isinstance(step, torch.Tensor):
|
||||
if step.numel() != 1:
|
||||
raise RuntimeError(f'Optimizer step tensor must be scalar, got shape: {tuple(step.shape)}')
|
||||
return int(step.item())
|
||||
return int(step)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _canonicalize_optimizer_steps_for_checkpoint(optimizer):
|
||||
"""Normalize NPU scalar step tensors while Megatron builds optimizer checkpoint state.
|
||||
|
||||
Megatron-Core deduplicates param-group steps with set(). Equal NPU scalar
|
||||
tensors can still hash as distinct objects, so use their numeric value only
|
||||
while sharded_state_dict() is being built and restore the optimizer in place.
|
||||
"""
|
||||
saved_steps = []
|
||||
numeric_steps = set()
|
||||
for param_groups in _iter_optimizer_param_groups(optimizer):
|
||||
for param_group in param_groups:
|
||||
if len(param_group.get('params', [])) == 0 or 'step' not in param_group:
|
||||
continue
|
||||
step = param_group['step']
|
||||
numeric_step = _step_to_int(step)
|
||||
saved_steps.append((param_group, step))
|
||||
numeric_steps.add(numeric_step)
|
||||
|
||||
if len(numeric_steps) > 1:
|
||||
raise RuntimeError(f'Inconsistent optimizer steps before checkpoint save: {sorted(numeric_steps)}')
|
||||
|
||||
canonical_step = next(iter(numeric_steps), None)
|
||||
try:
|
||||
if canonical_step is not None:
|
||||
for param_group, _step in saved_steps:
|
||||
param_group['step'] = canonical_step
|
||||
if any(isinstance(step, torch.Tensor) for _param_group, step in saved_steps):
|
||||
logger.warning(f'Canonicalized optimizer param-group step to {canonical_step} for checkpoint save.')
|
||||
yield
|
||||
finally:
|
||||
for param_group, step in saved_steps:
|
||||
param_group['step'] = step
|
||||
|
||||
|
||||
def optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs):
|
||||
with _canonicalize_optimizer_steps_for_checkpoint(optimizer):
|
||||
return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs)
|
||||
|
||||
|
||||
def _iter_distributed_optimizers(optimizer):
|
||||
visited = set()
|
||||
|
||||
def visit(obj):
|
||||
if obj is None or id(obj) in visited:
|
||||
return
|
||||
visited.add(id(obj))
|
||||
|
||||
if hasattr(obj, 'load_parameter_state_from_dp_reshardable') or hasattr(
|
||||
obj, 'load_parameter_state_from_fully_reshardable'):
|
||||
yield obj
|
||||
return
|
||||
|
||||
for child in getattr(obj, 'chained_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
for child in getattr(obj, 'sub_optimizers', []) or []:
|
||||
yield from visit(child)
|
||||
|
||||
yield from visit(optimizer)
|
||||
|
||||
|
||||
def _has_mindspeed_patched_load_state_dict(distributed_optimizer):
|
||||
load_state_dict = getattr(type(distributed_optimizer), 'load_state_dict', None)
|
||||
return getattr(load_state_dict, '__module__', '').startswith('mindspeed.')
|
||||
|
||||
|
||||
_MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS = {
|
||||
'dp_reshardable': 'load_parameter_state_from_dp_reshardable',
|
||||
'fully_reshardable': 'load_parameter_state_from_fully_reshardable',
|
||||
}
|
||||
|
||||
|
||||
def _current_npu_device():
|
||||
if hasattr(torch, 'npu'):
|
||||
return torch.device('npu', torch.npu.current_device())
|
||||
return torch.cuda.current_device()
|
||||
|
||||
|
||||
def _restore_mindspeed_optimizer_step_tensors(optimizer):
|
||||
restored_count = 0
|
||||
for param_groups in _iter_optimizer_param_groups(optimizer):
|
||||
for param_group in param_groups:
|
||||
step = param_group.get('step')
|
||||
if isinstance(step, torch.Tensor):
|
||||
continue
|
||||
if isinstance(step, (int, float)):
|
||||
param_group['step'] = torch.tensor(int(step), dtype=torch.int64, device=_current_npu_device())
|
||||
restored_count += 1
|
||||
if restored_count:
|
||||
logger.warning(f'Restored {restored_count} MindSpeed optimizer param-group step values to NPU tensors.')
|
||||
|
||||
|
||||
def _split_chained_optimizer_state_dict(chained_optimizers, state_dict):
|
||||
if isinstance(state_dict, dict):
|
||||
state_dicts = [v for _k, v in sorted(state_dict.items())]
|
||||
else:
|
||||
state_dicts = list(state_dict)
|
||||
if len(chained_optimizers) != len(state_dicts):
|
||||
raise RuntimeError(
|
||||
f'Expected {len(chained_optimizers)} entries in optimizer state dict, but got {len(state_dicts)}.')
|
||||
return state_dicts
|
||||
|
||||
|
||||
def _load_chained_optimizer_state_dict(optimizer, state_dict):
|
||||
chained_optimizers = getattr(optimizer, 'chained_optimizers', None)
|
||||
if not chained_optimizers or len(chained_optimizers) <= 1:
|
||||
return False
|
||||
|
||||
state_dicts = _split_chained_optimizer_state_dict(chained_optimizers, state_dict)
|
||||
for child_optimizer, child_state_dict in zip(chained_optimizers, state_dicts):
|
||||
load_optimizer_state_dict(child_optimizer, child_state_dict)
|
||||
synchronize_steps = getattr(optimizer, '_synchronize_steps', None)
|
||||
if synchronize_steps is not None:
|
||||
synchronize_steps()
|
||||
return True
|
||||
|
||||
|
||||
def load_optimizer_state_dict(optimizer, state_dict):
|
||||
if _load_chained_optimizer_state_dict(optimizer, state_dict):
|
||||
return
|
||||
|
||||
distributed_optimizers = list(_iter_distributed_optimizers(optimizer))
|
||||
mindspeed_patched = any(
|
||||
_has_mindspeed_patched_load_state_dict(distributed_optimizer)
|
||||
for distributed_optimizer in distributed_optimizers)
|
||||
sharding_type = state_dict.get('param_state_sharding_type') if isinstance(state_dict, dict) else None
|
||||
native_loader_name = _MEGATRON_RESHARDABLE_PARAM_STATE_LOADERS.get(sharding_type)
|
||||
if native_loader_name is None:
|
||||
optimizer.load_state_dict(state_dict)
|
||||
if mindspeed_patched:
|
||||
_restore_mindspeed_optimizer_step_tensors(optimizer)
|
||||
return
|
||||
|
||||
if not mindspeed_patched:
|
||||
optimizer.load_state_dict(state_dict)
|
||||
return
|
||||
|
||||
if len(distributed_optimizers) != 1:
|
||||
raise RuntimeError(f'MindSpeed optimizer checkpoint compatibility supports exactly one distributed optimizer, '
|
||||
f'got {len(distributed_optimizers)}.')
|
||||
distributed_optimizer = distributed_optimizers[0]
|
||||
if not hasattr(distributed_optimizer, native_loader_name):
|
||||
raise RuntimeError(f'Distributed optimizer does not support sharding type {sharding_type}.')
|
||||
|
||||
state_dict_without_param_state = dict(state_dict)
|
||||
param_state = state_dict_without_param_state.pop('param_state', None)
|
||||
state_dict_without_param_state.pop('param_state_sharding_type', None)
|
||||
if param_state is None:
|
||||
raise RuntimeError(f'Optimizer checkpoint missing param_state for sharding type {sharding_type}.')
|
||||
|
||||
logger.warning(f'Loading optimizer param_state with ms-swift compatibility path because MindSpeed '
|
||||
f'DistributedOptimizer.load_state_dict does not support {sharding_type}.')
|
||||
# Let MindSpeed restore the generic optimizer state; load the missing
|
||||
# reshardable param_state with Megatron-Core's native implementation.
|
||||
optimizer.load_state_dict(state_dict_without_param_state)
|
||||
_restore_mindspeed_optimizer_step_tensors(optimizer)
|
||||
getattr(distributed_optimizer, native_loader_name)(param_state)
|
||||
|
||||
|
||||
__all__ = ['load_optimizer_state_dict', 'optimizer_sharded_state_dict']
|
||||
@@ -0,0 +1,52 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
_ORIGINAL_MINDSPEED_TE_CP_CLASS = None
|
||||
|
||||
|
||||
def patch_mindspeed_te_cp_implementation(megatron_args: dict[str, Any]) -> None:
|
||||
"""
|
||||
Route NPU CP to the legacy MindSpeed TE adaptor when the new strategy factory
|
||||
only supports kvallgather.
|
||||
"""
|
||||
# MindSpeed 0.15.3 replaced the TE context-parallel attention class with a
|
||||
# new implementation. That new class does not yet cover all CP algorithms,
|
||||
# so the default non-kvallgather path can fail during Megatron training.
|
||||
# For those algorithms, temporarily route TE attention back to the legacy
|
||||
# MindSpeedCPDotProductAttention adaptor. Once MindSpeed's new CP class has
|
||||
# feature parity, this compatibility patch can be removed.
|
||||
try:
|
||||
import mindspeed.te.pytorch.attention.dot_product_attention.dot_product_attention as ms_te_dpa
|
||||
from mindspeed.core.context_parallel.adaptor import MindSpeedCPDotProductAttention
|
||||
except ImportError as e:
|
||||
logger.warning(f'Failed to import MindSpeed CP modules before repatch: {e}')
|
||||
return
|
||||
|
||||
global _ORIGINAL_MINDSPEED_TE_CP_CLASS
|
||||
if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None:
|
||||
_ORIGINAL_MINDSPEED_TE_CP_CLASS = getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None)
|
||||
|
||||
if _ORIGINAL_MINDSPEED_TE_CP_CLASS is None:
|
||||
logger.warning('MindSpeedTEDotProductAttention is unavailable before repatch; skip CP workaround.')
|
||||
return
|
||||
|
||||
cp_algo = megatron_args.get('context_parallel_algo', 'megatron_cp_algo')
|
||||
use_legacy_cp_te = int(megatron_args.get('context_parallel_size', 1)) > 1 and cp_algo != 'kvallgather_cp_algo'
|
||||
target_cls = MindSpeedCPDotProductAttention if use_legacy_cp_te else _ORIGINAL_MINDSPEED_TE_CP_CLASS
|
||||
|
||||
if getattr(ms_te_dpa, 'MindSpeedTEDotProductAttention', None) is target_cls:
|
||||
return
|
||||
|
||||
ms_te_dpa.MindSpeedTEDotProductAttention = target_cls
|
||||
logger.info(
|
||||
'Patched MindSpeedTEDotProductAttention to %s for context_parallel_size=%s, context_parallel_algo=%s.',
|
||||
target_cls.__name__,
|
||||
megatron_args.get('context_parallel_size', 1),
|
||||
cp_algo,
|
||||
)
|
||||
@@ -0,0 +1,552 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch_npu
|
||||
from torch import nn
|
||||
from transformers.models.qwen2 import modeling_qwen2
|
||||
from transformers.models.qwen3 import modeling_qwen3
|
||||
from transformers.models.qwen3_moe import modeling_qwen3_moe
|
||||
from transformers.models.qwen3_vl_moe import modeling_qwen3_vl_moe
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
from .utils import apply_patch_map, import_optional_module
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Common NPU helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _resolve_unsqueeze_dim(position_ids=None, unsqueeze_dim=1):
|
||||
if isinstance(position_ids, int) and unsqueeze_dim == 1:
|
||||
return position_ids
|
||||
return unsqueeze_dim
|
||||
|
||||
|
||||
def _get_hidden_size(module, hidden_states: torch.Tensor) -> int:
|
||||
return getattr(module, 'hidden_size', getattr(module, 'hidden_dim', hidden_states.shape[-1]))
|
||||
|
||||
|
||||
class NpuRMSNorm(nn.Module):
|
||||
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.variance_epsilon)[0]
|
||||
|
||||
def extra_repr(self):
|
||||
return f'{tuple(self.weight.shape)}, eps={self.variance_epsilon}'
|
||||
|
||||
|
||||
class NpuGmmFunction(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, weight, group_list, split_size):
|
||||
ctx.save_for_backward(x, weight)
|
||||
ctx.group_list = group_list
|
||||
ctx.split_size = split_size
|
||||
|
||||
outputs = torch_npu.npu_grouped_matmul([x], [weight], group_list=group_list, group_type=0, split_item=2)
|
||||
return outputs[0]
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_outputs):
|
||||
x, weight = ctx.saved_tensors
|
||||
group_list = ctx.group_list
|
||||
wt = weight.permute(0, 2, 1)
|
||||
xt = x.permute(1, 0)
|
||||
dx = torch_npu.npu_grouped_matmul([grad_outputs], [wt], group_list=group_list, group_type=0, split_item=2)
|
||||
split_size = ctx.split_size
|
||||
xt_list = torch.split(xt, split_size, dim=1)
|
||||
grad_outputs_list = torch.split(grad_outputs, split_size, dim=0)
|
||||
with torch.npu.amp.autocast(enabled=False):
|
||||
dw = torch.stack([torch.matmul(xt_list[i], grad_outputs_list[i]) for i in range(len(xt_list))])
|
||||
|
||||
return dx[0], dw, None, None
|
||||
|
||||
|
||||
class GmmFunction(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, weight, group_list):
|
||||
ctx.save_for_backward(x, weight)
|
||||
ctx.group_list = group_list
|
||||
|
||||
fwd_output = torch_npu.npu_grouped_matmul([x], [weight],
|
||||
bias=None,
|
||||
group_list=group_list,
|
||||
split_item=2,
|
||||
group_type=0,
|
||||
group_list_type=1)[0]
|
||||
return fwd_output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
input_tensor, weight = ctx.saved_tensors
|
||||
group_list = ctx.group_list
|
||||
|
||||
weight = torch.transpose(weight, 1, 2)
|
||||
grad_input = torch_npu.npu_grouped_matmul([grad_output], [weight],
|
||||
bias=None,
|
||||
group_list=group_list,
|
||||
split_item=2,
|
||||
group_type=0,
|
||||
group_list_type=1)[0]
|
||||
grad_weight = torch_npu.npu_grouped_matmul(
|
||||
[input_tensor.T],
|
||||
[grad_output],
|
||||
bias=None,
|
||||
group_list=group_list,
|
||||
split_item=3,
|
||||
group_type=2,
|
||||
group_list_type=1,
|
||||
)[0]
|
||||
return grad_input, grad_weight, None
|
||||
|
||||
|
||||
def _normalize_packed_expert_weights(module, input_dtype: torch.dtype, hidden_dim: int):
|
||||
gate_up_proj = module.gate_up_proj.to(input_dtype)
|
||||
down_proj = module.down_proj.to(input_dtype)
|
||||
|
||||
if gate_up_proj.shape[1] == hidden_dim:
|
||||
gate_up_weight = gate_up_proj
|
||||
elif gate_up_proj.shape[2] == hidden_dim:
|
||||
gate_up_weight = gate_up_proj.transpose(1, 2)
|
||||
else:
|
||||
raise RuntimeError(f'Unsupported gate_up_proj shape for NPU MoE patch: {tuple(gate_up_proj.shape)}.')
|
||||
|
||||
if down_proj.shape[2] == hidden_dim:
|
||||
down_weight = down_proj
|
||||
elif down_proj.shape[1] == hidden_dim:
|
||||
down_weight = down_proj.transpose(1, 2)
|
||||
else:
|
||||
raise RuntimeError(f'Unsupported down_proj shape for NPU MoE patch: {tuple(down_proj.shape)}.')
|
||||
|
||||
return gate_up_weight, down_weight
|
||||
|
||||
|
||||
def npu_packed_moe_experts_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_indices_or_routing_weights: torch.Tensor,
|
||||
routing_weights_or_router_indices: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
if router_indices_or_routing_weights.dtype in {torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8}:
|
||||
router_indices = router_indices_or_routing_weights
|
||||
routing_weights = routing_weights_or_router_indices
|
||||
else:
|
||||
routing_weights = router_indices_or_routing_weights
|
||||
router_indices = routing_weights_or_router_indices
|
||||
|
||||
output_shape = hidden_states.shape
|
||||
hidden_dim = output_shape[-1]
|
||||
hidden_states = hidden_states.reshape(-1, hidden_dim)
|
||||
|
||||
if routing_weights.shape != router_indices.shape:
|
||||
routing_weights = torch.gather(routing_weights, dim=-1, index=router_indices.to(torch.long))
|
||||
routing_weights = routing_weights.to(hidden_states.dtype)
|
||||
router_indices = router_indices.to(torch.int32)
|
||||
|
||||
permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(hidden_states, router_indices)
|
||||
tokens_per_expert = torch.histc(
|
||||
router_indices.to(torch.float), bins=self.num_experts, min=0, max=self.num_experts).to(torch.int64)
|
||||
gate_up_weight, down_weight = _normalize_packed_expert_weights(self, hidden_states.dtype, hidden_dim)
|
||||
|
||||
intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_weight, tokens_per_expert)
|
||||
intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
|
||||
output = GmmFunction.apply(intermediate_activations, down_weight, tokens_per_expert)
|
||||
next_states = torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=routing_weights)
|
||||
return next_states.view(*output_shape)
|
||||
|
||||
|
||||
def _topk_from_router_logits(module, hidden_states: torch.Tensor, router_logits: torch.Tensor):
|
||||
routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
|
||||
routing_weights, router_indices = torch.topk(routing_weights, module.top_k, dim=-1)
|
||||
if getattr(module, 'norm_topk_prob', True):
|
||||
routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
|
||||
routing_weights = routing_weights.to(hidden_states.dtype)
|
||||
return routing_weights, router_indices
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen2/Qwen3 dense patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def npu_apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
"""Applies Rotary Position Embedding to the query and key tensors."""
|
||||
unsqueeze_dim = _resolve_unsqueeze_dim(position_ids, unsqueeze_dim)
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
q_embed = torch_npu.npu_rotary_mul(q, cos, sin)
|
||||
k_embed = torch_npu.npu_rotary_mul(k, cos, sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
def npu_swiglu_forward(self, hidden_state):
|
||||
return self.down_proj(
|
||||
torch_npu.npu_swiglu(torch.cat((self.gate_proj(hidden_state), self.up_proj(hidden_state)), dim=-1), dim=-1))
|
||||
|
||||
|
||||
QWEN2_PATCHES = {
|
||||
'Qwen2RMSNorm': NpuRMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
|
||||
'Qwen2MLP.forward': npu_swiglu_forward,
|
||||
}
|
||||
|
||||
QWEN3_PATCHES = {
|
||||
'Qwen3RMSNorm': NpuRMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
|
||||
'Qwen3MLP.forward': npu_swiglu_forward,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen3.5 dense patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class NpuQwen3_5RMSNorm(nn.Module):
|
||||
|
||||
def __init__(self, dim, eps=1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.zeros(dim))
|
||||
|
||||
def forward(self, x):
|
||||
scale = (1.0 + self.weight).to(dtype=x.dtype)
|
||||
return torch_npu.npu_rms_norm(x, scale, epsilon=self.eps)[0]
|
||||
|
||||
def extra_repr(self):
|
||||
return f'{tuple(self.weight.shape)}, eps={self.eps}'
|
||||
|
||||
|
||||
def npu_apply_rotary_pos_emb_qwen3_5(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
unsqueeze_dim = _resolve_unsqueeze_dim(position_ids, unsqueeze_dim)
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
|
||||
rotary_dim = cos.shape[-1]
|
||||
q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
|
||||
k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
|
||||
|
||||
q_rot = torch_npu.npu_rotary_mul(q_rot, cos, sin)
|
||||
k_rot = torch_npu.npu_rotary_mul(k_rot, cos, sin)
|
||||
|
||||
q_embed = torch.cat([q_rot, q_pass], dim=-1)
|
||||
k_embed = torch.cat([k_rot, k_pass], dim=-1)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
_MISSING = object()
|
||||
_TRANSFORMERS_FLA_PROBE_MODULES = ('transformers.utils', 'transformers.utils.import_utils')
|
||||
|
||||
|
||||
def _patch_transformers_flash_linear_attention_available(available: bool) -> dict[str, object]:
|
||||
|
||||
def _is_flash_linear_attention_available() -> bool:
|
||||
return available
|
||||
|
||||
originals = {}
|
||||
for module_name in _TRANSFORMERS_FLA_PROBE_MODULES:
|
||||
module = import_optional_module(module_name)
|
||||
if module is None:
|
||||
continue
|
||||
originals[module_name] = getattr(module, 'is_flash_linear_attention_available', _MISSING)
|
||||
setattr(module, 'is_flash_linear_attention_available', _is_flash_linear_attention_available)
|
||||
return originals
|
||||
|
||||
|
||||
def _restore_transformers_flash_linear_attention_available(originals: dict[str, object]) -> None:
|
||||
for module_name, original in originals.items():
|
||||
module = import_optional_module(module_name)
|
||||
if module is None:
|
||||
continue
|
||||
if original is _MISSING:
|
||||
delattr(module, 'is_flash_linear_attention_available')
|
||||
else:
|
||||
setattr(module, 'is_flash_linear_attention_available', original)
|
||||
|
||||
|
||||
def patch_qwen3_5_chunk_gated_delta_rule_with_mindspeed() -> None:
|
||||
try:
|
||||
from ..chunk_gated_delta_rule import chunk_gated_delta_rule
|
||||
except ImportError as exc:
|
||||
logger.warning('Failed to import embedded MindSpeed chunk_gated_delta_rule: %s', exc)
|
||||
return
|
||||
|
||||
patched_modules = []
|
||||
for module_name in ('transformers.models.qwen3_5.modeling_qwen3_5',
|
||||
'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe'):
|
||||
module = import_optional_module(module_name)
|
||||
if module is None:
|
||||
continue
|
||||
|
||||
setattr(module, 'is_flash_linear_attention_available', lambda: True)
|
||||
setattr(module, 'is_fast_path_available', True)
|
||||
# FLA's fused RMSNormGated initializes with torch.cuda.current_device(),
|
||||
# so keep the native Qwen3.5 torch implementation on NPU.
|
||||
setattr(module, 'FusedRMSNormGated', None)
|
||||
setattr(module, 'chunk_gated_delta_rule', chunk_gated_delta_rule)
|
||||
patched_modules.append(module_name)
|
||||
|
||||
if patched_modules:
|
||||
logger.info('Patched Qwen3.5 chunk_gated_delta_rule to embedded MindSpeed implementation: %s.',
|
||||
', '.join(patched_modules))
|
||||
|
||||
|
||||
QWEN3_5_PATCHES = {
|
||||
'Qwen3_5RMSNorm': NpuQwen3_5RMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb_qwen3_5,
|
||||
'Qwen3_5MLP.forward': npu_swiglu_forward,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen3-MoE patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _qwen3_moe_forward_transformers_457(self, hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
|
||||
if getattr(self, 'norm_topk_prob', False):
|
||||
routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
|
||||
routing_weights = routing_weights.to(hidden_states.dtype)
|
||||
|
||||
expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
|
||||
|
||||
input_dtype = hidden_states.dtype
|
||||
up_weight_list = [expert.up_proj.weight.t().to(input_dtype) for expert in self.experts]
|
||||
gate_weight_list = [expert.gate_proj.weight.t().to(input_dtype) for expert in self.experts]
|
||||
down_weight_list = [expert.down_proj.weight.t().to(input_dtype) for expert in self.experts]
|
||||
w1 = torch.stack(up_weight_list)
|
||||
w2 = torch.stack(gate_weight_list)
|
||||
w3 = torch.stack(down_weight_list)
|
||||
|
||||
routing_map = selected_experts
|
||||
flatten_indices = routing_map.view(-1)
|
||||
sorted_indices = torch.sort(flatten_indices.float(), stable=True)[1]
|
||||
permuted_tokens = hidden_states.index_select(0, sorted_indices // self.top_k)
|
||||
|
||||
tokens_per_experts = torch.sum(expert_mask, dim=(1, 2))
|
||||
group_list = torch.cumsum(tokens_per_experts, dim=0)
|
||||
|
||||
cpu_group_list = group_list.to('cpu', non_blocking=False)
|
||||
cpu_group_list = [0] + cpu_group_list.tolist()
|
||||
split_size = [cpu_group_list[i + 1] - cpu_group_list[i] for i in range(len(cpu_group_list) - 1)]
|
||||
|
||||
up_res = NpuGmmFunction.apply(permuted_tokens, w1, group_list, split_size)
|
||||
gate_res = NpuGmmFunction.apply(permuted_tokens, w2, group_list, split_size)
|
||||
act_res = torch_npu.npu_swiglu(torch.cat([gate_res, up_res], dim=-1))
|
||||
down_res = NpuGmmFunction.apply(act_res, w3, group_list, split_size)
|
||||
|
||||
num_unpermuted_tokens = routing_weights.numel()
|
||||
unpermuted_tokens = torch.zeros(
|
||||
[num_unpermuted_tokens, down_res.shape[-1]],
|
||||
dtype=down_res.dtype,
|
||||
device=down_res.device,
|
||||
)
|
||||
unpermuted_tokens.index_copy_(0, sorted_indices, down_res)
|
||||
unpermuted_tokens = unpermuted_tokens.reshape(-1, self.top_k, down_res.size(-1))
|
||||
unpermuted_tokens = unpermuted_tokens * routing_weights.unsqueeze(-1)
|
||||
final_hidden_states = unpermuted_tokens.sum(dim=1).to(hidden_states.dtype)
|
||||
final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
|
||||
|
||||
return final_hidden_states, router_logits
|
||||
|
||||
|
||||
def _qwen3_moe_forward_transformers_5(self, hidden_states: torch.Tensor, routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
final_hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
|
||||
return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
|
||||
|
||||
|
||||
def npu_qwen3_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_dim = hidden_states.shape[-1]
|
||||
gate_output = self.gate(hidden_states.view(-1, hidden_dim))
|
||||
|
||||
if isinstance(gate_output, tuple):
|
||||
# Transformers 5.x: gate is a router module and returns
|
||||
# (router_logits, routing_weights, selected_experts).
|
||||
_, routing_weights, selected_experts = gate_output
|
||||
return _qwen3_moe_forward_transformers_5(self, hidden_states, routing_weights, selected_experts)
|
||||
|
||||
# Transformers 4.57.x: gate is nn.Linear and returns router logits.
|
||||
return _qwen3_moe_forward_transformers_457(self, hidden_states, gate_output)
|
||||
|
||||
|
||||
QWEN3_MOE_PATCHES = {
|
||||
'Qwen3MoeRMSNorm': NpuRMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
|
||||
'Qwen3MoeSparseMoeBlock.forward': npu_qwen3_moe_sparse_block_forward,
|
||||
}
|
||||
|
||||
QWEN3_MOE_TRANSFORMERS_5_PATCHES = {
|
||||
'Qwen3MoeExperts.forward': npu_packed_moe_experts_forward,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen3-VL-MoE patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _qwen3_vl_moe_forward_transformers_457(self, hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor) -> torch.Tensor:
|
||||
batch_size = hidden_states.shape[0]
|
||||
hidden_size = _get_hidden_size(self, hidden_states)
|
||||
hidden_states = hidden_states.reshape(-1, hidden_size)
|
||||
|
||||
routing_weights, router_indices = _topk_from_router_logits(self, hidden_states, router_logits)
|
||||
hidden_states = hidden_states.reshape(batch_size, -1, hidden_size)
|
||||
return self.experts(hidden_states, routing_weights, router_indices)
|
||||
|
||||
|
||||
def _qwen3_vl_moe_forward_transformers_5(self, hidden_states: torch.Tensor, routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, sequence_length, hidden_size = hidden_states.shape
|
||||
hidden_states = hidden_states.reshape(-1, hidden_size)
|
||||
routed_out = self.experts(hidden_states, selected_experts, routing_weights)
|
||||
return routed_out.reshape(batch_size, sequence_length, hidden_size)
|
||||
|
||||
|
||||
def npu_qwen3_vl_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_size = _get_hidden_size(self, hidden_states)
|
||||
gate_output = self.gate(hidden_states.reshape(-1, hidden_size))
|
||||
|
||||
if isinstance(gate_output, tuple):
|
||||
# Transformers 5.x: gate is a router module and returns
|
||||
# (router_logits, routing_weights, selected_experts).
|
||||
_, routing_weights, selected_experts = gate_output
|
||||
return _qwen3_vl_moe_forward_transformers_5(self, hidden_states, routing_weights, selected_experts)
|
||||
|
||||
# Transformers 4.57.x: gate is nn.Linear and experts use the old
|
||||
# (hidden_states, routing_weights, router_indices) call order.
|
||||
return _qwen3_vl_moe_forward_transformers_457(self, hidden_states, gate_output)
|
||||
|
||||
|
||||
QWEN3_VL_MOE_PATCHES = {
|
||||
'Qwen3VLMoeTextExperts.forward': npu_packed_moe_experts_forward,
|
||||
'Qwen3VLMoeTextSparseMoeBlock.forward': npu_qwen3_vl_moe_sparse_block_forward,
|
||||
'Qwen3VLMoeTextRMSNorm': NpuRMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen3.5-MoE patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _add_shared_expert(self, hidden_states: torch.Tensor, expert_output: torch.Tensor) -> torch.Tensor:
|
||||
if not (hasattr(self, 'shared_expert') and hasattr(self, 'shared_expert_gate')):
|
||||
return expert_output
|
||||
|
||||
shared_expert_output = self.shared_expert(hidden_states)
|
||||
shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
|
||||
return expert_output + shared_expert_output
|
||||
|
||||
|
||||
def _qwen3_5_moe_forward_transformers_5(self, hidden_states: torch.Tensor, routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
expert_output = self.experts(hidden_states, selected_experts, routing_weights)
|
||||
expert_output = _add_shared_expert(self, hidden_states, expert_output)
|
||||
return expert_output.reshape(batch_size, sequence_length, hidden_dim)
|
||||
|
||||
|
||||
def _qwen3_5_moe_forward_linear_gate(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, sequence_length, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
routing_weights, selected_experts = _topk_from_router_logits(self, hidden_states, router_logits)
|
||||
expert_output = self.experts(hidden_states, selected_experts, routing_weights)
|
||||
expert_output = _add_shared_expert(self, hidden_states, expert_output)
|
||||
return expert_output.reshape(batch_size, sequence_length, hidden_dim)
|
||||
|
||||
|
||||
def npu_qwen3_5_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_dim = hidden_states.shape[-1]
|
||||
gate_output = self.gate(hidden_states.view(-1, hidden_dim))
|
||||
|
||||
if isinstance(gate_output, tuple):
|
||||
# Transformers 5.x: Qwen3.5-MoE has packed experts plus shared expert.
|
||||
_, routing_weights, selected_experts = gate_output
|
||||
return _qwen3_5_moe_forward_transformers_5(self, hidden_states, routing_weights, selected_experts)
|
||||
|
||||
return _qwen3_5_moe_forward_linear_gate(self, hidden_states, gate_output)
|
||||
|
||||
|
||||
QWEN3_5_MOE_PATCHES = {
|
||||
'Qwen3_5MoeRMSNorm': NpuQwen3_5RMSNorm,
|
||||
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb_qwen3_5,
|
||||
'Qwen3_5MoeMLP.forward': npu_swiglu_forward,
|
||||
'Qwen3_5MoeExperts.forward': npu_packed_moe_experts_forward,
|
||||
'Qwen3_5MoeSparseMoeBlock.forward': npu_qwen3_5_moe_sparse_block_forward,
|
||||
}
|
||||
|
||||
QWEN3_5_MOE_OPTIONAL_PATCHES = {}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch table and apply entry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_patch_map(root, patches: dict[str, object], optional_patches: dict[str, object] | None = None):
|
||||
patch_map = dict(patches)
|
||||
for path, value in (optional_patches or {}).items():
|
||||
current = root
|
||||
for part in path.split('.'):
|
||||
if not hasattr(current, part):
|
||||
break
|
||||
current = getattr(current, part)
|
||||
else:
|
||||
patch_map[path] = value
|
||||
return patch_map
|
||||
|
||||
|
||||
_APPLIED = False
|
||||
|
||||
|
||||
def apply_patch() -> None:
|
||||
global _APPLIED
|
||||
if _APPLIED:
|
||||
return
|
||||
|
||||
patch_groups = [
|
||||
('qwen2', modeling_qwen2, QWEN2_PATCHES, {}),
|
||||
('qwen3', modeling_qwen3, QWEN3_PATCHES, {}),
|
||||
('qwen3_moe', modeling_qwen3_moe, QWEN3_MOE_PATCHES, QWEN3_MOE_TRANSFORMERS_5_PATCHES),
|
||||
('qwen3_vl_moe', modeling_qwen3_vl_moe, QWEN3_VL_MOE_PATCHES, {}),
|
||||
]
|
||||
|
||||
# Qwen3.5 GDN is patched to swift's embedded MindSpeed implementation below.
|
||||
# Skip Transformers' import-time FLA probe so FLA is not a hard dependency.
|
||||
fla_probe_originals = _patch_transformers_flash_linear_attention_available(False)
|
||||
try:
|
||||
modeling_qwen3_5 = import_optional_module('transformers.models.qwen3_5.modeling_qwen3_5')
|
||||
modeling_qwen3_5_moe = import_optional_module('transformers.models.qwen3_5_moe.modeling_qwen3_5_moe')
|
||||
finally:
|
||||
_restore_transformers_flash_linear_attention_available(fla_probe_originals)
|
||||
if modeling_qwen3_5 is not None:
|
||||
patch_qwen3_5_chunk_gated_delta_rule_with_mindspeed()
|
||||
|
||||
if modeling_qwen3_5 is not None:
|
||||
patch_groups.append(('qwen3_5', modeling_qwen3_5, QWEN3_5_PATCHES, {}))
|
||||
|
||||
if modeling_qwen3_5_moe is not None:
|
||||
patch_groups.append(('qwen3_5_moe', modeling_qwen3_5_moe, QWEN3_5_MOE_PATCHES, QWEN3_5_MOE_OPTIONAL_PATCHES))
|
||||
|
||||
for _group_name, module, patches, optional_patches in patch_groups:
|
||||
apply_patch_map(module, _build_patch_map(module, patches, optional_patches))
|
||||
|
||||
_APPLIED = True
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def import_optional_module(module_name: str) -> Any | None:
|
||||
try:
|
||||
return importlib.import_module(module_name)
|
||||
except ImportError as exc:
|
||||
logger.debug('Failed to import optional module %s: %s', module_name, exc)
|
||||
return None
|
||||
|
||||
|
||||
def apply_patch_map(root: Any, patch_map: dict[str, Any]) -> None:
|
||||
for path, value in patch_map.items():
|
||||
current = root
|
||||
parts = path.split('.')
|
||||
for part in parts[:-1]:
|
||||
current = getattr(current, part)
|
||||
setattr(current, parts[-1], value)
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""Facade for SWIFT's vLLM-Ascend NPU compatibility patches.
|
||||
|
||||
Keep this file thin. The real patches are split by responsibility:
|
||||
|
||||
* ``vllm_ascend_moe``: MoE routing and GRPO weight-sync layout handling.
|
||||
* ``vllm_ascend_memory``: small torch-npu/vLLM-Ascend memory API compatibility.
|
||||
Callers should import from this module so the public entrypoints stay stable,
|
||||
while reviewers can audit each patch family in its own file. The caller is
|
||||
still responsible for guarding these entrypoints with an NPU/device check.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from swift.model.npu_patch.vllm_ascend_memory import patch_vllm_ascend_memory_runtime
|
||||
from swift.model.npu_patch.vllm_ascend_moe import (patch_vllm_ascend_moe_expert_weight_loader,
|
||||
patch_vllm_ascend_moe_runtime, should_skip_vllm_ascend_moe_post_load,
|
||||
use_vllm_ascend_moe_preprocessed_weight)
|
||||
from swift.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _patch_flash_attn_optional_import() -> None:
|
||||
"""Clear a stub ``flash_attn`` module that can block optional imports.
|
||||
|
||||
Some stacks insert a non-package ``flash_attn`` placeholder into
|
||||
``sys.modules``. vLLM import paths then treat it as the real package and
|
||||
fail on submodule imports. Removing the placeholder lets normal optional
|
||||
dependency checks proceed.
|
||||
"""
|
||||
module = sys.modules.get('flash_attn')
|
||||
if module is None or hasattr(module, '__path__'):
|
||||
return
|
||||
for module_name in list(sys.modules):
|
||||
if module_name == 'flash_attn' or module_name.startswith('flash_attn.'):
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
|
||||
def patch_vllm_ascend_runtime(*, colocate: bool = False) -> None:
|
||||
"""Apply vLLM-Ascend patches needed by SWIFT NPU rollout.
|
||||
|
||||
``colocate=False`` covers patches that are also safe for standalone
|
||||
vLLM-Ascend server/native inference, such as optional import cleanup, MoE
|
||||
routing, and ``mem_get_info`` binding compatibility.
|
||||
|
||||
``colocate`` is kept in the public signature for callers that share this
|
||||
entrypoint between server and colocate modes. Process-group creation is
|
||||
left to upstream vLLM/vLLM-Ascend; SWIFT only keeps the narrow runtime
|
||||
compatibility patches below.
|
||||
"""
|
||||
_patch_flash_attn_optional_import()
|
||||
patch_vllm_ascend_moe_runtime()
|
||||
patch_vllm_ascend_memory_runtime()
|
||||
|
||||
|
||||
__all__ = [
|
||||
'patch_vllm_ascend_moe_expert_weight_loader',
|
||||
'patch_vllm_ascend_runtime',
|
||||
'should_skip_vllm_ascend_moe_post_load',
|
||||
'use_vllm_ascend_moe_preprocessed_weight',
|
||||
]
|
||||
@@ -0,0 +1,91 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""Small vLLM-Ascend memory API compatibility patches.
|
||||
|
||||
This module intentionally avoids colocate memory-policy changes. It only
|
||||
normalizes API differences that are safe for both standalone vLLM-Ascend
|
||||
inference and SWIFT GRPO rollout.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
|
||||
_ORIGIN_TORCH_NPU_MEM_GET_INFO = None
|
||||
_BOUND_TORCH_NPU_MEM_GET_INFO_DEVICE = None
|
||||
|
||||
|
||||
def _patch_vllm_ascend_mem_get_info() -> None:
|
||||
"""Patch ``NPUPlatform.mem_get_info`` for torch-npu binding differences.
|
||||
|
||||
vLLM-Ascend calls ``current_platform.mem_get_info(device)`` during worker
|
||||
initialization. Without this wrapper, some versions expose
|
||||
``NPUPlatform.mem_get_info`` in a way that gets Python method binding plus
|
||||
the explicit device argument at the same time, producing:
|
||||
|
||||
TypeError: mem_get_info() got multiple values for argument 'device'
|
||||
|
||||
Defining a classmethod here gives vLLM-Ascend one stable call surface. It
|
||||
keeps the device-aware torch-npu query when available and falls back to the
|
||||
no-argument query only when torch-npu rejects the keyword. This does not
|
||||
change memory profiling policy.
|
||||
"""
|
||||
try:
|
||||
from vllm_ascend.platform import NPUPlatform
|
||||
except (ImportError, AttributeError):
|
||||
return
|
||||
if getattr(NPUPlatform, '_swift_mem_get_info_patched', False):
|
||||
return
|
||||
|
||||
@classmethod
|
||||
def mem_get_info(cls, device=None):
|
||||
if device is None:
|
||||
return torch.npu.mem_get_info()
|
||||
try:
|
||||
return torch.npu.mem_get_info(device=device)
|
||||
except TypeError:
|
||||
return torch.npu.mem_get_info()
|
||||
|
||||
NPUPlatform.mem_get_info = mem_get_info
|
||||
NPUPlatform._swift_mem_get_info_patched = True
|
||||
|
||||
|
||||
def patch_vllm_ascend_memory_runtime() -> None:
|
||||
"""Apply memory patches that do not depend on colocated training."""
|
||||
_patch_vllm_ascend_mem_get_info()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def vllm_ascend_mem_get_info_context(vllm_device: str):
|
||||
"""Bind bare ``torch.npu.mem_get_info()`` calls to vLLM's device.
|
||||
|
||||
Most vLLM memory accounting goes through ``NPUPlatform.mem_get_info`` and is
|
||||
handled by ``patch_vllm_ascend_memory_runtime`` above. Some vLLM-Ascend
|
||||
paths still call ``torch.npu.mem_get_info()`` directly, or assign it to
|
||||
``torch.cuda.mem_get_info`` for CUDA-compatible worker code.
|
||||
|
||||
Keep this binding for the process lifetime after the context exits. vLLM
|
||||
sleep/wake paths can call bare ``torch.npu.mem_get_info()`` after engine
|
||||
construction, so restoring here would regress the original behavior in
|
||||
``swift.infer_engine.utils.patch_npu_vllm``. Re-entering with another device
|
||||
rebinds from the original function instead of stacking nested partials.
|
||||
"""
|
||||
global _ORIGIN_TORCH_NPU_MEM_GET_INFO, _BOUND_TORCH_NPU_MEM_GET_INFO_DEVICE
|
||||
|
||||
if (_ORIGIN_TORCH_NPU_MEM_GET_INFO is None
|
||||
or getattr(torch.npu.mem_get_info, '_swift_bound_mem_get_info_device', None) is None):
|
||||
_ORIGIN_TORCH_NPU_MEM_GET_INFO = torch.npu.mem_get_info
|
||||
|
||||
if _BOUND_TORCH_NPU_MEM_GET_INFO_DEVICE != vllm_device:
|
||||
mem_get_info = partial(_ORIGIN_TORCH_NPU_MEM_GET_INFO, device=vllm_device)
|
||||
mem_get_info._swift_bound_mem_get_info_device = vllm_device
|
||||
torch.npu.mem_get_info = mem_get_info
|
||||
_BOUND_TORCH_NPU_MEM_GET_INFO_DEVICE = vllm_device
|
||||
|
||||
yield
|
||||
|
||||
|
||||
__all__ = [
|
||||
'patch_vllm_ascend_memory_runtime',
|
||||
'vllm_ascend_mem_get_info_context',
|
||||
]
|
||||
@@ -0,0 +1,394 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""vLLM-Ascend MoE patches used by SWIFT NPU rollout.
|
||||
|
||||
There are two independent responsibilities in this file:
|
||||
|
||||
* runtime routing: avoid the unstable custom non-quantized MoE routing op on
|
||||
stacks where vLLM-Ascend still dispatches that branch to
|
||||
``aclnnMoeInitRoutingCustom``;
|
||||
* weight sync: adapt 2D HF/Megatron MoE expert weights to the already-processed
|
||||
3D vLLM-Ascend expert parameter layout during GRPO colocate updates.
|
||||
|
||||
Both patches are guarded by vLLM-Ascend implementation checks and only touch the
|
||||
specific MoE paths they need.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import torch
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
_VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR = '_swift_vllm_ascend_moe_weight_sync_layout'
|
||||
_VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR = '_swift_vllm_ascend_moe_skip_post_load'
|
||||
_VLLM_ASCEND_MOE_PROCESSED_LAYOUT = 'megatron_processed'
|
||||
_VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT = 'fsdp2_preprocessed'
|
||||
_QWEN_MOE_MODEL_TYPES = {'qwen3_moe', 'qwen3_5_moe'}
|
||||
|
||||
|
||||
def _patch_vllm_ascend_device_op_nonquant_routing() -> None:
|
||||
"""Use the stable torch-npu routing op for non-quantized MoE when needed.
|
||||
|
||||
Some released vLLM-Ascend versions route the non-quantized MoE case
|
||||
(``scale is None`` and ``quant_mode == -1``) through
|
||||
``npu_moe_init_routing_custom`` / ``aclnnMoeInitRoutingCustom``, which is
|
||||
not stable for the parameter combination used by Qwen-style MoE rollout.
|
||||
|
||||
This is intentionally gated by implementation detection instead of a fixed
|
||||
version threshold: source builds or future/backported versions may already
|
||||
dispatch the non-quantized path to ``torch_npu.npu_moe_init_routing_v2``.
|
||||
When that fixed branch is present, skip patching and keep the upstream
|
||||
implementation intact.
|
||||
|
||||
Do not probe the custom op by calling it first. On Ascend, a missing custom
|
||||
binary can be reported asynchronously: even if Python catches the immediate
|
||||
RuntimeError and falls back, the failed launch can poison the stream and hang
|
||||
later at an unrelated event synchronization. Therefore, when source
|
||||
inspection shows that the non-quantized branch still routes to the custom op,
|
||||
dispatch that branch directly to ``torch_npu.npu_moe_init_routing_v2``.
|
||||
"""
|
||||
try:
|
||||
import torch_npu
|
||||
from vllm_ascend.device import device_op
|
||||
except (ImportError, AttributeError):
|
||||
return
|
||||
|
||||
adaptor_cls = getattr(device_op, 'BaseDeviceAdaptor', None)
|
||||
if adaptor_cls is None:
|
||||
return
|
||||
origin_routing = getattr(adaptor_cls, 'npu_moe_init_routing', None)
|
||||
if origin_routing is None or getattr(origin_routing, '_swift_nonquant_routing_patched', False):
|
||||
return
|
||||
try:
|
||||
origin_source = inspect.getsource(origin_routing)
|
||||
except (OSError, TypeError):
|
||||
origin_source = ''
|
||||
if 'npu_moe_init_routing_v2' in origin_source and 'quant_mode == -1' in origin_source:
|
||||
return
|
||||
origin_signature = inspect.signature(origin_routing)
|
||||
routing_defaults = {
|
||||
'scale': None,
|
||||
'active_num': None,
|
||||
'expert_num': None,
|
||||
'expert_tokens_num_type': 1,
|
||||
'expert_tokens_num_flag': True,
|
||||
'active_expert_range': None,
|
||||
'quant_mode': -1,
|
||||
}
|
||||
missing_params = set(routing_defaults).difference(origin_signature.parameters)
|
||||
if missing_params:
|
||||
raise RuntimeError('Unsupported vLLM-Ascend npu_moe_init_routing signature: '
|
||||
f'signature={origin_signature}, missing={sorted(missing_params)}.')
|
||||
|
||||
def is_nonquant_routing(routing_kwargs) -> bool:
|
||||
return routing_kwargs['scale'] is None and routing_kwargs['quant_mode'] == -1
|
||||
|
||||
def npu_moe_init_routing_v2(hidden_states, topk_ids, routing_kwargs):
|
||||
active_num = routing_kwargs['active_num']
|
||||
expert_num = routing_kwargs['expert_num']
|
||||
active_expert_range = routing_kwargs['active_expert_range']
|
||||
return torch_npu.npu_moe_init_routing_v2(
|
||||
hidden_states,
|
||||
topk_ids,
|
||||
scale=None,
|
||||
offset=None,
|
||||
active_num=0 if active_num is None else active_num,
|
||||
expert_capacity=-1,
|
||||
expert_num=expert_num,
|
||||
drop_pad_mode=0,
|
||||
expert_tokens_num_type=routing_kwargs['expert_tokens_num_type'],
|
||||
expert_tokens_num_flag=routing_kwargs['expert_tokens_num_flag'],
|
||||
active_expert_range=[0, expert_num] if active_expert_range is None else active_expert_range,
|
||||
quant_mode=routing_kwargs['quant_mode'],
|
||||
row_idx_type=0,
|
||||
)
|
||||
|
||||
def patched_npu_moe_init_routing(hidden_states, topk_ids, *args, **kwargs):
|
||||
try:
|
||||
bound = origin_signature.bind(hidden_states, topk_ids, *args, **kwargs)
|
||||
except TypeError as e:
|
||||
raise RuntimeError('Failed to bind vLLM-Ascend npu_moe_init_routing arguments: '
|
||||
f'signature={origin_signature}, args={args}, kwargs={kwargs}.') from e
|
||||
bound.apply_defaults()
|
||||
routing_kwargs = {key: bound.arguments.get(key, default) for key, default in routing_defaults.items()}
|
||||
|
||||
if not is_nonquant_routing(routing_kwargs):
|
||||
return origin_routing(hidden_states, topk_ids, *args, **kwargs)
|
||||
logger.warning_once(
|
||||
'Using torch_npu.npu_moe_init_routing_v2 for vLLM-Ascend non-quantized MoE routing. '
|
||||
'The installed vLLM-Ascend implementation still dispatches this branch to '
|
||||
'npu_moe_init_routing_custom, whose missing custom-op binary fails asynchronously on this stack.')
|
||||
return npu_moe_init_routing_v2(hidden_states, topk_ids, routing_kwargs)
|
||||
|
||||
patched_npu_moe_init_routing._swift_nonquant_routing_patched = True
|
||||
patched_npu_moe_init_routing._swift_origin = origin_routing
|
||||
adaptor_cls.npu_moe_init_routing = staticmethod(patched_npu_moe_init_routing)
|
||||
|
||||
|
||||
def patch_vllm_ascend_moe_runtime() -> None:
|
||||
"""Apply MoE runtime patches that are independent of GRPO weight sync."""
|
||||
_patch_vllm_ascend_device_op_nonquant_routing()
|
||||
|
||||
|
||||
def _is_qwen_moe_model(model) -> bool:
|
||||
return getattr(getattr(model, 'config', None), 'model_type', None) in _QWEN_MOE_MODEL_TYPES
|
||||
|
||||
|
||||
def configure_vllm_ascend_moe_weight_sync(vllm_model, train_model, *, is_fsdp2: bool) -> None:
|
||||
"""Record the vLLM-Ascend MoE sync layout required by this training backend."""
|
||||
fsdp2_qwen_moe = is_fsdp2 and _is_qwen_moe_model(train_model)
|
||||
layout = _VLLM_ASCEND_MOE_PROCESSED_LAYOUT
|
||||
# Current vLLM-Ascend 0.18 non-quantized Qwen MoE forward keeps
|
||||
# ``need_trans=False`` and feeds ``w13_weight`` directly to
|
||||
# ``npu_grouped_matmul``. After FSDP2 runtime sync, write Qwen MoE weights
|
||||
# directly into the runtime [hidden, I_tp] direction and skip checkpoint
|
||||
# post-load processing; otherwise post-load transposes them back to
|
||||
# [I_tp, hidden] and the first rollout fails with a hidden-size mismatch
|
||||
# such as 2048 vs 192/384.
|
||||
setattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR, layout)
|
||||
setattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, fsdp2_qwen_moe)
|
||||
|
||||
|
||||
def configure_vllm_ascend_moe_preprocessed_weight_sync(vllm_model) -> None:
|
||||
"""Record that reload writes the layout expected before vLLM-Ascend post-processing."""
|
||||
setattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR, _VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT)
|
||||
setattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, False)
|
||||
|
||||
|
||||
def use_vllm_ascend_moe_preprocessed_weight(vllm_model) -> bool:
|
||||
"""Return whether runtime sync should write the pre-process MoE layout."""
|
||||
return getattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR,
|
||||
_VLLM_ASCEND_MOE_PROCESSED_LAYOUT) == _VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT
|
||||
|
||||
|
||||
def should_skip_vllm_ascend_moe_post_load(vllm_model) -> bool:
|
||||
"""Return whether vLLM post-load processing should be skipped after sync."""
|
||||
return bool(getattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, False))
|
||||
|
||||
|
||||
def expand_fused_moe_expert_names_for_vllm_ascend(name: str):
|
||||
"""Map Transformers fused Qwen MoE expert names to vLLM checkpoint names.
|
||||
|
||||
FSDP2 can expose Qwen-style MoE expert weights as fused tensors:
|
||||
|
||||
mlp.experts.gate_up_proj: [experts, 2 * intermediate, hidden]
|
||||
mlp.experts.down_proj : [experts, hidden, intermediate]
|
||||
|
||||
vLLM's Qwen MoE ``load_weights`` path expects checkpoint-style names such as
|
||||
``mlp.experts.0.gate_proj.weight`` / ``up_proj`` / ``down_proj`` and maps
|
||||
those names onto its internal ``w13_weight`` / ``w2_weight`` parameters.
|
||||
Use expert 0 only as a name anchor; the paired vLLM-Ascend weight-loader
|
||||
patch below copies all local experts from the full 3D tensor.
|
||||
"""
|
||||
gate_up_suffix = '.mlp.experts.gate_up_proj'
|
||||
down_suffix = '.mlp.experts.down_proj'
|
||||
if name.endswith(gate_up_suffix):
|
||||
prefix = name[:-len('gate_up_proj')]
|
||||
return [
|
||||
f'{prefix}0.gate_proj.weight',
|
||||
f'{prefix}0.up_proj.weight',
|
||||
]
|
||||
if name.endswith(down_suffix):
|
||||
prefix = name[:-len('down_proj')]
|
||||
return [f'{prefix}0.down_proj.weight']
|
||||
return None
|
||||
|
||||
|
||||
def expand_fused_moe_expert_weight_for_vllm_ascend(name: str, param):
|
||||
"""Expand one FSDP2 fused Qwen MoE expert tensor for vLLM-Ascend weight sync."""
|
||||
if not isinstance(param, torch.Tensor) or param.dim() != 3:
|
||||
return None
|
||||
expanded_names = expand_fused_moe_expert_names_for_vllm_ascend(name)
|
||||
if expanded_names is None:
|
||||
return None
|
||||
if name.endswith('.mlp.experts.gate_up_proj'):
|
||||
gate_proj, up_proj = param.chunk(2, dim=1)
|
||||
return [
|
||||
(expanded_names[0], gate_proj.contiguous()),
|
||||
(expanded_names[1], up_proj.contiguous()),
|
||||
]
|
||||
if name.endswith('.mlp.experts.down_proj'):
|
||||
return [(expanded_names[0], param)]
|
||||
return None
|
||||
|
||||
|
||||
def patch_vllm_ascend_moe_expert_weight_loader(experts,
|
||||
name: str,
|
||||
param,
|
||||
*,
|
||||
load_preprocessed_weight: bool = False) -> None:
|
||||
"""Patch one processed vLLM-Ascend MoE expert parameter loader.
|
||||
|
||||
vLLM-Ascend transposes unquantized MoE weights after each model load
|
||||
so grouped matmul can consume them efficiently. During GRPO weight sync,
|
||||
however, SWIFT can send regular HF/Megatron expert weights, for example:
|
||||
|
||||
gate_proj/up_proj: [intermediate, hidden] -> w13_weight
|
||||
down_proj : [hidden, intermediate] -> w2_weight
|
||||
|
||||
FSDP2 Qwen MoE may expose the same weights as fused 3D tensors. SWIFT
|
||||
expands those tensors to checkpoint-style gate/up/down names before calling
|
||||
vLLM ``load_weights``:
|
||||
|
||||
gate_proj/up_proj: [experts, intermediate, hidden]
|
||||
down_proj : [experts, hidden, intermediate]
|
||||
|
||||
Full-weight server reload still writes the pre-processed layout and then
|
||||
calls ``process_weights_after_loading`` once, letting vLLM-Ascend transpose
|
||||
complete weights afterwards:
|
||||
|
||||
w13_weight before process: [local_experts, 2 * intermediate_per_tp, hidden]
|
||||
w2_weight before process : [local_experts, hidden, intermediate_per_tp]
|
||||
|
||||
Megatron colocate runtime sync loads into the already-processed layout used
|
||||
by the existing Megatron rollout path:
|
||||
|
||||
w13_weight after process: [local_experts, hidden, 2 * intermediate_per_tp]
|
||||
w2_weight after process : [local_experts, intermediate_per_tp, hidden]
|
||||
|
||||
``load_preprocessed_weight`` selects the server full-reload target. FSDP2
|
||||
Qwen MoE colocate runtime sync keeps the processed target and deliberately
|
||||
skips the post-load transpose because current vLLM-Ascend non-quantized
|
||||
grouped matmul consumes the [hidden, I_tp] direction in this path.
|
||||
|
||||
This wrapper keeps the normal vLLM loader for initial checkpoint load,
|
||||
quantized experts, and non-Ascend backends. It only handles the 3D
|
||||
vLLM-Ascend expert tensors when a 2D or fused 3D runtime-sync tensor is
|
||||
loaded into ``w13_weight`` or ``w2_weight``.
|
||||
"""
|
||||
if 'w13_weight' not in name and 'w2_weight' not in name:
|
||||
return
|
||||
quant_method = getattr(experts, 'quant_method', None)
|
||||
quant_method_module = type(quant_method).__module__ if quant_method is not None else ''
|
||||
if not quant_method_module.startswith('vllm_ascend'):
|
||||
return
|
||||
|
||||
def make_ascend_moe_weight_loader(experts, origin_weight_loader):
|
||||
|
||||
def load_processed_ascend_weight(param, loaded_weight, weight_name, shard_id, expert_id, return_success=False):
|
||||
quant_method = getattr(experts, 'quant_method', None)
|
||||
quant_method_module = type(quant_method).__module__ if quant_method is not None else ''
|
||||
# Only the GRPO runtime-sync path needs special handling here.
|
||||
# SWIFT provides HF/Megatron tensors, while vLLM-Ascend stores MoE
|
||||
# experts as 3D per-local-expert tensors. Initial checkpoint load
|
||||
# and other layouts continue to use the original vLLM loader.
|
||||
is_runtime_sync_into_processed_param = (
|
||||
param.data.dim() == 3 and loaded_weight.dim() in {2, 3}
|
||||
and quant_method_module.startswith('vllm_ascend'))
|
||||
if not is_runtime_sync_into_processed_param:
|
||||
return origin_weight_loader(param, loaded_weight, weight_name, shard_id, expert_id, return_success)
|
||||
|
||||
is_w13_shard = shard_id in {'w1', 'w3'} and 'w13_weight' in weight_name
|
||||
is_w2_shard = shard_id == 'w2' and 'w2_weight' in weight_name
|
||||
|
||||
loaded_expert_sample = loaded_weight[0] if loaded_weight.dim() == 3 else loaded_weight
|
||||
|
||||
def prepare_fsdp2_preprocessed_target_layout():
|
||||
"""FSDP2 path: write weights before vLLM-Ascend post-load processing."""
|
||||
if is_w13_shard and param.data.shape[1] == loaded_expert_sample.shape[-1]:
|
||||
param.data = param.data.transpose(1, 2).contiguous()
|
||||
elif is_w2_shard and param.data.shape[2] == loaded_expert_sample.shape[0]:
|
||||
param.data = param.data.transpose(1, 2).contiguous()
|
||||
|
||||
def prepare_megatron_processed_target_layout():
|
||||
"""Megatron path: write weights into vLLM-Ascend runtime layout."""
|
||||
if (is_w13_shard and param.data.shape[-1] == loaded_expert_sample.shape[-1]
|
||||
and param.data.shape[-2] != loaded_expert_sample.shape[-1]):
|
||||
param.data = param.data.transpose(1, 2).contiguous()
|
||||
elif (is_w2_shard and param.data.shape[-2] == loaded_expert_sample.shape[0]
|
||||
and param.data.shape[-1] != loaded_expert_sample.shape[0]):
|
||||
param.data = param.data.transpose(1, 2).contiguous()
|
||||
|
||||
tp_rank = experts.tp_rank
|
||||
|
||||
def copy_fsdp2_preprocessed_expert(local_expert_id: int, loaded_expert_weight) -> bool:
|
||||
"""Copy FSDP2 fused expert weights into pre-process vLLM-Ascend layout."""
|
||||
param_data = param.data[local_expert_id]
|
||||
if is_w13_shard:
|
||||
# Target: [2 * intermediate_per_tp, hidden].
|
||||
shard_size = param_data.shape[0] // 2
|
||||
loaded_expert_weight = loaded_expert_weight.narrow(0, shard_size * tp_rank, shard_size)
|
||||
offset = 0 if shard_id == 'w1' else shard_size
|
||||
param_data[offset:offset + shard_size].copy_(loaded_expert_weight.contiguous())
|
||||
return True
|
||||
|
||||
if is_w2_shard:
|
||||
# Target: [hidden, intermediate_per_tp].
|
||||
shard_size = param_data.shape[1]
|
||||
loaded_expert_weight = loaded_expert_weight.narrow(1, shard_size * tp_rank, shard_size)
|
||||
param_data.copy_(loaded_expert_weight.contiguous())
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def copy_megatron_processed_expert(local_expert_id: int, loaded_expert_weight) -> bool:
|
||||
"""Copy Megatron/HF expert shards into processed vLLM-Ascend layout."""
|
||||
param_data = param.data[local_expert_id]
|
||||
if is_w13_shard:
|
||||
# Target: [hidden, 2 * intermediate_per_tp].
|
||||
shard_size = param_data.shape[1] // 2
|
||||
loaded_expert_weight = loaded_expert_weight.narrow(0, shard_size * tp_rank, shard_size)
|
||||
offset = 0 if shard_id == 'w1' else shard_size
|
||||
param_data[:, offset:offset + shard_size].copy_(loaded_expert_weight.transpose(0, 1).contiguous())
|
||||
return True
|
||||
|
||||
if is_w2_shard:
|
||||
# Target: [intermediate_per_tp, hidden].
|
||||
shard_size = param_data.shape[0]
|
||||
loaded_expert_weight = loaded_expert_weight.narrow(1, shard_size * tp_rank, shard_size)
|
||||
param_data.copy_(loaded_expert_weight.transpose(0, 1).contiguous())
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
if load_preprocessed_weight:
|
||||
prepare_fsdp2_preprocessed_target_layout()
|
||||
copy_one_expert = copy_fsdp2_preprocessed_expert
|
||||
else:
|
||||
prepare_megatron_processed_target_layout()
|
||||
copy_one_expert = copy_megatron_processed_expert
|
||||
|
||||
if loaded_weight.dim() == 3:
|
||||
copied = False
|
||||
for global_expert_id, loaded_expert_weight in enumerate(loaded_weight):
|
||||
local_expert_id = experts._map_global_expert_id_to_local_expert_id(global_expert_id)
|
||||
if local_expert_id == -1:
|
||||
continue
|
||||
copied = copy_one_expert(local_expert_id, loaded_expert_weight) or copied
|
||||
return copied if return_success else None
|
||||
|
||||
local_expert_id = experts._map_global_expert_id_to_local_expert_id(expert_id)
|
||||
if local_expert_id == -1:
|
||||
return False if return_success else None
|
||||
|
||||
if copy_one_expert(local_expert_id, loaded_weight):
|
||||
return True if return_success else None
|
||||
|
||||
return origin_weight_loader(param, loaded_weight, weight_name, shard_id, expert_id, return_success)
|
||||
|
||||
load_processed_ascend_weight._swift_ascend_moe_weight_loader = True
|
||||
load_processed_ascend_weight._swift_origin_weight_loader = origin_weight_loader
|
||||
load_processed_ascend_weight._swift_load_preprocessed_weight = load_preprocessed_weight
|
||||
return load_processed_ascend_weight
|
||||
|
||||
if not hasattr(experts, 'weight_loader'):
|
||||
return
|
||||
weight_loader = getattr(param, 'weight_loader', experts.weight_loader)
|
||||
origin_weight_loader = getattr(weight_loader, '_swift_origin_weight_loader', weight_loader)
|
||||
if (not getattr(weight_loader, '_swift_ascend_moe_weight_loader', False)
|
||||
or getattr(weight_loader, '_swift_load_preprocessed_weight', None) != load_preprocessed_weight):
|
||||
param.weight_loader = make_ascend_moe_weight_loader(experts, origin_weight_loader)
|
||||
|
||||
|
||||
__all__ = [
|
||||
'configure_vllm_ascend_moe_preprocessed_weight_sync',
|
||||
'configure_vllm_ascend_moe_weight_sync',
|
||||
'expand_fused_moe_expert_names_for_vllm_ascend',
|
||||
'expand_fused_moe_expert_weight_for_vllm_ascend',
|
||||
'patch_vllm_ascend_moe_expert_weight_loader',
|
||||
'patch_vllm_ascend_moe_runtime',
|
||||
'should_skip_vllm_ascend_moe_post_load',
|
||||
'use_vllm_ascend_moe_preprocessed_weight',
|
||||
]
|
||||
Reference in New Issue
Block a user