From 9fc96ada8e443fd081ce0e68f7db058c3cad7590 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:46:22 +0800 Subject: [PATCH 1/4] feat: register Qwen3MoeForCausalLM mcore converter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stock HF/SGLang qwen3_moe serves canonical attention names (self_attn.{q,k,v,o}_proj + q/k_norm), so keep them verbatim instead of the bailing-flavored renames, and split Megatron's fused linear_qkv with GQA-aware group strides — the base equal-thirds split only holds when q/k/v head counts match. Row-count validation raises on any mismatch instead of mis-splitting silently. --- awex/models/qwen3_moe.py | 106 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 awex/models/qwen3_moe.py diff --git a/awex/models/qwen3_moe.py b/awex/models/qwen3_moe.py new file mode 100644 index 0000000..41bc42a --- /dev/null +++ b/awex/models/qwen3_moe.py @@ -0,0 +1,106 @@ +# 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 + +logger = logging.getLogger(__name__) + + +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))) + 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, +} From b678b6b4aaf2b66fa1815cfadfec3bd3007d080f Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 15:46:27 +0800 Subject: [PATCH 2/4] fix: group IPC payload tensors by dtype only Shape-keyed grouping degenerates to nearly one group (= one CUDA IPC handle) per parameter when payload entries have unique lengths. Observed on a 30B MoE model: 4888 params -> 4636 groups, and the receiver's cudaIpcOpenMemHandle fails with a driver OOM despite >50GB free GPU memory. Reconstruction slices each entry by element offset/size and reshapes from per-entry metadata, so same-shape packing was never required. Key groups by dtype alone and flatten entries before concatenation; group count drops to O(dtypes x size-cap) and the handle-exhaustion failure disappears. --- awex/tests/test_tensor_utils.py | 39 ++++++++++++++++++--------------- awex/util/tensor_util.py | 22 +++++++++++++------ 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/awex/tests/test_tensor_utils.py b/awex/tests/test_tensor_utils.py index 452bd8f..9b26023 100644 --- a/awex/tests/test_tensor_utils.py +++ b/awex/tests/test_tensor_utils.py @@ -49,8 +49,9 @@ def test_single_tensor(self): assert len(result_groups) == 1 assert len(result_metadata) == 1 - # Check that the group contains the original tensor - assert torch.equal(result_groups[0], tensor) + # Check that the group contains the original tensor's data + # (groups are flattened 1-D concatenations since grouping is by dtype) + assert torch.equal(result_groups[0], tensor.reshape(-1)) # Check metadata metadata = result_metadata[0] @@ -73,8 +74,10 @@ def test_tensors_same_shape_and_dtype(self): assert len(result_groups) == 1 assert len(result_metadata) == 3 - # Check that all tensors are concatenated in one group - expected_concatenated = torch.cat([tensor1, tensor2, tensor3], dim=0) + # Check that all tensors are concatenated (flattened) in one group + expected_concatenated = torch.cat( + [tensor1.reshape(-1), tensor2.reshape(-1), tensor3.reshape(-1)], dim=0 + ) assert torch.equal(result_groups[0], expected_concatenated) # Check metadata for each tensor @@ -86,7 +89,7 @@ def test_tensors_same_shape_and_dtype(self): assert metadata["size"] == tensor1.numel() def test_tensors_different_shapes(self): - """Test function with tensors of different shapes.""" + """Test function with tensors of different shapes (same dtype).""" tensor1 = torch.randn(2, 3, dtype=torch.float32) tensor2 = torch.randn(4, 5, dtype=torch.float32) tensor3 = torch.randn(2, 3, dtype=torch.float32) @@ -94,21 +97,20 @@ def test_tensors_different_shapes(self): result_groups, result_metadata = group_tensors_by_shape_and_dtype(tensors) - assert len(result_groups) == 2 # Two different shapes + # Grouping is by dtype only: mixed shapes share one flattened group + assert len(result_groups) == 1 assert len(result_metadata) == 3 - # Check that tensors with same shape are grouped together - # tensor1 and tensor3 should be in one group - # tensor2 should be in another group - shape_groups = {} - for metadata in result_metadata: - shape = metadata["shape"] - if shape not in shape_groups: - shape_groups[shape] = [] - shape_groups[shape].append(metadata["original_index"]) + # Per-entry metadata keeps the original shapes for reconstruction + shapes = {m["original_index"]: m["shape"] for m in result_metadata} + assert shapes[0] == torch.Size([2, 3]) + assert shapes[1] == torch.Size([4, 5]) + assert shapes[2] == torch.Size([2, 3]) - assert set(shape_groups[torch.Size([2, 3])]) == {0, 2} - assert set(shape_groups[torch.Size([4, 5])]) == {1} + reconstructed = reconstruct_tensors_from_groups(result_groups, result_metadata) + assert torch.equal(reconstructed[0], tensor1) + assert torch.equal(reconstructed[1], tensor2) + assert torch.equal(reconstructed[2], tensor3) def test_tensors_different_dtypes(self): """Test function with tensors of different dtypes.""" @@ -163,7 +165,8 @@ def test_complex_tensor_shapes(self): result_groups, result_metadata = group_tensors_by_shape_and_dtype(tensors) - assert len(result_groups) == 2 # Two different shapes + # Grouping is by dtype only: all three share one flattened group + assert len(result_groups) == 1 assert len(result_metadata) == 3 # Verify reconstruction works correctly diff --git a/awex/util/tensor_util.py b/awex/util/tensor_util.py index 8087502..545586f 100644 --- a/awex/util/tensor_util.py +++ b/awex/util/tensor_util.py @@ -210,12 +210,16 @@ def group_tensors_by_shape_and_dtype( logger.info( f"Start to group tensors, total size: {total_size}, num tensors: {len(tensors)}" ) - # 1. Group tensors by shape and dtype - tensor_groups: Dict[ - Tuple[torch.Size, torch.dtype], List[Tuple[int, torch.Tensor]] - ] = {} + # 1. Group tensors by dtype only. Reconstruction slices each entry by + # element offset/size and reshapes from per-entry metadata, so same-shape + # packing is not required — and shape-keyed buckets degenerate to one + # group (= one CUDA IPC handle) per parameter on sparse delta payloads, + # where every parameter's values/indices have unique lengths. Observed on + # Qwen3-30B: 4888 params -> 4636 groups, receiver cudaIpcOpenMemHandle + # failing with driver OOM despite >50GB free. + tensor_groups: Dict[torch.dtype, List[Tuple[int, torch.Tensor]]] = {} for i, tensor in enumerate(tensors): - key = (tensor.shape, tensor.dtype) + key = tensor.dtype if key not in tensor_groups: tensor_groups[key] = [] tensor_groups[key].append((i, tensor)) @@ -241,7 +245,9 @@ def group_tensors_by_shape_and_dtype( if current_group_size > max_tensor_size: # Finalize current group and start new one # Use clone() to ensure a copy so caller can safely release original tensors - concatenated = torch.cat(current_group, dim=0).clone() + concatenated = torch.cat( + [t.reshape(-1) for t in current_group], dim=0 + ).clone() final_tensor_groups.append(concatenated) # Record metadata for tensors in this group offset_elements = 0 @@ -266,7 +272,9 @@ def group_tensors_by_shape_and_dtype( # Finalize any remaining group if current_group: # Use clone() to ensure a copy so caller can safely release original tensors - concatenated = torch.cat(current_group, dim=0).clone() + concatenated = torch.cat( + [t.reshape(-1) for t in current_group], dim=0 + ).clone() final_tensor_groups.append(concatenated) # Record metadata for tensors in this group offset_elements = 0 From 2b102f766d4f1bc60f87a5bd0e43acdb65293838 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Tue, 21 Jul 2026 16:55:44 +0800 Subject: [PATCH 3/4] fix: validate GQA head-count divisibility in qwen3-moe qkv split Address review: fail fast with clear messages when num_key_value_heads is not divisible by attn_tp_size or num_attention_heads is not divisible by num_key_value_heads, instead of relying on the downstream row-count mismatch error. --- awex/models/qwen3_moe.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/awex/models/qwen3_moe.py b/awex/models/qwen3_moe.py index 41bc42a..72742d6 100644 --- a/awex/models/qwen3_moe.py +++ b/awex/models/qwen3_moe.py @@ -58,6 +58,17 @@ def _split_gqa_qkv( 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 From 2d3b77485945bb78572002940a46829684f04472 Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Wed, 22 Jul 2026 17:03:17 +0800 Subject: [PATCH 4/4] fix: route SGLang weight naming through converter registry The base SGlangToHFWeightConverter split fused qkv into equal thirds, which is wrong for GQA models. The train side reports canonical per-expert HF names, so Qwen3-MoE expert weights in SGLang fused form matched nothing in the transfer plan. Key changes: - Make the base SGlangToHFWeightConverter qkv split GQA-aware (head-count-weighted sizes instead of equal thirds; fused paths and MHA behavior unchanged) - Register a Qwen3-MoE SGLang converter that unfuses qkv and accepts q_norm/k_norm layer norms - Extract reconstruct_ipc_weights into tensor_util, shared by the NCCL reader colocate receive path - Add CPU name-contract tests pinning the inference/train name set alignment for Qwen3-MoE, including EP global expert ids (Adapter-side changes from the original commit are dropped along with the removed engine/sglang integration layer.) --- awex/converter/sglang_converter.py | 27 ++- awex/models/qwen3_moe.py | 35 +++- awex/reader/nccl_reader.py | 20 +- awex/tests/test_qwen3_moe_sglang_converter.py | 194 ++++++++++++++++++ awex/util/tensor_util.py | 33 +++ 5 files changed, 283 insertions(+), 26 deletions(-) create mode 100644 awex/tests/test_qwen3_moe_sglang_converter.py diff --git a/awex/converter/sglang_converter.py b/awex/converter/sglang_converter.py index 4e2925f..8953059 100644 --- a/awex/converter/sglang_converter.py +++ b/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: diff --git a/awex/models/qwen3_moe.py b/awex/models/qwen3_moe.py index 72742d6..e26b8f4 100644 --- a/awex/models/qwen3_moe.py +++ b/awex/models/qwen3_moe.py @@ -20,10 +20,36 @@ 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). @@ -48,9 +74,7 @@ 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 - ) + return int(self.hf_config.hidden_size // self.hf_config.num_attention_heads) def _split_gqa_qkv( self, parameter: torch.Tensor @@ -82,9 +106,7 @@ def _split_gqa_qkv( ) 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:] - ) + 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:] ) @@ -114,4 +136,5 @@ def _convert_attention_param( CONFIG = { "model_name": "Qwen3MoeForCausalLM", "mcore_converter": _build_mcore_converter_qwen3_moe, + "sglang_converter": SGlangToHFWeightConverterQwen3Moe, } diff --git a/awex/reader/nccl_reader.py b/awex/reader/nccl_reader.py index c548bd0..76c5dba 100644 --- a/awex/reader/nccl_reader.py +++ b/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()}" diff --git a/awex/tests/test_qwen3_moe_sglang_converter.py b/awex/tests/test_qwen3_moe_sglang_converter.py new file mode 100644 index 0000000..ff5a693 --- /dev/null +++ b/awex/tests/test_qwen3_moe_sglang_converter.py @@ -0,0 +1,194 @@ +# 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. + +"""Name-contract tests for the Qwen3-MoE SGLang->HF weight converter. + +The colocate transfer plan matches inference-side and train-side parameters +by name. The train side (Megatron converter) emits canonical per-parameter +HF names, so the inference side must expand SGLang fused parameters — +including MoE experts — to the exact same name set. These tests pin that +contract on CPU without requiring SGLang or Megatron. +""" + +from types import SimpleNamespace + +import torch + +from awex.models.qwen3_moe import SGlangToHFWeightConverterQwen3Moe +from awex.models.registry import get_infer_weights_converter + +# Tiny Qwen3-MoE-like geometry: GQA with 8 query heads and 2 KV heads. +NUM_HEADS = 8 +NUM_KV_HEADS = 2 +HEAD_DIM = 4 +HIDDEN = 16 +MOE_INTERMEDIATE = 8 +NUM_EXPERTS = 4 + + +def _model_config(): + return SimpleNamespace( + num_attention_heads=NUM_HEADS, + num_key_value_heads=NUM_KV_HEADS, + head_dim=HEAD_DIM, + hidden_size=HIDDEN, + num_experts=NUM_EXPERTS, + architectures=["Qwen3MoeForCausalLM"], + ) + + +def _rank_info(tp_rank=0, ep_rank=0): + return SimpleNamespace(tp_rank=tp_rank, ep_rank=ep_rank) + + +def _infer_engine_config(tp_size=1, ep_size=1): + return SimpleNamespace(tp_size=tp_size, ep_size=ep_size, device_backend="cuda") + + +def _make_converter(tp_size=1, ep_size=1, tp_rank=0, ep_rank=0): + return SGlangToHFWeightConverterQwen3Moe( + _model_config(), + _infer_engine_config(tp_size=tp_size, ep_size=ep_size), + _rank_info(tp_rank=tp_rank, ep_rank=ep_rank), + ) + + +def _sglang_named_params(num_local_experts=NUM_EXPERTS): + """Parameter names/shapes as exposed by SGLang for one decoder layer.""" + qkv_rows = (NUM_HEADS + 2 * NUM_KV_HEADS) * HEAD_DIM + return { + "model.embed_tokens.weight": torch.randn(32, HIDDEN), + "model.layers.0.self_attn.qkv_proj.weight": torch.randn(qkv_rows, HIDDEN), + "model.layers.0.self_attn.o_proj.weight": torch.randn( + HIDDEN, NUM_HEADS * HEAD_DIM + ), + "model.layers.0.self_attn.q_norm.weight": torch.randn(HEAD_DIM), + "model.layers.0.self_attn.k_norm.weight": torch.randn(HEAD_DIM), + "model.layers.0.input_layernorm.weight": torch.randn(HIDDEN), + "model.layers.0.post_attention_layernorm.weight": torch.randn(HIDDEN), + "model.layers.0.mlp.gate.weight": torch.randn(NUM_EXPERTS, HIDDEN), + "model.layers.0.mlp.experts.w13_weight": torch.randn( + num_local_experts, 2 * MOE_INTERMEDIATE, HIDDEN + ), + "model.layers.0.mlp.experts.w2_weight": torch.randn( + num_local_experts, HIDDEN, MOE_INTERMEDIATE + ), + "model.norm.weight": torch.randn(HIDDEN), + "lm_head.weight": torch.randn(32, HIDDEN), + } + + +def _expected_hf_names(expert_ids): + """Canonical HF names the Megatron train-side converter reports.""" + names = { + "model.embed_tokens.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.q_norm.weight", + "model.layers.0.self_attn.k_norm.weight", + "model.layers.0.input_layernorm.weight", + "model.layers.0.post_attention_layernorm.weight", + "model.layers.0.mlp.gate.weight", + "model.norm.weight", + "lm_head.weight", + } + for expert_id in expert_ids: + names.add(f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight") + names.add(f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight") + names.add(f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight") + return names + + +def test_registry_resolves_qwen3_moe_sglang_converter(): + converter = get_infer_weights_converter( + "sglang", + "Qwen3MoeForCausalLM", + _model_config(), + _rank_info(), + _infer_engine_config(), + ) + assert isinstance(converter, SGlangToHFWeightConverterQwen3Moe) + + +def test_convert_param_name_set_matches_train_side_contract(): + converter = _make_converter() + converted_names = set() + for name, param in _sglang_named_params().items(): + for hf_name, _ in converter.convert_param(name, param): + converted_names.add(hf_name) + assert converted_names == _expected_hf_names(range(NUM_EXPERTS)) + + +def test_convert_param_expert_names_use_global_ids_with_ep(): + # With ep_size=2 each rank holds half the experts; converted names must + # carry global expert ids offset by ep_rank. + num_local = NUM_EXPERTS // 2 + converter = _make_converter(ep_size=2, ep_rank=1) + converted_names = set() + for name, param in _sglang_named_params(num_local_experts=num_local).items(): + for hf_name, _ in converter.convert_param(name, param): + converted_names.add(hf_name) + assert converted_names == _expected_hf_names(range(num_local, NUM_EXPERTS)) + + +def test_qkv_split_is_gqa_aware(): + converter = _make_converter() + q = torch.randn(NUM_HEADS * HEAD_DIM, HIDDEN) + k = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN) + v = torch.randn(NUM_KV_HEADS * HEAD_DIM, HIDDEN) + fused = torch.cat([q, k, v], dim=0) + + result = dict( + converter.convert_param("model.layers.0.self_attn.qkv_proj.weight", fused) + ) + assert torch.equal(result["model.layers.0.self_attn.q_proj.weight"], q) + assert torch.equal(result["model.layers.0.self_attn.k_proj.weight"], k) + assert torch.equal(result["model.layers.0.self_attn.v_proj.weight"], v) + + +def test_qkv_split_rejects_indivisible_rows(): + converter = _make_converter() + bad = torch.randn((NUM_HEADS + 2 * NUM_KV_HEADS) * HEAD_DIM + 1, HIDDEN) + try: + converter.convert_param("model.layers.0.self_attn.qkv_proj.weight", bad) + except ValueError as e: + assert "not divisible" in str(e) + else: + raise AssertionError("expected ValueError for indivisible qkv rows") + + +def test_expert_split_values_match_fused_slices(): + converter = _make_converter() + w13 = torch.randn(NUM_EXPERTS, 2 * MOE_INTERMEDIATE, HIDDEN) + w2 = torch.randn(NUM_EXPERTS, HIDDEN, MOE_INTERMEDIATE) + + gate_up = dict( + converter.convert_param("model.layers.0.mlp.experts.w13_weight", w13) + ) + down = dict(converter.convert_param("model.layers.0.mlp.experts.w2_weight", w2)) + + for expert_id in range(NUM_EXPERTS): + prefix = f"model.layers.0.mlp.experts.{expert_id}" + assert torch.equal( + gate_up[f"{prefix}.gate_proj.weight"], w13[expert_id, :MOE_INTERMEDIATE] + ) + assert torch.equal( + gate_up[f"{prefix}.up_proj.weight"], w13[expert_id, MOE_INTERMEDIATE:] + ) + assert torch.equal(down[f"{prefix}.down_proj.weight"], w2[expert_id]) diff --git a/awex/util/tensor_util.py b/awex/util/tensor_util.py index 545586f..7393acc 100644 --- a/awex/util/tensor_util.py +++ b/awex/util/tensor_util.py @@ -350,6 +350,39 @@ def reconstruct_tensors_from_groups( return result_tensors +@torch.no_grad() +def reconstruct_ipc_weights( + serialized_weights: bytes, + ipc_backend: str = "cuda", + device_id=None, +) -> Tuple[Dict[str, torch.Tensor], int]: + """Deserialize IPC-shared weight groups and rebuild the name->tensor map. + + Shared by the colocate receive paths (NCCL reader and the SGLang + adapter): deserialize the pickled (groups, metadata, names) payload, + synchronize the device, and slice per-parameter tensors back out of the + dtype-grouped buffers. + + Args: + serialized_weights: ForkingPickler payload produced by the sender. + ipc_backend: "cuda" for CUDA IPC handles; "cpu"/"npu" for + plain pickled tensors that must be moved to ``device_id``. + device_id: Target device for cpu/npu backends. + + Returns: + Tuple of (name -> tensor mapping, number of shared groups). + """ + if 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()) + return dict(zip(names, tensors)), len(group_shared) + + def release_tensors(tensors: Union[Iterable[torch.Tensor], torch.Tensor]): if not isinstance(tensors, Iterable): tensors = [tensors]