[6/8] feat(models): port Qwen3.5 / Qwen3-Next GDN to the head-sharded layer

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)

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))

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",

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)

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",

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. "

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

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")}

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()

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

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"):

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))``,