diff --git a/miles/backends/megatron_utils/checkpoint.py b/miles/backends/megatron_utils/checkpoint.py index 3b1571b7adb..b5808a8f86e 100644 --- a/miles/backends/megatron_utils/checkpoint.py +++ b/miles/backends/megatron_utils/checkpoint.py @@ -121,6 +121,8 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_con if has_local_checkpoint_manager or _is_megatron_checkpoint(load_path): if not has_local_checkpoint_manager and is_dsv4_model(args): assert_checkpoint_is_current(load_path) + if not has_local_checkpoint_manager and _has_linear_attention(ddp_model): + assert_linear_attn_checkpoint_is_current(load_path) result = _load_checkpoint_megatron( ddp_model=ddp_model, optimizer=optimizer, @@ -199,6 +201,35 @@ def publish_as_release(save_dir: str) -> None: shutil.move(source_dir, target_dir) +LINEAR_ATTN_LAYOUT_KEY = "weight_layout_version" + + +def _has_linear_attention(ddp_model) -> bool: + return any(LINEAR_ATTN_LAYOUT_KEY in name for chunk in ddp_model for name, _ in chunk.named_buffers()) + + +def assert_linear_attn_checkpoint_is_current(load_dir: str) -> None: + """Reject torch_dist checkpoints written before the shared linear-attention layer (group-major GDN + q/k/v and conv rows, Kimi-K3 KDA under linear_attn.*). Megatron skips tensors a checkpoint lacks, so + such a checkpoint would load scrambled or half-initialized instead of failing.""" + from torch.distributed.checkpoint import FileSystemReader + + path = Path(load_dir) + step_file = path / "latest_checkpointed_iteration.txt" + if step_file.is_file(): + step = step_file.read_text().strip() + path = path / (step if step == "release" else f"iter_{int(step):07d}") + if not path.is_dir(): + return + if not any( + key.endswith(LINEAR_ATTN_LAYOUT_KEY) for key in FileSystemReader(path).read_metadata().state_dict_metadata + ): + raise ValueError( + f"{load_dir} was written before the shared linear-attention layer's weight layout and would load " + "scrambled. Re-convert it from the HuggingFace checkpoint (tools/convert_hf_to_torch_dist.py)." + ) + + def _is_megatron_checkpoint(path: str | Path) -> bool: return (Path(path) / "latest_checkpointed_iteration.txt").is_file() or bool( re.fullmatch(r"iter_\d{7}", Path(path).name) diff --git a/miles/backends/megatron_utils/megatron_to_hf/linear_attn_layout.py b/miles/backends/megatron_utils/megatron_to_hf/linear_attn_layout.py index 8f5d6804ecf..f0f5523ba73 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/linear_attn_layout.py +++ b/miles/backends/megatron_utils/megatron_to_hf/linear_attn_layout.py @@ -8,9 +8,12 @@ """ from dataclasses import dataclass +from functools import cache import torch +from miles.utils.hf_utils.config import load_hf_config + @dataclass(frozen=True) class LinearAttnHeads: @@ -89,3 +92,8 @@ def qkv_flat_to_group_major(weight: torch.Tensor, heads: LinearAttnHeads) -> tor def qkv_group_major_to_flat(weight: torch.Tensor, heads: LinearAttnHeads) -> torch.Tensor: assert weight.shape[0] == heads.qkv_dim, (weight.shape, heads) return torch.cat(split_group_major(weight, _qkv_rows(heads), heads)) + + +@cache +def gdn_heads_of(hf_checkpoint: str) -> LinearAttnHeads: + return gdn_heads(load_hf_config(hf_checkpoint)) diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py index 70d155f76e4..3e22e66ea17 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py @@ -2,6 +2,8 @@ import torch +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import gdn_heads_of, qkv_group_major_to_flat + def _convert_mtp_layer(args, name, param, layer_idx): """Convert MTP layer parameters from Megatron to HuggingFace format.""" @@ -168,15 +170,15 @@ def convert_qwen3_5_to_hf(args, name, param): return [(f"{prefix}.self_attn.q_norm.weight", param)] elif rest == "self_attention.k_layernorm.weight": return [(f"{prefix}.self_attn.k_norm.weight", param)] + elif rest in ("self_attention.linear_attn.in_proj_qkv.weight", "self_attention.linear_attn.conv1d.weight"): + hf_rest = rest[len("self_attention.") :] + return [(f"{prefix}.{hf_rest}", qkv_group_major_to_flat(param, gdn_heads_of(args.hf_checkpoint)))] elif rest.startswith("self_attention.") and rest[len("self_attention.") :] in [ "input_layernorm.weight", - # linear attn (Qwen3.5 uses separate in_proj_b/in_proj_a) "linear_attn.A_log", - "linear_attn.conv1d.weight", "linear_attn.dt_bias", "linear_attn.in_proj_a.weight", "linear_attn.in_proj_b.weight", - "linear_attn.in_proj_qkv.weight", "linear_attn.in_proj_z.weight", "linear_attn.norm.weight", "linear_attn.out_proj.weight", diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_8_next.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_8_next.py index c2642a20a33..e46bd7df2dd 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_8_next.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_8_next.py @@ -37,6 +37,10 @@ "module.module.decoder.hc_head_contraction.input_mix_weight_up": "model.language_model.hyper_connection_mixer.input_mix_weight_up.weight", } +# Qwen3.8-Next runs the replicated GDN, which holds these in HF's own layout rather than the +# head-sharded layer's group-major rows, so they skip the Qwen3.5 permutation +_HF_LAYOUT_GDN = ("self_attention.linear_attn.in_proj_qkv.weight", "self_attention.linear_attn.conv1d.weight") + _LAYER_PATTERN = re.compile(r"module\.module\.decoder\.layers\.(\d+)\.(.+)") @@ -50,5 +54,8 @@ def convert_qwen3_8_next_to_hf(args, name, param): hf_suffix = _HC_SUFFIX_MAPPING.get(rest) if hf_suffix is not None: return [(f"model.language_model.layers.{layer_idx}.{hf_suffix}", param)] + if rest in _HF_LAYOUT_GDN: + hf_rest = rest[len("self_attention.") :] + return [(f"model.language_model.layers.{layer_idx}.{hf_rest}", param)] return convert_qwen3_5_to_hf(args, name, param) diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py index 94266a1603b..14e8c6f1a6c 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py @@ -2,6 +2,8 @@ import torch +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import gdn_heads_of, qkv_group_major_to_flat + def convert_qwen3_next_to_hf(args, name, param): if name == "module.module.embedding.word_embeddings.weight": @@ -118,11 +120,17 @@ def convert_qwen3_next_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)] elif rest == "self_attention.k_layernorm.weight": return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)] + elif rest == "self_attention.linear_attn.conv1d.weight": + return [ + ( + f"model.layers.{layer_idx}.linear_attn.conv1d.weight", + qkv_group_major_to_flat(param, gdn_heads_of(args.hf_checkpoint)), + ) + ] elif rest.startswith("self_attention.") and rest[len("self_attention.") :] in [ "input_layernorm.weight", # linear attn "linear_attn.A_log", - "linear_attn.conv1d.weight", "linear_attn.dt_bias", "linear_attn.in_proj_ba.weight", "linear_attn.in_proj_qkvz.weight", diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index ebea08b2752..b3dfd8c0437 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -50,6 +50,8 @@ logger = logging.getLogger(__name__) +LINEAR_ATTENTION_BACKENDS = ("fla", "flashqla") + FULLY_ASYNC_ROLLOUT_PATH = "miles.rollout.fully_async_rollout.FullyAsyncRolloutFn" @@ -440,7 +442,7 @@ def add_train_arguments(parser): parser.add_argument( "--linear-attention-backend", type=str, - choices=["fla", "flashqla"], + choices=LINEAR_ATTENTION_BACKENDS, default="fla", help=( "Backend for Qwen GDN linear-attention layers. " diff --git a/miles/utils/external_utils/command_utils/base_backend.py b/miles/utils/external_utils/command_utils/base_backend.py index cc4cd51890d..8f0a3a9c8d2 100644 --- a/miles/utils/external_utils/command_utils/base_backend.py +++ b/miles/utils/external_utils/command_utils/base_backend.py @@ -32,10 +32,11 @@ CONVERSION_VERSION_FILE = "miles_conversion_version.txt" # Bump when a conversion this checkout writes would not load on a checkout holding another value: -# the skip path trusts a cache only when the stamp equals this checkout's own value, so any weight -# layout change bumps it. Values 1 and 2 were written briefly by an earlier numbering and are never -# trusted. 3: the layouts as of this commit. -CONVERSION_VERSION = "3" +# the skip path trusts a cache only when the stamp equals the running checkout's value, so every +# weight layout change bumps it. Values 1 and 2 were written briefly by an earlier numbering and +# are never trusted. 3: the layouts before the head-sharded linear-attention layer. +# 4: Qwen3.5 / Qwen3-Next GDN weights in group-major rows. +CONVERSION_VERSION = "4" @dataclass diff --git a/miles_plugins/mbridge/linear_attn.py b/miles_plugins/mbridge/linear_attn.py new file mode 100644 index 00000000000..9f78bfc034c --- /dev/null +++ b/miles_plugins/mbridge/linear_attn.py @@ -0,0 +1,6 @@ +class LinearAttentionBridgeMixin: + """Keeps the linear-attention weight_layout_version marker out of HF weight loading.""" + + def _weight_name_mapping_mcore_local_to_global(self, model): + mapping = super()._weight_name_mapping_mcore_local_to_global(model) + return {k: v for k, v in mapping.items() if not v.endswith("weight_layout_version")} diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index 3748cf93249..b292f6db2d4 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -5,10 +5,16 @@ from mbridge.core import register_model from mbridge.models import Qwen2MoEBridge +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import ( + gdn_heads, + qkv_flat_to_group_major, + qkv_group_major_to_flat, +) +from miles_plugins.mbridge.linear_attn import LinearAttentionBridgeMixin @register_model(["qwen3_5", "qwen3_5_moe", "qwen3_6", "qwen3_6_moe"]) -class Qwen3_5Bridge(Qwen2MoEBridge): +class Qwen3_5Bridge(LinearAttentionBridgeMixin, Qwen2MoEBridge): """ Bridge for Qwen3.5 / Qwen3.6 models (both dense and MoE variants). These share the ``qwen3_5_moe`` HF config schema: VLM layout under @@ -21,6 +27,12 @@ class Qwen3_5Bridge(Qwen2MoEBridge): ``_mtp_experts_fused()`` autodetects from the safetensor index. """ + # the head-sharded layer holds these group-major; HF keeps them flat + _GDN_GROUP_MAJOR = ( + "self_attention.linear_attn.in_proj_qkv.weight", + "self_attention.linear_attn.conv1d.weight", + ) + _DIRECT_MAPPING = { "embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight", "decoder.final_layernorm.weight": "model.language_model.norm.weight", @@ -340,6 +352,9 @@ def _weight_to_mcore_format( # from Bridge's global pre-cast to self.dtype. return hf_weights[0].to(dtype=torch.float32).contiguous() + if mcore_weights_name.endswith(self._GDN_GROUP_MAJOR): + return qkv_flat_to_group_major(hf_weights[0], gdn_heads(self._get_text_config())) + if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: # merge qkv assert len(hf_weights) == 3 @@ -389,7 +404,10 @@ def _weight_to_mcore_format( def _weight_to_hf_format( self, mcore_weights_name: str, mcore_weights: torch.Tensor ) -> tuple[list[str], list[torch.Tensor]]: - return super()._weight_to_hf_format(mcore_weights_name, mcore_weights) + names, weights = super()._weight_to_hf_format(mcore_weights_name, mcore_weights) + if mcore_weights_name.endswith(self._GDN_GROUP_MAJOR): + weights = [qkv_group_major_to_flat(weights[0], gdn_heads(self._get_text_config()))] + return names, weights def _build_config(self): text_config = self._get_text_config() diff --git a/miles_plugins/mbridge/qwen3_8_next.py b/miles_plugins/mbridge/qwen3_8_next.py index 8df2e48fbe2..d0720c5e1a4 100644 --- a/miles_plugins/mbridge/qwen3_8_next.py +++ b/miles_plugins/mbridge/qwen3_8_next.py @@ -15,6 +15,10 @@ class Qwen38NextBridge(Qwen3_5Bridge): """Weight mapping + Megatron config for Qwen3.8-Next.""" + # Qwen3.8-Next still runs the replicated GDN, which holds q/k/v and the conv in HF's own + # layout rather than the head-sharded layer's group-major rows + _GDN_GROUP_MAJOR = () + _DIRECT_MAPPING = Qwen3_5Bridge._DIRECT_MAPPING.copy() # no final norm in the checkpoint: the final mixer's hc_norm is the final norm diff --git a/miles_plugins/mbridge/qwen3_next.py b/miles_plugins/mbridge/qwen3_next.py index 77940752f36..6202e5eccfb 100644 --- a/miles_plugins/mbridge/qwen3_next.py +++ b/miles_plugins/mbridge/qwen3_next.py @@ -1,10 +1,19 @@ import torch + from mbridge.core import register_model from mbridge.models import Qwen2MoEBridge +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import ( + gdn_heads, + qkv_flat_to_group_major, + qkv_group_major_to_flat, +) +from miles_plugins.mbridge.linear_attn import LinearAttentionBridgeMixin + +_GDN_CONV = "self_attention.linear_attn.conv1d.weight" @register_model("qwen3_next") -class Qwen3NextBridge(Qwen2MoEBridge): +class Qwen3NextBridge(LinearAttentionBridgeMixin, Qwen2MoEBridge): _ATTENTION_MAPPING = ( Qwen2MoEBridge._ATTENTION_MAPPING | { @@ -41,6 +50,9 @@ class Qwen3NextBridge(Qwen2MoEBridge): def _weight_to_mcore_format( self, mcore_weights_name: str, hf_weights: list[torch.Tensor] ) -> tuple[list[str], list[torch.Tensor]]: + if mcore_weights_name.endswith(_GDN_CONV): + return qkv_flat_to_group_major(hf_weights[0], gdn_heads(self.hf_config)) + if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: # merge qkv assert len(hf_weights) == 3 @@ -75,6 +87,14 @@ def _weight_to_mcore_format( return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) + def _weight_to_hf_format( + self, mcore_weights_name: str, mcore_weights: torch.Tensor + ) -> tuple[list[str], list[torch.Tensor]]: + names, weights = super()._weight_to_hf_format(mcore_weights_name, mcore_weights) + if mcore_weights_name.endswith(_GDN_CONV): + weights = [qkv_group_major_to_flat(weights[0], gdn_heads(self.hf_config))] + return names, weights + def _build_config(self): mtp_args = {} if hasattr(self.hf_config, "num_nextn_predict_layers"): diff --git a/miles_plugins/models/linear_attn.py b/miles_plugins/models/linear_attn.py index 7c6ca596a64..969e62ba938 100644 --- a/miles_plugins/models/linear_attn.py +++ b/miles_plugins/models/linear_attn.py @@ -1,4 +1,4 @@ -"""Head-sharded linear attention (KDA, GDN) as Megatron modules. +"""Head-sharded linear attention (GDN, KDA) as Megatron modules. :class:`LinearAttentionLayer` is the ``self_attention`` drop-in: input norm, one TP collective in (identity / all-reduce, or all-gather / reduce-scatter under sequence parallelism), the context-parallel @@ -41,6 +41,7 @@ except ImportError: pass +WEIGHT_LAYOUT_VERSION = 1 INT32_ELEMENTS = 2**31 - 1 _CHUNK_ELEMENTS = 2**30 _CHANNEL_ALIGN = 128 @@ -74,6 +75,30 @@ def kda_kernel(): return chunk_kda +def gdn_recurrence(q, k, v, beta_logits, decay, A_log, dt_bias, *, backend, cu_seqlens, cp_context): + """q/k ``[b, s, G, hk]``, v ``[b, s, H, hv]`` (fla groups value heads per key head), decay ``[b, s, H]`` + -> ``[b, s, H, hv]``.""" + if cp_context is not None and backend != "fla": + raise NotImplementedError(f"GDN context parallelism requires the 'fla' backend, got {backend!r}.") + beta = beta_logits.sigmoid() + g = -A_log.float().exp() * F.softplus(decay.float() + dt_bias) + if backend == "flashqla": + q, k, v, g, beta = (t.contiguous() for t in (q, k, v, g, beta)) + out, _ = gdn_kernel(backend)( + q, + k, + v, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + **({"cp_context": cp_context} if cp_context is not None else {}), + ) + return out + + def kda_recurrence(q, k, v, beta_logits, decay, A_log, dt_bias, *, gate_lower_bound, cu_seqlens, cp_context): """q/k ``[b, s, H, hk]``, v ``[b, s, H, hv]``, decay ``[b, s, H * hv]`` (the low-rank forget gate, gated inside the kernel) -> ``[b, s, H, hv]``.""" @@ -210,6 +235,9 @@ def __init__( dtype = config.params_dtype self._sharded_params: dict[str, int] = {} + self.register_buffer( + "weight_layout_version", torch.tensor([WEIGHT_LAYOUT_VERSION], dtype=torch.int32, device=device) + ) self._build_projections() with get_cuda_rng_tracker().fork(): self.conv1d = ShardedShortConv( @@ -301,6 +329,30 @@ def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor | None, cp_context=N return core.reshape(batch, seq_len, -1) +class GatedDeltaNet(LinearAttention): + """Gated DeltaNet: one softplus-gated decay per value head, through fla's or FlashQLA's chunked + kernel. Models subclass it with their projections.""" + + def __init__(self, config, heads, conv_kernel_size, norm_eps, tp_group, backend="fla", norm_activation="silu"): + gdn_kernel(backend) + super().__init__(config, heads, conv_kernel_size, norm_eps, tp_group, norm_activation) + self.backend = backend + + def recurrence(self, q, k, v, beta_logits, decay, cu_seqlens, cp_context): + return gdn_recurrence( + q, + k, + v, + beta_logits, + decay, + self.A_log, + self.dt_bias, + backend=self.backend, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + + class KimiDeltaAttention(LinearAttention): """KDA with ``q_proj`` / ``k_proj`` / ``v_proj`` fused into one group-major ``in_proj_qkv``, ``g_proj`` (output gate), ``b_proj`` (beta), and the low-rank forget gate ``f_b_proj(f_a_proj(x))``, diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index fca9f7b4855..5be83430028 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -1,25 +1,16 @@ 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.process_groups_config import ProcessGroupCollection 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 +from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextRMSNorm -try: - from fla.modules import FusedRMSNormGated, ShortConvolution -except ImportError: - pass - -from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype -from miles_plugins.models.cp_utils import build_gdn_cp_context -from miles_plugins.models.linear_attn import gdn_kernel - -from .hf_attention import HuggingfaceAttention +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import gdn_heads +from miles.utils.hf_utils.config import load_hf_config +from miles_plugins.models.linear_attn import GatedDeltaNet, LinearAttentionLayer, Projections def _get_text_config(hf_config): @@ -29,202 +20,40 @@ def _get_text_config(hf_config): return hf_config -# Adapted from Qwen3NextGatedDeltaNet but with separate in_proj_qkv and in_proj_z -class Qwen3_5GatedDeltaNet(nn.Module): - """ - Qwen3.5 GatedDeltaNet with varlen support. - Unlike Qwen3Next which uses a combined in_proj_qkvz, Qwen3.5 uses - separate in_proj_qkv (for Q,K,V) and in_proj_z (for Z). - """ - - def __init__(self, config, layer_idx: int, args=None): - super().__init__() - self.gdn_backend = getattr(args, "linear_attention_backend", "fla") - self.chunk_gated_delta_rule = gdn_kernel(self.gdn_backend) - 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] - # Qwen3.8-Next gates the output norm with sigmoid while hidden_act stays silu - self.output_gate_activation = getattr(config, "output_gate_type", None) or 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, - ) - - # Separate projections for QKV and Z (unlike Qwen3Next which combines QKVZ) - projection_size_qkv = self.key_dim * 2 + self.value_dim - projection_size_z = self.value_dim - self.in_proj_qkv = nn.Linear(self.hidden_size, projection_size_qkv, bias=False) - self.in_proj_z = nn.Linear(self.hidden_size, projection_size_z, bias=False) - self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) - self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) - - # time step projection - self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) - - A = torch.empty(self.num_v_heads).uniform_(0, 16) - # Qwen3.5 ships A_log in fp32; Qwen3.6 reverted the HF weight to bf16. - # We still hold it in fp32 here for engineering simplicity (no special - # path for 3.6) — this matches the SGLang implementation, which also - # keeps A_log in fp32 regardless of the HF dtype. - self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) - mark_param_dtype(self.A_log, torch.float32) - - # HF stores this norm in fp32, but unlike A_log its precision impact is - # negligible and sglang runs it in bf16 on the rollout side — follow - # config.dtype (bf16) to stay equivalent to rollout. - self.norm = FusedRMSNormGated( - self.head_v_dim, - eps=self.layer_norm_epsilon, - activation=self.output_gate_activation, - device=torch.cuda.current_device(), - dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(), - ) +class Qwen3_5GatedDeltaNet(GatedDeltaNet): + """Qwen3.5 / 3.6 / 3.8 GDN: ``in_proj_qkv`` (group-major rows), ``in_proj_z``, ``in_proj_b``, ``in_proj_a``.""" - self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + def _build_projections(self): + hidden, local = self.config.hidden_size, self.local + self.in_proj_qkv = self.sharded_linear("in_proj_qkv", hidden, local.num_k_heads * local.group_qkv_dim) + self.in_proj_z = self.sharded_linear("in_proj_z", hidden, local.value_dim) + self.in_proj_b = self.sharded_linear("in_proj_b", hidden, local.num_v_heads) + self.in_proj_a = self.sharded_linear("in_proj_a", hidden, local.num_v_heads) - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor = None, - ): - batch_size, seq_len, _ = hidden_states.shape + def project(self, x): + return Projections(self.in_proj_qkv(x), self.in_proj_z(x), self.in_proj_b(x), self.in_proj_a(x)) - cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) - # Projections (flat layout: [Q_all, K_all, V_all]) - mixed_qkv = self.in_proj_qkv(hidden_states) - z = self.in_proj_z(hidden_states) - z = z.reshape(batch_size, seq_len, -1, self.head_v_dim) - b = self.in_proj_b(hidden_states) - a = self.in_proj_a(hidden_states) +class Attention(LinearAttentionLayer): + """GDN ``self_attention``; ``core_cls`` picks the family's HF projection layout.""" - # Convolution on the flat QKV (pass cp_context for boundary handling) - conv_cu_seqlens = cp_context.cu_seqlens if cp_context is not None else cu_seqlens - mixed_qkv, _ = self.conv1d( - x=mixed_qkv, - cu_seqlens=conv_cu_seqlens, - cp_context=cp_context, - ) - - # Split into Q, K, V (flat split, matching HF layout) - query, key, value = torch.split( - mixed_qkv, - [self.key_dim, self.key_dim, self.value_dim], - dim=-1, - ) - query = query.reshape(batch_size, seq_len, -1, self.head_k_dim) - key = key.reshape(batch_size, seq_len, -1, self.head_k_dim) - value = value.reshape(batch_size, seq_len, -1, self.head_v_dim) + core_cls: type[GatedDeltaNet] = Qwen3_5GatedDeltaNet - 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) - - if cp_context is not None: - if self.gdn_backend != "fla": - raise NotImplementedError( - f"GDN context parallelism requires the 'fla' backend, got {self.gdn_backend!r}." - ) - core_attn_out, _ = self.chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cp_context.cu_seqlens, - cp_context=cp_context, - ) - else: - if self.gdn_backend == "flashqla": - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - g = g.contiguous() - beta = beta.contiguous() - core_attn_out, _ = self.chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) - - 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(batch_size, seq_len, -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, - name: str | None = None, - ): - super().__init__( - args, + def __init__(self, args, config, layer_number: int, cp_comm_type=None, pg_collection=None, name=None): + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + text_config = _get_text_config(load_hf_config(args.hf_checkpoint)) + linear_attn = self.core_cls( config, - layer_number, - cp_comm_type, - pg_collection, - name=name, - ) - # Qwen3.5 is a VLM model with nested text_config - self.hf_config = _get_text_config(self.hf_config) - self.hf_config._attn_implementation = "flash_attention_2" - - self.linear_attn = Qwen3_5GatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args) - - # Use a simple RMSNorm - try: - from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextRMSNorm - - self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps) - except ImportError: - from torch.nn import RMSNorm - - self.input_layernorm = RMSNorm(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, + heads=gdn_heads(text_config), + conv_kernel_size=text_config.linear_conv_kernel_dim, + norm_eps=text_config.rms_norm_eps, + tp_group=pg_collection.tp, + backend=args.linear_attention_backend, + norm_activation=text_config.hidden_act, ) - return hidden_states + input_layernorm = Qwen3NextRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) + super().__init__(config, linear_attn, input_layernorm, pg_collection, allgather_cp=args.allgather_cp) def get_qwen3_5_spec(args, config, vp_stage): diff --git a/miles_plugins/models/qwen3_8_next/gdn.py b/miles_plugins/models/qwen3_8_next/gdn.py new file mode 100644 index 00000000000..05dde4d4ffd --- /dev/null +++ b/miles_plugins/models/qwen3_8_next/gdn.py @@ -0,0 +1,216 @@ +"""Qwen3.8-Next runs its GDN replicated on every TP rank through the HF-attention shim. +Moved from qwen3_5.py when Qwen3.5 / Qwen3-Next switched to the head-sharded layer.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.activations import ACT2FN + +try: + from fla.modules import FusedRMSNormGated, ShortConvolution +except ImportError: + pass + +from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype +from miles_plugins.models.cp_utils import build_gdn_cp_context +from miles_plugins.models.hf_attention import HuggingfaceAttention +from miles_plugins.models.linear_attn import gdn_kernel +from miles_plugins.models.qwen3_5 import _get_text_config + + +# Adapted from Qwen3NextGatedDeltaNet but with separate in_proj_qkv and in_proj_z +class ReplicatedGatedDeltaNet(nn.Module): + """ + Qwen3.5 GatedDeltaNet with varlen support. + Unlike Qwen3Next which uses a combined in_proj_qkvz, Qwen3.5 uses + separate in_proj_qkv (for Q,K,V) and in_proj_z (for Z). + """ + + def __init__(self, config, layer_idx: int, args=None): + super().__init__() + self.gdn_backend = getattr(args, "linear_attention_backend", "fla") + self.chunk_gated_delta_rule = gdn_kernel(self.gdn_backend) + 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] + # Qwen3.8-Next gates the output norm with sigmoid while hidden_act stays silu + self.output_gate_activation = getattr(config, "output_gate_type", None) or 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, + ) + + # Separate projections for QKV and Z (unlike Qwen3Next which combines QKVZ) + projection_size_qkv = self.key_dim * 2 + self.value_dim + projection_size_z = self.value_dim + self.in_proj_qkv = nn.Linear(self.hidden_size, projection_size_qkv, bias=False) + self.in_proj_z = nn.Linear(self.hidden_size, projection_size_z, bias=False) + self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + + # time step projection + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) + + A = torch.empty(self.num_v_heads).uniform_(0, 16) + # Qwen3.5 ships A_log in fp32; Qwen3.6 reverted the HF weight to bf16. + # We still hold it in fp32 here for engineering simplicity (no special + # path for 3.6) — this matches the SGLang implementation, which also + # keeps A_log in fp32 regardless of the HF dtype. + self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) + mark_param_dtype(self.A_log, torch.float32) + + # HF stores this norm in fp32, but unlike A_log its precision impact is + # negligible and sglang runs it in bf16 on the rollout side — follow + # config.dtype (bf16) to stay equivalent to rollout. + self.norm = FusedRMSNormGated( + self.head_v_dim, + eps=self.layer_norm_epsilon, + activation=self.output_gate_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 forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor = None, + ): + batch_size, seq_len, _ = hidden_states.shape + + cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) + + # Projections (flat layout: [Q_all, K_all, V_all]) + mixed_qkv = self.in_proj_qkv(hidden_states) + z = self.in_proj_z(hidden_states) + z = z.reshape(batch_size, seq_len, -1, self.head_v_dim) + b = self.in_proj_b(hidden_states) + a = self.in_proj_a(hidden_states) + + # Convolution on the flat QKV (pass cp_context for boundary handling) + conv_cu_seqlens = cp_context.cu_seqlens if cp_context is not None else cu_seqlens + mixed_qkv, _ = self.conv1d( + x=mixed_qkv, + cu_seqlens=conv_cu_seqlens, + cp_context=cp_context, + ) + + # Split into Q, K, V (flat split, matching HF layout) + query, key, value = torch.split( + mixed_qkv, + [self.key_dim, self.key_dim, self.value_dim], + dim=-1, + ) + query = query.reshape(batch_size, seq_len, -1, self.head_k_dim) + key = key.reshape(batch_size, seq_len, -1, self.head_k_dim) + value = value.reshape(batch_size, seq_len, -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) + + if cp_context is not None: + if self.gdn_backend != "fla": + raise NotImplementedError( + f"GDN context parallelism requires the 'fla' backend, got {self.gdn_backend!r}." + ) + core_attn_out, _ = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cp_context.cu_seqlens, + cp_context=cp_context, + ) + else: + if self.gdn_backend == "flashqla": + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + g = g.contiguous() + beta = beta.contiguous() + core_attn_out, _ = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + 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(batch_size, seq_len, -1) + + output = self.out_proj(core_attn_out) + return output + + +class ReplicatedGDNAttention(HuggingfaceAttention): + def __init__( + self, + args, + config, + layer_number: int, + cp_comm_type: str = "p2p", + pg_collection=None, + name: str | None = None, + ): + super().__init__( + args, + config, + layer_number, + cp_comm_type, + pg_collection, + name=name, + ) + # Qwen3.5 is a VLM model with nested text_config + self.hf_config = _get_text_config(self.hf_config) + self.hf_config._attn_implementation = "flash_attention_2" + + self.linear_attn = ReplicatedGatedDeltaNet(self.hf_config, self.hf_layer_idx, args=args) + + # Use a simple RMSNorm + try: + from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextRMSNorm + + self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps) + except ImportError: + from torch.nn import RMSNorm + + self.input_layernorm = RMSNorm(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 diff --git a/miles_plugins/models/qwen3_8_next/qwen3_8_next.py b/miles_plugins/models/qwen3_8_next/qwen3_8_next.py index 5b281546524..a563fbb36e9 100644 --- a/miles_plugins/models/qwen3_8_next/qwen3_8_next.py +++ b/miles_plugins/models/qwen3_8_next/qwen3_8_next.py @@ -15,8 +15,8 @@ from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from miles.utils.hf_utils.config import load_hf_config, register_hf_config_aliases -from miles_plugins.models.qwen3_5 import Attention as Qwen35LinearAttention from miles_plugins.models.qwen3_5 import _get_text_config +from miles_plugins.models.qwen3_8_next.gdn import ReplicatedGDNAttention as Qwen35LinearAttention from miles_plugins.models.qwen3_8_next.hyper_connection import ( Qwen38NextHCHeadContraction, Qwen38NextHyperConnection, diff --git a/miles_plugins/models/qwen3_next.py b/miles_plugins/models/qwen3_next.py index 53427783839..6c8b86c811f 100644 --- a/miles_plugins/models/qwen3_next.py +++ b/miles_plugins/models/qwen3_next.py @@ -1,229 +1,49 @@ 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 transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm -except ImportError: - pass +from miles_plugins.models.linear_attn import GatedDeltaNet, Projections +from miles_plugins.models.qwen3_5 import Attention as _Qwen3_5Attention -from miles_plugins.models.cp_utils import build_gdn_cp_context -from miles_plugins.models.linear_attn import gdn_kernel -from .hf_attention import HuggingfaceAttention +class Qwen3NextGatedDeltaNet(GatedDeltaNet): + """Qwen3-Next GDN: HF fuses ``in_proj_qkvz`` and ``in_proj_ba``, both group-major per key head, so the + projections stay fused and are split here.""" - -# 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, args=None): - super().__init__() - self.gdn_backend = getattr(args, "linear_attention_backend", "fla") - self.chunk_gated_delta_rule = gdn_kernel(self.gdn_backend) - 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(), + def _build_projections(self): + hidden, local = self.config.hidden_size, self.local + self.in_proj_qkvz = self.sharded_linear( + "in_proj_qkvz", hidden, local.num_k_heads * (local.group_qkv_dim + local.group_value_dim) ) - - 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, - ): - cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) - - 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) - - conv_cu_seqlens = cp_context.cu_seqlens if cp_context is not None else cu_seqlens - mixed_qkv, _ = self.conv1d( - x=mixed_qkv, - cu_seqlens=conv_cu_seqlens, - cp_context=cp_context, + self.in_proj_ba = self.sharded_linear("in_proj_ba", hidden, 2 * local.num_v_heads) + + def project(self, x): + batch, seq_len, _ = x.shape + local = self.local + qkv, z = ( + self.in_proj_qkvz(x) + .view(batch, seq_len, local.num_k_heads, -1) + .split([local.group_qkv_dim, local.group_value_dim], dim=-1) ) - - query, key, value = torch.split( - mixed_qkv, - [ - self.key_dim, - self.key_dim, - self.value_dim, - ], - dim=-1, + b, a = ( + self.in_proj_ba(x) + .view(batch, seq_len, local.num_k_heads, -1) + .split([local.v_per_k, local.v_per_k], 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) - - if cp_context is not None: - if self.gdn_backend != "fla": - raise NotImplementedError( - f"GDN context parallelism requires the 'fla' backend, got {self.gdn_backend!r}." - ) - core_attn_out, _ = self.chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cp_context.cu_seqlens, - cp_context=cp_context, - ) - else: - if self.gdn_backend == "flashqla": - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - g = g.contiguous() - beta = beta.contiguous() - core_attn_out, _ = self.chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) - - 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, - name: str | None = None, - ): - super().__init__( - args, - config, - layer_number, - cp_comm_type, - pg_collection, - name=name, + return Projections( + qkv.reshape(batch, seq_len, -1), + z.reshape(batch, seq_len, -1), + b.reshape(batch, seq_len, -1), + a.reshape(batch, seq_len, -1), ) - 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, args=args) - 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 +class Attention(_Qwen3_5Attention): + core_cls = Qwen3NextGatedDeltaNet def get_qwen3_next_spec(args, config, vp_stage): diff --git a/tests/e2e/precision/test_qwen3_5_cp_correctness.py b/tests/e2e/precision/test_qwen3_5_cp_correctness.py index e33dc125c3f..98085767abb 100644 --- a/tests/e2e/precision/test_qwen3_5_cp_correctness.py +++ b/tests/e2e/precision/test_qwen3_5_cp_correctness.py @@ -1,156 +1,99 @@ -"""Correctness test for Qwen3.5 GDN with native fla Context Parallel. +"""Head-sharded GDN and KDA under native fla context parallelism against the same layer on the whole +sequence. The sequence boundary falls inside a CP shard, so the recurrent state is passed mid-sequence. -Run with: - torchrun --nproc_per_node=2 tests/test_qwen3_5_cp_correctness.py # CP=2 - torchrun --nproc_per_node=4 tests/test_qwen3_5_cp_correctness.py # CP=4 - -Validates that GDN forward+backward with native fla CP produces results -consistent with the non-CP (single-rank full-sequence) baseline. + torchrun --nproc_per_node=2 tests/e2e/precision/test_qwen3_5_cp_correctness.py # CP=2 + torchrun --nproc_per_node=4 tests/e2e/precision/test_qwen3_5_cp_correctness.py # CP=4 """ import os import sys +from datetime import timedelta import torch import torch.distributed as dist - +from megatron.core import parallel_state +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig from tests.ci.ci_register import register_cuda_ci, register_rocm_ci -register_cuda_ci(est_time=300, suite="stage-c-4-gpu-h200", labels=["precision"], hardware=["hopper", "blackwell"]) -register_rocm_ci(est_time=200, suite="nightly-stage-c-4-gpu-mi350", labels=["precision"]) - - -def setup_dist(): - rank = int(os.environ["RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) - return rank, world_size, local_rank - - -def build_gdn_module(device, dtype=torch.bfloat16): - """Build a small Qwen3.5 GDN module for testing.""" - - class FakeConfig: - hidden_size = 256 - linear_num_value_heads = 4 - linear_num_key_heads = 2 - linear_key_head_dim = 64 - linear_value_head_dim = 64 - linear_conv_kernel_dim = 4 - hidden_act = "silu" - rms_norm_eps = 1e-6 - - FakeConfig.dtype = dtype - - from miles_plugins.models.qwen3_5 import Qwen3_5GatedDeltaNet - - return Qwen3_5GatedDeltaNet(FakeConfig, layer_idx=0).to(device=device, dtype=dtype) - - -def test_cp_forward_backward(rank, world_size): - device = torch.device(f"cuda:{rank}") - dtype = torch.bfloat16 - - # ---- Reference: full sequence on rank 0 (no CP) ---- - torch.manual_seed(42) - model_ref = build_gdn_module(device, dtype) - - total_seq_len = 128 * world_size # must be divisible by world_size - batch = 1 - - torch.manual_seed(123) - full_hidden = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype, requires_grad=True) - full_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device) +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import LinearAttnHeads - # Forward without CP - ref_out = model_ref(full_hidden, cu_seqlens=full_cu) - ref_loss = ref_out.sum() - ref_loss.backward() - ref_grad = full_hidden.grad.clone() +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "fast-gpu")) +from linear_attn_reference import ReplicatedGDN, ReplicatedKDA, build_layer, packed, rel_err # noqa: E402 - # ---- Test: CP across ranks ---- - torch.manual_seed(42) - model_cp = build_gdn_module(device, dtype) - # Copy weights from ref to ensure identical params - model_cp.load_state_dict(model_ref.state_dict()) +register_cuda_ci(est_time=300, suite="stage-c-4-gpu-h200", labels=["precision"], hardware=["hopper", "blackwell"]) +register_rocm_ci(est_time=200, suite="nightly-stage-c-4-gpu-mi350", labels=["precision"]) - # Set up CP context on the module - cp_group = dist.group.WORLD - model_cp.cp_group = cp_group - model_cp.cp_rank = rank - model_cp.cp_world_size = world_size +HIDDEN = 256 - # Each rank gets its local chunk - local_seq_len = total_seq_len // world_size - start = rank * local_seq_len - end = start + local_seq_len +def check(ref, config, rank, world_size): + layer = build_layer(ref, config, allgather_cp=True) + core = layer.linear_attn + total = 128 * world_size torch.manual_seed(123) - full_hidden_cp = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype) - local_hidden = full_hidden_cp[:, start:end, :].clone().contiguous().requires_grad_(True) - - # Global cu_seqlens (build_gdn_cp_context expects global boundaries) - global_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device) + full = torch.randn(1, total, HIDDEN, device="cuda", dtype=torch.bfloat16) + full_cu = torch.tensor([0, total // 3, total], dtype=torch.int32, device="cuda") - # Forward with CP - cp_out = model_cp(local_hidden, cu_seqlens=global_cu) - cp_loss = cp_out.sum() + x_ref = full.clone().requires_grad_() + ref_out = core.out_proj(core(x_ref, full_cu, None)) + ref_out.sum().backward() - # Reduce loss across ranks to match reference - dist.all_reduce(cp_loss, op=dist.ReduceOp.SUM) + local = total // world_size + x_local = full[:, rank * local : (rank + 1) * local].transpose(0, 1).contiguous().requires_grad_() + cp_out, _ = layer(x_local, packed_seq_params=packed(full_cu)) + loss = cp_out.sum() + dist.all_reduce(loss) + loss.backward() - cp_loss.backward() + def gather_seq(t): + parts = [torch.empty_like(t) for _ in range(world_size)] + dist.all_gather(parts, t.contiguous()) + return torch.cat(parts, dim=0) - # ---- Gather outputs for comparison ---- - gathered_out = [torch.zeros_like(cp_out) for _ in range(world_size)] - dist.all_gather(gathered_out, cp_out.contiguous()) - full_cp_out = torch.cat(gathered_out, dim=1) - - gathered_grad = [torch.zeros_like(local_hidden.grad) for _ in range(world_size)] - dist.all_gather(gathered_grad, local_hidden.grad.contiguous()) - full_cp_grad = torch.cat(gathered_grad, dim=1) - - if rank == 0: - # Compare outputs - out_diff = (ref_out.detach().float() - full_cp_out.detach().float()).abs() - out_max_diff = out_diff.max().item() - out_rel_diff = (out_diff / (ref_out.detach().float().abs() + 1e-8)).max().item() - - # Compare gradients - grad_diff = (ref_grad.float() - full_cp_grad.float()).abs() - grad_max_diff = grad_diff.max().item() - grad_rel_diff = (grad_diff / (ref_grad.float().abs() + 1e-8)).max().item() - - print(f"\n=== CP={world_size} Correctness Test ===") - print(f"Forward max abs diff: {out_max_diff:.6e} max rel diff: {out_rel_diff:.6e}") - print(f"Backward max abs diff: {grad_max_diff:.6e} max rel diff: {grad_rel_diff:.6e}") - - # bf16 tolerance: 1e-2 is generous for bf16 accumulated ops - fwd_ok = out_max_diff < 1e-2 - bwd_ok = grad_max_diff < 1e-2 - print(f"Forward PASS: {fwd_ok}") - print(f"Backward PASS: {bwd_ok}") - - if not (fwd_ok and bwd_ok): - print("FAILED!") - sys.exit(1) - else: - print(f"CP={world_size} test PASSED!") + return rel_err(ref_out.detach()[0], gather_seq(cp_out.detach())[:, 0]), rel_err( + x_ref.grad[0], gather_seq(x_local.grad)[:, 0] + ) def main(): - rank, world_size, _ = setup_dist() + rank, world_size = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group(backend="nccl", timeout=timedelta(seconds=120)) + parallel_state.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=world_size) + model_parallel_cuda_manual_seed(1234) + config = TransformerConfig( + num_layers=1, + hidden_size=HIDDEN, + num_attention_heads=4, + params_dtype=torch.bfloat16, + bf16=True, + use_cpu_initialization=False, + context_parallel_size=world_size, + ) + torch.manual_seed(42) + refs = { + "gdn": ReplicatedGDN("qwen3_5", HIDDEN, LinearAttnHeads(2, 4, 64, 64), torch.bfloat16).cuda(), + "kda": ReplicatedKDA(HIDDEN, LinearAttnHeads(4, 4, 64, 64), torch.bfloat16).cuda(), + } + failed = False try: - test_cp_forward_backward(rank, world_size) + for name, ref in refs.items(): + out_diff, grad_diff = check(ref, config, rank, world_size) + ok = out_diff < 1e-2 and grad_diff < 1e-2 + failed |= not ok + if rank == 0: + print( + f"[{name} CP={world_size}] rel err fwd {out_diff:.2e} dx {grad_diff:.2e} {'PASS' if ok else 'FAIL'}" + ) finally: + parallel_state.destroy_model_parallel() dist.destroy_process_group() + if failed: + sys.exit(1) if __name__ == "__main__": - # Self-bootstrap under torchrun when invoked as `python3 file.py` (the - # CUDA CI runner's mode). Already inside torchrun => RANK is set. if "RANK" not in os.environ: os.execvp("torchrun", ["torchrun", "--nproc_per_node=4", __file__]) main() diff --git a/tests/fast-gpu/linear_attn_reference.py b/tests/fast-gpu/linear_attn_reference.py index 1b339478df3..1626e830d5a 100644 --- a/tests/fast-gpu/linear_attn_reference.py +++ b/tests/fast-gpu/linear_attn_reference.py @@ -1,4 +1,4 @@ -"""Replicated, HF-layout reference for the head-sharded KDA layer, and helpers to load their +"""Replicated, HF-layout references for the head-sharded linear-attention layers, and helpers to load their weights into a :class:`LinearAttentionLayer`.""" from types import SimpleNamespace @@ -6,18 +6,73 @@ import torch import torch.distributed as dist import torch.nn as nn +import torch.nn.functional as F from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.ops.gated_delta_rule import chunk_gated_delta_rule from fla.ops.kda import chunk_kda from megatron.core.process_groups_config import ProcessGroupCollection from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import LinearAttnHeads, qkv_flat_to_group_major from miles_plugins.models.linear_attn import KimiDeltaAttention, LinearAttentionLayer +from miles_plugins.models.qwen3_5 import Qwen3_5GatedDeltaNet +from miles_plugins.models.qwen3_next import Qwen3NextGatedDeltaNet CONV = 4 EPS = 1e-6 KDA_LOWER_BOUND = -5.0 +class ReplicatedGDN(nn.Module): + """Pre-sharding GDN: every rank holds all heads. ``family`` picks the HF projection layout.""" + + def __init__(self, family: str, hidden: int, heads: LinearAttnHeads, dtype): + super().__init__() + self.family, self.heads = family, heads + h = heads + if family == "qwen3_5": + self.in_proj_qkv = nn.Linear(hidden, h.qkv_dim, bias=False) + self.in_proj_z = nn.Linear(hidden, h.value_dim, bias=False) + self.in_proj_b = nn.Linear(hidden, h.num_v_heads, bias=False) + self.in_proj_a = nn.Linear(hidden, h.num_v_heads, bias=False) + else: + self.in_proj_qkvz = nn.Linear(hidden, h.qkv_dim + h.value_dim, bias=False) + self.in_proj_ba = nn.Linear(hidden, 2 * h.num_v_heads, bias=False) + self.conv1d = ShortConvolution(hidden_size=h.qkv_dim, kernel_size=CONV, bias=False) + self.dt_bias = nn.Parameter(torch.rand(h.num_v_heads)) + self.A_log = nn.Parameter(torch.log(torch.empty(h.num_v_heads).uniform_(1, 16))) + self.norm = FusedRMSNormGated(h.head_v_dim, eps=EPS, activation="silu") + self.out_proj = nn.Linear(h.value_dim, hidden, bias=False) + self.to(dtype=dtype) + self.A_log.data = self.A_log.data.float() + + def _split(self, x): + h = self.heads + if self.family == "qwen3_5": + q, k, v = self.in_proj_qkv(x).split([h.key_dim, h.key_dim, h.value_dim], dim=-1) + return q, k, v, self.in_proj_z(x), self.in_proj_b(x), self.in_proj_a(x) + r, hv = h.v_per_k, h.head_v_dim + grouped = self.in_proj_qkvz(x).view(*x.shape[:-1], h.num_k_heads, 2 * h.head_k_dim + 2 * r * hv) + q, k, v, z = grouped.split([h.head_k_dim, h.head_k_dim, r * hv, r * hv], dim=-1) + b, a = self.in_proj_ba(x).view(*x.shape[:-1], h.num_k_heads, 2 * r).split([r, r], dim=-1) + return q.flatten(-2), k.flatten(-2), v.flatten(-2), z.flatten(-2), b.flatten(-2), a.flatten(-2) + + def forward(self, x, cu_seqlens): + h = self.heads + bsz, seq_len, _ = x.shape + q, k, v, z, b, a = self._split(x) + mixed, _ = self.conv1d(x=torch.cat([q, k, v], dim=-1), cu_seqlens=cu_seqlens) + q, k, v = mixed.split([h.key_dim, h.key_dim, h.value_dim], dim=-1) + q = q.reshape(bsz, seq_len, -1, h.head_k_dim).repeat_interleave(h.v_per_k, dim=2) + k = k.reshape(bsz, seq_len, -1, h.head_k_dim).repeat_interleave(h.v_per_k, dim=2) + v = v.reshape(bsz, seq_len, -1, h.head_v_dim) + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + out, _ = chunk_gated_delta_rule( + q, k, v, g=g, beta=b.sigmoid(), use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens + ) + out = self.norm(out.reshape(-1, h.head_v_dim), z.reshape(-1, h.head_v_dim)) + return self.out_proj(out.reshape(bsz, seq_len, -1)) + + class ReplicatedKDA(nn.Module): """Pre-sharding KDA in the Kimi-K3 / GLM-5.3-flash HF layout.""" @@ -78,17 +133,29 @@ def sharded_projections(ref, grad: bool = False) -> dict[str, torch.Tensor]: def w(module): return module.weight.grad if grad else module.weight - qkv = torch.cat([w(ref.q_proj), w(ref.k_proj), w(ref.v_proj)]) - return { - "in_proj_qkv": qkv_flat_to_group_major(qkv, ref.heads), - "g_proj": w(ref.g_proj), - "b_proj": w(ref.b_proj), - "f_b_proj": w(ref.f_b_proj), - } + if isinstance(ref, ReplicatedKDA): + qkv = torch.cat([w(ref.q_proj), w(ref.k_proj), w(ref.v_proj)]) + return { + "in_proj_qkv": qkv_flat_to_group_major(qkv, ref.heads), + "g_proj": w(ref.g_proj), + "b_proj": w(ref.b_proj), + "f_b_proj": w(ref.f_b_proj), + } + if ref.family == "qwen3_5": + return { + "in_proj_qkv": qkv_flat_to_group_major(w(ref.in_proj_qkv), ref.heads), + "in_proj_z": w(ref.in_proj_z), + "in_proj_b": w(ref.in_proj_b), + "in_proj_a": w(ref.in_proj_a), + } + return {"in_proj_qkvz": w(ref.in_proj_qkvz), "in_proj_ba": w(ref.in_proj_ba)} def conv_of(ref, grad: bool = False) -> torch.Tensor: - w = torch.cat([c.weight.grad if grad else c.weight for c in (ref.q_conv1d, ref.k_conv1d, ref.v_conv1d)]) + if isinstance(ref, ReplicatedKDA): + w = torch.cat([c.weight.grad if grad else c.weight for c in (ref.q_conv1d, ref.k_conv1d, ref.v_conv1d)]) + else: + w = ref.conv1d.weight.grad if grad else ref.conv1d.weight return qkv_flat_to_group_major(w, ref.heads) @@ -106,16 +173,23 @@ def build_layer(ref, config, allgather_cp: bool = True) -> LinearAttentionLayer: """A head-sharded layer holding this TP rank's shard of ``ref``'s weights, identity input norm.""" pg = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) tp = pg.tp - core = KimiDeltaAttention(config, ref.heads, CONV, EPS, tp, gate_lower_bound=KDA_LOWER_BOUND) + if isinstance(ref, ReplicatedKDA): + core = KimiDeltaAttention(config, ref.heads, CONV, EPS, tp, gate_lower_bound=KDA_LOWER_BOUND) + else: + core_cls = Qwen3_5GatedDeltaNet if ref.family == "qwen3_5" else Qwen3NextGatedDeltaNet + core = core_cls(config, ref.heads, CONV, EPS, tp) with torch.no_grad(): for name, full in sharded_projections(ref).items(): getattr(core, name).weight.copy_(shard(full, 0, tp)) core.conv1d.weight.copy_(shard(conv_of(ref), 0, tp)) core.A_log.copy_(shard(ref.A_log, 0, tp)) core.dt_bias.copy_(shard(ref.dt_bias, 0, tp)) - core.norm.weight.copy_(ref.o_norm.weight) - core.out_proj.weight.copy_(shard(ref.o_proj.weight, 1, tp)) - core.f_a_proj.weight.copy_(ref.f_a_proj.weight) + norm = ref.o_norm if isinstance(ref, ReplicatedKDA) else ref.norm + core.norm.weight.copy_(norm.weight) + out_proj = ref.o_proj if isinstance(ref, ReplicatedKDA) else ref.out_proj + core.out_proj.weight.copy_(shard(out_proj.weight, 1, tp)) + if isinstance(ref, ReplicatedKDA): + core.f_a_proj.weight.copy_(ref.f_a_proj.weight) return LinearAttentionLayer(config, core, nn.Identity(), pg, allgather_cp=allgather_cp) diff --git a/tests/fast-gpu/test_linear_attn_head_sharded.py b/tests/fast-gpu/test_linear_attn_head_sharded.py index a2b4dfd6272..6cb9433b288 100644 --- a/tests/fast-gpu/test_linear_attn_head_sharded.py +++ b/tests/fast-gpu/test_linear_attn_head_sharded.py @@ -1,4 +1,4 @@ -"""Head-sharded KDA against its replicated reference at +"""Head-sharded GDN (Qwen3.5, Qwen3-Next layouts) and KDA against their replicated references at TP = world size: forward output and every input / parameter gradient, gathered to the full layout. torchrun --nproc_per_node=2 tests/fast-gpu/test_linear_attn_head_sharded.py @@ -19,6 +19,7 @@ sys.path.insert(0, os.path.dirname(__file__)) from linear_attn_reference import ( # noqa: E402 + ReplicatedGDN, ReplicatedKDA, build_layer, conv_of, @@ -32,6 +33,8 @@ HIDDEN = 256 CASES = { + "qwen3_5": LinearAttnHeads(num_k_heads=4, num_v_heads=8, head_k_dim=64, head_v_dim=64), + "qwen3_next": LinearAttnHeads(num_k_heads=4, num_v_heads=8, head_k_dim=64, head_v_dim=64), "kda": LinearAttnHeads(num_k_heads=8, num_v_heads=8, head_k_dim=64, head_v_dim=64), } @@ -39,7 +42,11 @@ def check(name, config, tp_group): torch.manual_seed(7) heads = CASES[name] - ref = ReplicatedKDA(HIDDEN, heads, torch.bfloat16).cuda() + ref = ( + ReplicatedKDA(HIDDEN, heads, torch.bfloat16) + if name == "kda" + else ReplicatedGDN(name, HIDDEN, heads, torch.bfloat16) + ).cuda() layer = build_layer(ref, config) core = layer.linear_attn @@ -54,7 +61,8 @@ def check(name, config, tp_group): out_ref.backward(grad_out) out_new.backward(grad_out.transpose(0, 1)) - norm_ref, out_proj_ref = ref.o_norm, ref.o_proj + norm_ref = ref.o_norm if name == "kda" else ref.norm + out_proj_ref = ref.o_proj if name == "kda" else ref.out_proj errors = { "out": rel_err(out_ref, out_new.transpose(0, 1)), "dx": rel_err(x_ref.grad, x_new.grad.transpose(0, 1)), @@ -66,7 +74,8 @@ def check(name, config, tp_group): } for proj, full_grad in sharded_projections(ref, grad=True).items(): errors[proj] = rel_err(full_grad, gather(getattr(core, proj).weight.grad, 0, tp_group)) - errors["f_a_proj"] = rel_err(ref.f_a_proj.weight.grad, core.f_a_proj.weight.grad) + if name == "kda": + errors["f_a_proj"] = rel_err(ref.f_a_proj.weight.grad, core.f_a_proj.weight.grad) return errors diff --git a/tests/fast/backends/megatron_utils/test_linear_attn_checkpoint_layout.py b/tests/fast/backends/megatron_utils/test_linear_attn_checkpoint_layout.py new file mode 100644 index 00000000000..d40572c71a8 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_linear_attn_checkpoint_layout.py @@ -0,0 +1,75 @@ +"""assert_linear_attn_checkpoint_is_current: a torch_dist checkpoint without the linear-attention layout +marker is rejected before Megatron would load it half-initialized.""" + +import pytest +import torch +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint import FileSystemWriter + +from miles.backends.megatron_utils.checkpoint import ( + LINEAR_ATTN_LAYOUT_KEY, + _has_linear_attention, + assert_linear_attn_checkpoint_is_current, + publish_as_release, +) + +LAYER = "decoder.layers.0.self_attention.linear_attn" + + +def _write(path, keys): + path.mkdir(parents=True) + dcp.save({key: torch.zeros(2) for key in keys}, storage_writer=FileSystemWriter(path), no_dist=True) + + +def test_release_checkpoint_with_the_marker_passes(tmp_path): + _write(tmp_path / "release", [f"{LAYER}.in_proj_qkv.weight", f"{LAYER}.{LINEAR_ATTN_LAYOUT_KEY}"]) + (tmp_path / "latest_checkpointed_iteration.txt").write_text("release") + assert_linear_attn_checkpoint_is_current(str(tmp_path)) + + +def test_iteration_checkpoint_without_the_marker_is_rejected(tmp_path): + _write(tmp_path / "iter_0000012", [f"{LAYER}.in_proj_qkv.weight", f"{LAYER}.conv1d.weight"]) + (tmp_path / "latest_checkpointed_iteration.txt").write_text("12") + with pytest.raises(ValueError, match="Re-convert it from the HuggingFace checkpoint"): + assert_linear_attn_checkpoint_is_current(str(tmp_path)) + + +def test_only_models_with_linear_attention_are_checked(): + class Core(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer(LINEAR_ATTN_LAYOUT_KEY, torch.ones(1)) + + with_layer = torch.nn.Sequential(torch.nn.Linear(2, 2), Core()) + assert _has_linear_attention([with_layer]) + assert not _has_linear_attention([torch.nn.Sequential(torch.nn.Linear(2, 2))]) + + +def _converted(save_dir, marker): + iteration = save_dir / "iter_0000001" + iteration.mkdir() + (iteration / "weights.txt").write_text(marker) + + +def test_a_first_conversion_becomes_the_release(tmp_path): + _converted(tmp_path, "new") + + publish_as_release(str(tmp_path)) + + assert (tmp_path / "release" / "weights.txt").read_text() == "new" + assert not (tmp_path / "iter_0000001").exists() + + +def test_a_reconversion_replaces_the_previous_release(tmp_path): + """shutil.move into an existing directory nests, which would leave the old weights in place.""" + stale = tmp_path / "release" + stale.mkdir() + (stale / "weights.txt").write_text("stale") + (stale / "only-in-the-old-one.txt").write_text("stale") + _converted(tmp_path, "new") + + publish_as_release(str(tmp_path)) + + assert (tmp_path / "release" / "weights.txt").read_text() == "new" + assert not (tmp_path / "release" / "only-in-the-old-one.txt").exists() + assert not (tmp_path / "release" / "iter_0000001").exists() diff --git a/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py b/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py index bd07baf94b4..8cb3efdb798 100644 --- a/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py +++ b/tests/fast/backends/megatron_utils/test_qwen3_5_mtp_bridge_mapping.py @@ -112,6 +112,16 @@ def _build_base_config(self, **kwargs): sys.modules["mbridge.core"] = mbridge_core_mod sys.modules["mbridge.models"] = mbridge_models_mod + # the real mixin is dependency-free; loading its file directly skips the package + # __init__, which imports every bridge and their real megatron/mbridge deps + mixin_path = Path(__file__).resolve().parents[4] / "miles_plugins" / "mbridge" / "linear_attn.py" + spec = importlib.util.spec_from_file_location("miles_plugins.mbridge.linear_attn", mixin_path) + mixin_mod = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(mixin_mod) + sys.modules["miles_plugins.mbridge"] = types.ModuleType("miles_plugins.mbridge") + sys.modules["miles_plugins.mbridge.linear_attn"] = mixin_mod + def load_bridge_module(): install_bridge_stubs() diff --git a/tests/fast/backends/megatron_utils/test_qwen3_8_next_gdn_layout.py b/tests/fast/backends/megatron_utils/test_qwen3_8_next_gdn_layout.py new file mode 100644 index 00000000000..a12a0d3de78 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_qwen3_8_next_gdn_layout.py @@ -0,0 +1,35 @@ +"""Qwen3.8-Next runs the replicated GDN in HF's layout, so its weight conversion must not take the +head-sharded layer's group-major permutation the way Qwen3.5's does.""" + +import sys + +import pytest + +GDN_NAMES = ("self_attention.linear_attn.in_proj_qkv.weight", "self_attention.linear_attn.conv1d.weight") + + +@pytest.mark.parametrize("name", GDN_NAMES) +def test_the_export_passes_the_replicated_gdn_weights_through(name, monkeypatch): + torch = pytest.importorskip("torch") + monkeypatch.delitem(sys.modules, "miles.backends.megatron_utils.megatron_to_hf.qwen3_5", raising=False) + from miles.backends.megatron_utils.megatron_to_hf import qwen3_5 as qwen3_5_export + from miles.backends.megatron_utils.megatron_to_hf.qwen3_8_next import convert_qwen3_8_next_to_hf + + def _fail(*_args, **_kwargs): + raise AssertionError("Qwen3.8-Next must not take Qwen3.5's group-major permutation") + + monkeypatch.setattr(qwen3_5_export, "qkv_group_major_to_flat", _fail) + param = torch.zeros(4, 2) + converted = convert_qwen3_8_next_to_hf(None, f"module.module.decoder.layers.3.{name}", param) + + hf_rest = name[len("self_attention.") :] + assert converted == [(f"model.language_model.layers.3.{hf_rest}", param)] + + +def test_the_bridge_opts_out_of_the_group_major_names(): + pytest.importorskip("mbridge") + from miles_plugins.mbridge.qwen3_5 import Qwen3_5Bridge + from miles_plugins.mbridge.qwen3_8_next import Qwen38NextBridge + + assert Qwen3_5Bridge._GDN_GROUP_MAJOR == GDN_NAMES + assert Qwen38NextBridge._GDN_GROUP_MAJOR == () diff --git a/tests/manual/bench_linear_attn.py b/tests/manual/bench_linear_attn.py new file mode 100644 index 00000000000..8580b77ae98 --- /dev/null +++ b/tests/manual/bench_linear_attn.py @@ -0,0 +1,174 @@ +"""Speed, memory and precision of the head-sharded linear-attention layer at the attention configs of the +latest models, across sequence lengths and one TP x CP layout (world size = TP * CP). + + torchrun --nproc_per_node= tests/manual/bench_linear_attn.py [seq_len ...] + +Per (model, seq_len): wall ms / peak GiB of fwd+bwd for the head-sharded layer (sequence parallel, +zigzag CP relayout as in training) and, at CP=1, the replicated layer it replaced; relative error of +the output and input gradient against the replicated layer (CP=1) or against the same layer without +CP (CP>1, contiguous shards). +""" + +import os +import sys +import time +from datetime import timedelta + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + reduce_from_tensor_model_parallel_region, + reduce_scatter_to_sequence_parallel_region, + scatter_to_sequence_parallel_region, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig + +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import LinearAttnHeads + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "fast-gpu")) +from linear_attn_reference import ReplicatedGDN, ReplicatedKDA, build_layer, packed, rel_err # noqa: E402 + +MODELS = { + "qwen3.8-flash GDN": (2560, LinearAttnHeads(16, 48, 128, 128), "gdn"), + "glm5.3-flash KDA": (4096, LinearAttnHeads(64, 64, 128, 128), "kda"), + "kimi-k3 KDA": (7168, LinearAttnHeads(96, 96, 128, 128), "kda"), +} + + +def wall_ms(fn, iters): + fn() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + start = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) / iters * 1e3, torch.cuda.max_memory_allocated() / 2**30 + + +def zigzag_shard(full, cp_rank, cp_size): + chunks = full.chunk(2 * cp_size, dim=0) + return torch.cat([chunks[cp_rank], chunks[2 * cp_size - 1 - cp_rank]]).contiguous() + + +def contiguous_shard(full, cp_rank, cp_size): + return full.chunk(cp_size, dim=0)[cp_rank].contiguous() + + +def tp_gather(x, tp_group): + return ( + gather_from_sequence_parallel_region(x, tensor_parallel_output_grad=False, group=tp_group) + if tp_group.size() > 1 + else x + ) + + +def sp_shard(x, tp_group): + return x.chunk(tp_group.size(), dim=0)[tp_group.rank()].contiguous() + + +def run(name, seq_len, tp, cp, log): + hidden, heads, kind = MODELS[name] + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + config = TransformerConfig( + num_layers=1, + hidden_size=hidden, + num_attention_heads=heads.num_v_heads, + params_dtype=torch.bfloat16, + bf16=True, + use_cpu_initialization=False, + tensor_model_parallel_size=tp, + context_parallel_size=cp, + sequence_parallel=tp > 1, + ) + torch.manual_seed(7) + ref = ( + ReplicatedKDA(hidden, heads, torch.bfloat16) + if kind == "kda" + else ReplicatedGDN("qwen3_5", hidden, heads, torch.bfloat16) + ).cuda() + torch.manual_seed(11) + full = torch.randn(seq_len, 1, hidden, device="cuda", dtype=torch.bfloat16) + grad = torch.randn_like(full) + cu = torch.tensor([0, seq_len], device="cuda", dtype=torch.int32) + iters = max(2, min(10, 262144 // seq_len)) + row = f"[{name:18s} TP={tp} CP={cp} seq={seq_len:6d}]" + try: + layer = build_layer(ref, config, allgather_cp=False) + x = sp_shard(zigzag_shard(full, cp_group.rank(), cp), tp_group).requires_grad_() + new_ms, new_mem = wall_ms(lambda: layer(x, packed_seq_params=packed(cu))[0].float().sum().backward(), iters) + row += f" head-sharded {new_ms:8.2f} ms {new_mem:6.2f} GiB" + + if cp == 1: + + def replicated(x=x): + whole = ( + gather_from_sequence_parallel_region(x, tensor_parallel_output_grad=False, group=tp_group) + if tp > 1 + else x + ) + out = ref(whole.transpose(0, 1), cu).transpose(0, 1) + return scatter_to_sequence_parallel_region(out, group=tp_group) if tp > 1 else out + + old_ms, old_mem = wall_ms(lambda: replicated().float().sum().backward(), iters) + row += f" replicated {old_ms:8.2f} ms {old_mem:6.2f} GiB speed x{old_ms / new_ms:4.2f} mem x{old_mem / new_mem:4.2f}" + x_old, x_new = x.detach().clone().requires_grad_(), x.detach().clone().requires_grad_() + g = sp_shard(grad, tp_group) + out_old = replicated(x_old) + out_old.backward(g) + out_new = layer(x_new, packed_seq_params=packed(cu))[0] + out_new.backward(g) + row += f" rel err out {rel_err(out_old, out_new):.1e} dx {rel_err(x_old.grad, x_new.grad):.1e}" + else: + contig = build_layer(ref, config, allgather_cp=True) + x_cp = sp_shard(contiguous_shard(full, cp_group.rank(), cp), tp_group).requires_grad_() + out_cp = contig(x_cp, packed_seq_params=packed(cu))[0] + out_cp.backward(sp_shard(contiguous_shard(grad, cp_group.rank(), cp), tp_group)) + core = contig.linear_attn + x_ref = sp_shard(full, tp_group).requires_grad_() + whole = ( + gather_from_sequence_parallel_region(x_ref, tensor_parallel_output_grad=True, group=tp_group) + if tp > 1 + else x_ref + ) + out_ref = core.out_proj(core(whole.transpose(0, 1), cu, None)).transpose(0, 1) + out_ref = ( + reduce_scatter_to_sequence_parallel_region(out_ref, group=tp_group) + if tp > 1 + else reduce_from_tensor_model_parallel_region(out_ref, group=tp_group) + ) + out_ref.backward(sp_shard(grad, tp_group)) + start = cp_group.rank() * (seq_len // cp) + ref_out, ref_dx = tp_gather(out_ref.detach(), tp_group), tp_gather(x_ref.grad, tp_group) + cp_out, cp_dx = tp_gather(out_cp.detach(), tp_group), tp_gather(x_cp.grad, tp_group) + window = slice(start, start + seq_len // cp) + row += f" vs no-CP rel err out {rel_err(ref_out[window], cp_out):.1e} dx {rel_err(ref_dx[window], cp_dx):.1e}" + except torch.cuda.OutOfMemoryError: + row += " OOM" + torch.cuda.empty_cache() + log(row) + + +def main(): + tp, cp = int(sys.argv[1]), int(sys.argv[2]) + seq_lens = [int(s) for s in sys.argv[3:]] or [4096, 8192, 16384, 32768, 65536, 131072] + rank = int(os.environ["RANK"]) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", timeout=timedelta(minutes=30)) + parallel_state.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp) + model_parallel_cuda_manual_seed(1234) + log = (lambda s: print(s, flush=True)) if rank == 0 else (lambda s: None) + for name in MODELS: + for seq_len in seq_lens: + run(name, seq_len, tp, cp, log) + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/manual/perf_gdn_head_sharded.py b/tests/manual/perf_gdn_head_sharded.py new file mode 100644 index 00000000000..52724d80b84 --- /dev/null +++ b/tests/manual/perf_gdn_head_sharded.py @@ -0,0 +1,107 @@ +"""One GDN layer, fwd+bwd wall time, GPU time and peak memory: replicated (pre-sharding) vs head-sharded. + + torchrun --nproc_per_node=2 tests/manual/perf_gdn_head_sharded.py ... + +With ``sp`` both run as they do under ``--sequence-parallel``: the replicated layer all-gathers its input +and slices its output (what the old HF-attention shim did), the head-sharded layer takes the SP shard. +""" + +import os +import sys +import time +from datetime import timedelta + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig +from torch.profiler import ProfilerActivity, profile + +from miles.backends.megatron_utils.megatron_to_hf.linear_attn_layout import LinearAttnHeads + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "fast-gpu")) +from linear_attn_reference import ReplicatedGDN, build_layer, packed # noqa: E402 + + +def bench(fn, iters=10): + for _ in range(3): + fn() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + start = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + wall = (time.perf_counter() - start) / iters * 1e3 + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + gpu = sum(e.self_device_time_total for e in prof.key_averages()) / iters / 1e3 + return wall, gpu, torch.cuda.max_memory_allocated() / 2**30 + + +def main(): + sequence_parallel = sys.argv[1] == "sp" + hidden, k_heads, v_heads = (int(a) for a in sys.argv[2:5]) + seq_lens = [int(a) for a in sys.argv[5:]] + rank, world = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl", timeout=timedelta(seconds=120)) + parallel_state.initialize_model_parallel(tensor_model_parallel_size=world) + model_parallel_cuda_manual_seed(1234) + config = TransformerConfig( + num_layers=1, + hidden_size=hidden, + num_attention_heads=16, + params_dtype=torch.bfloat16, + bf16=True, + use_cpu_initialization=False, + tensor_model_parallel_size=world, + sequence_parallel=sequence_parallel, + ) + tp_group = parallel_state.get_tensor_model_parallel_group() + heads = LinearAttnHeads(k_heads, v_heads, 128, 128) + ref = ReplicatedGDN("qwen3_5", hidden, heads, torch.bfloat16).cuda() + layer = build_layer(ref, config) + if rank == 0: + print( + f"GDN layer fwd+bwd, hidden={hidden} heads={k_heads}/{v_heads}, TP={world}, SP={sequence_parallel}; wall ms / GPU ms / peak GiB" + ) + for seq_len in seq_lens: + local_len = seq_len // world if sequence_parallel else seq_len + x_sbh = torch.randn(local_len, 1, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + cu = torch.tensor([0, seq_len], device="cuda", dtype=torch.int32) + + def replicated(x=x_sbh, cu=cu): + full = ( + gather_from_sequence_parallel_region(x, tensor_parallel_output_grad=False, group=tp_group) + if sequence_parallel + else x + ) + out = ref(full.transpose(0, 1), cu).transpose(0, 1) + out = scatter_to_sequence_parallel_region(out, group=tp_group) if sequence_parallel else out + out.float().sum().backward() + + old = bench(replicated) + new = bench(lambda x=x_sbh, cu=cu: layer(x, packed_seq_params=packed(cu))[0].float().sum().backward()) + if rank == 0: + print( + f" seq {seq_len:6d} replicated {old[0]:7.2f} / {old[1]:7.2f} / {old[2]:5.2f}" + f" head-sharded {new[0]:7.2f} / {new[1]:7.2f} / {new[2]:5.2f}" + f" wall x{old[0] / new[0]:.2f} gpu x{old[1] / new[1]:.2f} mem x{old[2] / new[2]:.2f}" + ) + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +if __name__ == "__main__": + if "RANK" not in os.environ: + os.execvp("torchrun", ["torchrun", "--nproc_per_node=2", __file__, *sys.argv[1:]]) + main() diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index c1d76a7be00..22957e8fa2d 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -15,6 +15,7 @@ from miles.backends.megatron_utils.fp32_param_utils import enforce_marked_param_dtypes from miles.backends.megatron_utils.initialize import init from miles.backends.megatron_utils.model_provider import get_model_provider_func +from miles.utils.arguments import LINEAR_ATTENTION_BACKENDS from miles.utils.logging_utils import configure_logger_raw from miles.utils.memory_utils import print_memory from miles_plugins.models.deepseek_v4.arguments import add_dsv4_arguments @@ -26,6 +27,8 @@ def add_conversion_args(parser): add_dsv4_arguments(parser) add_dsa_arguments(parser) parser.add_argument("--hf-checkpoint", type=str, required=True, help="HuggingFace model path") + parser.add_argument("--linear-attention-backend", choices=LINEAR_ATTENTION_BACKENDS, default="fla") + parser.add_argument("--allgather-cp", action="store_true") parser.add_argument( "--megatron-to-hf-mode", choices=["raw", "bridge"], diff --git a/tools/convert_torch_dist_to_hf.py b/tools/convert_torch_dist_to_hf.py index 75caa890b00..6ae76540b92 100644 --- a/tools/convert_torch_dist_to_hf.py +++ b/tools/convert_torch_dist_to_hf.py @@ -201,6 +201,8 @@ def copy_assets(origin_hf_dir, output_dir): print(f"loading model from {args.input_dir}") t = time.time() megatron_args = torch.load(os.path.join(args.input_dir, "common.pt"), weights_only=False)["args"] + if args.origin_hf_dir: + megatron_args.hf_checkpoint = args.origin_hf_dir dist_cp.state_dict_loader._load_state_dict( state_dict, storage_reader=WrappedStorageReader(args.input_dir),