230 lines
9.0 KiB
Python
230 lines
9.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import copy
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
|
|
from megatron.core.transformer.spec_utils import ModuleSpec
|
|
from megatron.core.transformer.transformer_block import get_num_layers_to_build
|
|
from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
|
|
from transformers import AutoConfig
|
|
from transformers.activations import ACT2FN
|
|
|
|
try:
|
|
from fla.modules import FusedRMSNormGated, ShortConvolution
|
|
from fla.ops.gated_delta_rule import chunk_gated_delta_rule
|
|
from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm
|
|
except ImportError:
|
|
pass
|
|
|
|
from .hf_attention import HuggingfaceAttention
|
|
|
|
|
|
# adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564
|
|
class Qwen3NextGatedDeltaNet(nn.Module):
|
|
"""
|
|
Qwen3NextGatedDeltaNet with varlen support
|
|
"""
|
|
|
|
def __init__(self, config, layer_idx: int):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.num_v_heads = config.linear_num_value_heads
|
|
self.num_k_heads = config.linear_num_key_heads
|
|
self.head_k_dim = config.linear_key_head_dim
|
|
self.head_v_dim = config.linear_value_head_dim
|
|
self.key_dim = self.head_k_dim * self.num_k_heads
|
|
self.value_dim = self.head_v_dim * self.num_v_heads
|
|
|
|
self.conv_kernel_size = config.linear_conv_kernel_dim
|
|
self.layer_idx = layer_idx
|
|
self.activation = config.hidden_act
|
|
self.act = ACT2FN[config.hidden_act]
|
|
self.layer_norm_epsilon = config.rms_norm_eps
|
|
|
|
# QKV
|
|
self.conv_dim = self.key_dim * 2 + self.value_dim
|
|
self.conv1d = ShortConvolution(
|
|
hidden_size=self.conv_dim,
|
|
bias=False,
|
|
kernel_size=self.conv_kernel_size,
|
|
)
|
|
|
|
# projection of the input hidden states
|
|
projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
|
|
projection_size_ba = self.num_v_heads * 2
|
|
self.in_proj_qkvz = nn.Linear(self.hidden_size, projection_size_qkvz, bias=False)
|
|
self.in_proj_ba = nn.Linear(self.hidden_size, projection_size_ba, bias=False)
|
|
|
|
# time step projection (discretization)
|
|
# instantiate once and copy inv_dt in init_weights of PretrainedModel
|
|
self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
|
|
|
|
A = torch.empty(self.num_v_heads).uniform_(0, 16)
|
|
self.A_log = nn.Parameter(torch.log(A))
|
|
|
|
self.norm = FusedRMSNormGated(
|
|
self.head_v_dim,
|
|
eps=self.layer_norm_epsilon,
|
|
activation=self.activation,
|
|
device=torch.cuda.current_device(),
|
|
dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(),
|
|
)
|
|
|
|
self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
|
|
|
|
def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
|
|
"""
|
|
Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`.
|
|
"""
|
|
|
|
new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
|
|
self.num_k_heads,
|
|
2 * self.head_k_dim + 2 * self.head_v_dim * self.num_v_heads // self.num_k_heads,
|
|
)
|
|
new_tensor_shape_ba = mixed_ba.size()[:-1] + (self.num_k_heads, 2 * self.num_v_heads // self.num_k_heads)
|
|
|
|
mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
|
|
mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
|
|
split_arg_list_qkvz = [
|
|
self.head_k_dim,
|
|
self.head_k_dim,
|
|
(self.num_v_heads // self.num_k_heads * self.head_v_dim),
|
|
(self.num_v_heads // self.num_k_heads * self.head_v_dim),
|
|
]
|
|
split_arg_list_ba = [self.num_v_heads // self.num_k_heads, self.num_v_heads // self.num_k_heads]
|
|
query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=3)
|
|
b, a = torch.split(mixed_ba, split_arg_list_ba, dim=3)
|
|
# [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
|
|
value = value.reshape(value.size(0), value.size(1), -1, self.head_v_dim)
|
|
z = z.reshape(z.size(0), z.size(1), -1, self.head_v_dim)
|
|
b = b.reshape(b.size(0), b.size(1), self.num_v_heads)
|
|
a = a.reshape(a.size(0), a.size(1), self.num_v_heads)
|
|
return query, key, value, z, b, a
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cu_seqlens: torch.Tensor = None,
|
|
):
|
|
projected_states_qkvz = self.in_proj_qkvz(hidden_states)
|
|
projected_states_ba = self.in_proj_ba(hidden_states)
|
|
query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba)
|
|
query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value))
|
|
|
|
mixed_qkv = torch.cat((query, key, value), dim=-1)
|
|
|
|
mixed_qkv, _ = self.conv1d(
|
|
x=mixed_qkv,
|
|
cu_seqlens=cu_seqlens,
|
|
)
|
|
|
|
query, key, value = torch.split(
|
|
mixed_qkv,
|
|
[
|
|
self.key_dim,
|
|
self.key_dim,
|
|
self.value_dim,
|
|
],
|
|
dim=-1,
|
|
)
|
|
query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim)
|
|
key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim)
|
|
value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim)
|
|
|
|
beta = b.sigmoid()
|
|
# If the model is loaded in fp16, without the .float() here, A might be -inf
|
|
g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
|
|
if self.num_v_heads // self.num_k_heads > 1:
|
|
query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
|
|
key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
|
|
|
|
core_attn_out, last_recurrent_state = chunk_gated_delta_rule(
|
|
query,
|
|
key,
|
|
value,
|
|
g=g,
|
|
beta=beta,
|
|
initial_state=None,
|
|
output_final_state=False,
|
|
use_qk_l2norm_in_kernel=True,
|
|
)
|
|
|
|
z_shape_og = z.shape
|
|
# reshape input data into 2D tensor
|
|
core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
|
|
z = z.reshape(-1, z.shape[-1])
|
|
core_attn_out = self.norm(core_attn_out, z)
|
|
core_attn_out = core_attn_out.reshape(z_shape_og)
|
|
core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1)
|
|
|
|
output = self.out_proj(core_attn_out)
|
|
return output
|
|
|
|
|
|
class Attention(HuggingfaceAttention):
|
|
def __init__(
|
|
self,
|
|
args,
|
|
config,
|
|
layer_number: int,
|
|
cp_comm_type: str = "p2p",
|
|
pg_collection=None,
|
|
):
|
|
super().__init__(
|
|
args,
|
|
config,
|
|
layer_number,
|
|
cp_comm_type,
|
|
pg_collection,
|
|
)
|
|
if Qwen3NextAttention is None:
|
|
raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.")
|
|
|
|
self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx)
|
|
self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps)
|
|
|
|
def hf_forward(self, hidden_states, packed_seq_params):
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
hidden_states = self.linear_attn(
|
|
hidden_states=hidden_states,
|
|
cu_seqlens=packed_seq_params.cu_seqlens_q,
|
|
)
|
|
return hidden_states
|
|
|
|
|
|
def get_qwen3_next_spec(args, config, vp_stage):
|
|
# always use the moe path
|
|
if not args.num_experts:
|
|
config.moe_layer_freq = [0] * config.num_layers
|
|
|
|
# Define the decoder block spec
|
|
kwargs = {
|
|
"use_transformer_engine": True,
|
|
}
|
|
if vp_stage is not None:
|
|
kwargs["vp_stage"] = vp_stage
|
|
transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs)
|
|
|
|
assert config.pipeline_model_parallel_layout is None, "not support this at the moment"
|
|
|
|
# Slice the layer specs to only include the layers that are built in this pipeline stage.
|
|
# Note: MCore layer_number starts at 1
|
|
num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
|
|
offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
|
|
|
|
hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
|
|
|
|
for layer_id in range(num_layers_to_build):
|
|
if hf_config.layer_types[layer_id + offset] == "linear_attention":
|
|
layer_specs = copy.deepcopy(transformer_layer_spec.layer_specs[layer_id])
|
|
layer_specs.submodules.self_attention = ModuleSpec(
|
|
module=Attention,
|
|
params={"args": args},
|
|
)
|
|
transformer_layer_spec.layer_specs[layer_id] = layer_specs
|
|
return transformer_layer_spec
|