wehub-resource-sync
553 lines
23 KiB
Python
553 lines
23 KiB
Python
# 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
|