feat: support Qwen3-MoE weight conversion and dtype-only IPC grouping

awex/converter/sglang_converter.py

@@ -122,27 +122,44 @@ def _convert_attention_param(
                 # Keep fused format
                 return [(name, parameter)]
             else:
-                # Split into separate Q, K, V projections
+                # Split into separate Q, K, V projections. The fused dim0 is
+                # proportional to num_heads : num_kv_heads : num_kv_heads,
+                # which only degenerates to equal thirds for MHA — GQA models
+                # (num_kv_heads < num_heads) need head-count-weighted sizes.
+                # The ratio is TP-invariant since both head counts are divided
+                # by the same TP degree.
+                num_heads = int(self.total_num_heads)
+                num_kv_heads = int(self.total_kv_heads or num_heads)
+                total_units = num_heads + 2 * num_kv_heads
                 shape0 = parameter.shape[0]
-                stride = shape0 // 3
+                if (shape0 * num_heads) % total_units != 0 or (
+                    shape0 * num_kv_heads
+                ) % total_units != 0:
+                    raise ValueError(
+                        f"qkv dim0 {shape0} of {name} is not divisible into "
+                        f"q/k/v with num_heads={num_heads}, "
+                        f"num_kv_heads={num_kv_heads}"
+                    )
+                q_size = shape0 * num_heads // total_units
+                kv_size = shape0 * num_kv_heads // total_units
                 return [
                     (
                         name.replace("qkv_proj", "q_proj").replace(
                             "query_key_value", "q_proj"
                         ),
-                        parameter.narrow(0, 0, stride),
+                        parameter.narrow(0, 0, q_size),
                     ),
                     (
                         name.replace("qkv_proj", "k_proj").replace(
                             "query_key_value", "k_proj"
                         ),
-                        parameter.narrow(0, stride, stride),
+                        parameter.narrow(0, q_size, kv_size),
                     ),
                     (
                         name.replace("qkv_proj", "v_proj").replace(
                             "query_key_value", "v_proj"
                         ),
-                        parameter.narrow(0, 2 * stride, stride),
+                        parameter.narrow(0, q_size + kv_size, kv_size),
                     ),
                 ]
         elif "o_proj" in name or "dense" in name:

awex/models/qwen3_moe.py

@@ -0,0 +1,140 @@
+# Licensed to the Awex developers under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import List, Tuple
+
+import torch
+
+from awex import logging
+from awex.converter.sglang_converter import SGlangToHFWeightConverter
+
+logger = logging.getLogger(__name__)
+
+
+class SGlangToHFWeightConverterQwen3Moe(SGlangToHFWeightConverter):
+    """SGLang -> HF converter for Qwen3-MoE.
+
+    Splits the SGLang fused qkv_proj into canonical q/k/v projections
+    (GQA-aware split in the base class) so that inference-side weight
+    metadata matches the per-parameter HF names emitted by the Megatron
+    train-side converter. Expert parameters (experts.w13_weight /
+    experts.w2_weight) are expanded per expert by the base class.
+    """
+
+    def _fuse_qkv(self, name: str) -> bool:
+        # Train side reports canonical self_attn.{q,k,v}_proj names, so the
+        # inference side must unfuse qkv_proj for transfer-plan matching.
+        return False
+
+    def _convert_layer_norm_param(
+        self, name: str, parameter: torch.Tensor, layer_number: str
+    ) -> List[Tuple[str, torch.Tensor]]:
+        # Qwen3 uses self_attn.{q,k}_norm; the base class only recognizes
+        # the bailing-style {query,key}_layernorm names.
+        if "q_norm" in name or "k_norm" in name:
+            return [(name, parameter)]
+        return super()._convert_layer_norm_param(name, parameter, layer_number)
+
+
+def _build_mcore_converter_qwen3_moe():
+    # Lazily import Megatron converter to avoid MindSpeed patching in
+    # vLLM-only paths (same pattern as ling.py).
+    from awex.converter.mcore_converter import McoreToHFWeightConverter
+
+    class McoreToHFWeightConverterQwen3Moe(McoreToHFWeightConverter):
+        """Stock HF/sglang qwen3_moe serves canonical attention names
+        (self_attn.{q,k,v,o}_proj + self_attn.{q,k}_norm), so keep them
+        verbatim instead of the bailing-flavored renames, and split the
+        Megatron fused linear_qkv with GQA-aware group strides — the base
+        equal-thirds split only holds when q/k/v head counts match.
+        """
+
+        def _fuse_qkv(self, name: str) -> bool:
+            return False
+
+        @staticmethod
+        def _normalize_attn_name(name: str) -> str:
+            return name
+
+        def _gqa_head_dim(self) -> int:
+            head_dim = getattr(self.hf_config, "head_dim", None)
+            if head_dim:
+                return int(head_dim)
+            return int(self.hf_config.hidden_size // self.hf_config.num_attention_heads)
+
+        def _split_gqa_qkv(
+            self, parameter: torch.Tensor
+        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            hf = self.hf_config
+            head_dim = self._gqa_head_dim()
+            attn_tp = max(1, int(getattr(self.rank_info, "attn_tp_size", 1)))
+            if hf.num_key_value_heads % attn_tp != 0:
+                raise ValueError(
+                    f"num_key_value_heads ({hf.num_key_value_heads}) must be "
+                    f"divisible by attn_tp_size ({attn_tp})"
+                )
+            if hf.num_attention_heads % hf.num_key_value_heads != 0:
+                raise ValueError(
+                    f"num_attention_heads ({hf.num_attention_heads}) must be "
+                    f"divisible by num_key_value_heads "
+                    f"({hf.num_key_value_heads})"
+                )
+            num_groups = hf.num_key_value_heads // attn_tp
+            q_per_group = hf.num_attention_heads // hf.num_key_value_heads
+            group_rows = (q_per_group + 2) * head_dim
+            expected_rows = num_groups * group_rows
+            if parameter.shape[0] != expected_rows:
+                raise ValueError(
+                    "Unexpected linear_qkv rows for GQA split: "
+                    f"got {parameter.shape[0]}, expected {expected_rows} "
+                    f"(kv_heads={hf.num_key_value_heads}, attn_tp={attn_tp}, "
+                    f"q_per_group={q_per_group}, head_dim={head_dim})"
+                )
+            blocks = parameter.reshape(num_groups, group_rows, *parameter.shape[1:])
+            q_rows = q_per_group * head_dim
+            q = blocks[:, :q_rows].reshape(num_groups * q_rows, *parameter.shape[1:])
+            k = blocks[:, q_rows : q_rows + head_dim].reshape(
+                num_groups * head_dim, *parameter.shape[1:]
+            )
+            v = blocks[:, q_rows + head_dim :].reshape(
+                num_groups * head_dim, *parameter.shape[1:]
+            )
+            return q.contiguous(), k.contiguous(), v.contiguous()
+
+        def _convert_attention_param(
+            self, name: str, parameter: torch.Tensor, layer_number: str
+        ) -> List[Tuple[str, torch.Tensor]]:
+            if "self_attention.linear_qkv.weight" in name or (
+                "self_attention.linear_qkv.bias" in name
+            ):
+                suffix = "weight" if name.endswith("weight") else "bias"
+                q, k, v = self._split_gqa_qkv(parameter)
+                return [
+                    (f"self_attn.q_proj.{suffix}", q),
+                    (f"self_attn.k_proj.{suffix}", k),
+                    (f"self_attn.v_proj.{suffix}", v),
+                ]
+            return super()._convert_attention_param(name, parameter, layer_number)
+
+    return McoreToHFWeightConverterQwen3Moe
+
+
+CONFIG = {
+    "model_name": "Qwen3MoeForCausalLM",
+    "mcore_converter": _build_mcore_converter_qwen3_moe,
+    "sglang_converter": SGlangToHFWeightConverterQwen3Moe,
+}

awex/reader/nccl_reader.py

@@ -38,11 +38,7 @@
 )
 from awex.util.gpu import get_gpu_status, print_current_gpu_status
 from awex.util.system_util import count_open_fds
-from awex.util.tensor_util import (
-    cuda_ipc_deserialize,
-    ipc_deserialize,
-    reconstruct_tensors_from_groups,
-)
+from awex.util.tensor_util import reconstruct_ipc_weights
 
 logger = logging.getLogger(__name__)
 
@@ -309,17 +305,11 @@ def collect_training_weights(self, step_id, **kwargs):
         )
         logger.info(f"Open fds before deserialization: {count_open_fds()}")
         # Deserialize weights into tensors
-        if self.ipc_backend in ("cpu", "npu"):
-            group_shared, metadata, names = ipc_deserialize(serialized_weights)
-            group_shared = [t.to(device_id) for t in group_shared]
-        else:
-            group_shared, metadata, names = cuda_ipc_deserialize(serialized_weights)
-        device_util.synchronize(device_id=device_util.current_device())
-        tensors = reconstruct_tensors_from_groups(group_shared, metadata)
-        device_util.synchronize(device_id=device_util.current_device())
-        self.deserialized_weights = dict(zip(names, tensors))
+        self.deserialized_weights, num_groups = reconstruct_ipc_weights(
+            serialized_weights, ipc_backend=self.ipc_backend, device_id=device_id
+        )
         logger.info(
-            f"Deserialized {len(self.deserialized_weights)} parameters and {len(group_shared)} groups"
+            f"Deserialized {len(self.deserialized_weights)} parameters and {num_groups} groups"
         )
         logger.info(
             f"GPU status after deserialization for rank {self.rank_coordinate}:\n{get_gpu_status()}"