feat: support Qwen3-MoE weight conversion and dtype-only IPC grouping - #113
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the Awex SGLang adapter and plugin to support colocated weight transfer, metadata collection, and memory management, alongside adding support for Qwen3MoE and fixing several tensor conversion, sharding, and synchronization issues. The code review highlights critical issues that need to be addressed, including potential serialization errors when returning custom metadata objects via FastAPI, missing None checks for process handles and KV store lookups, and type mismatch bugs when handling string-serialized torch.dtype fields during chunked transfers. Additionally, the reviewer suggests adding divisibility checks for robust tensor splitting and removing unused attributes to clean up dead code.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| self._timeout_s: float = 120.0 | ||
| self._http_client = None |
There was a problem hiding this comment.
The self._http_client attribute is initialized and closed, but never used anywhere else in the class. Remove it to clean up dead code and reduce complexity.
| self._timeout_s: float = 120.0 | |
| self._http_client = None | |
| self._timeout_s = timeout_s |
19793a8 to
024c007
Compare
| self._adapter = None | ||
|
|
||
| def bind(self) -> None: | ||
| methods = [ |
There was a problem hiding this comment.
Implement the reader-facing worker task protocol
This bridge exposes dedicated metadata, initialization, and update RPCs, but AWEX already dispatches InferParamMetaResolver._get_model_param_info, WeightsReader._init_in_tp_worker, and _update_parameters_in_tp_worker through SGLangEngine.execute_task_in_model_worker(). The child process only needs to resolve the task and inject model plus a complete model_context. Because this bridge does not implement that protocol and instead owns conversion and updates itself, it bypasses converter selection, backend selection, validation, and memory sequencing in the standard reader. Replace these dedicated update RPCs with generic task dispatch and keep conversion and updates in WeightsReader/WorkerWeightsReader.
| "num_engines": 1, | ||
| } | ||
|
|
||
| def _unfuse_params( |
There was a problem hiding this comment.
Keep parameter conversion in the converter registry
_unfuse_params() creates a second conversion path outside get_infer_weights_converter. Model-specific SGLang converters therefore never run here: for example, the word_embeddings -> embed_tokens normalization added in ling_linear.py is skipped, as are MoE layouts and device-specific transpose/FP8 handling. This can make inference metadata names or shapes diverge from writer metadata. The plugin should pass model and model_context to the AWEX task and let the registry-selected converter perform conversion.
| self._timeout_s = timeout_s | ||
|
|
||
| from awex.meta.meta_server import MetaServerClient | ||
| infer_meta_key = f"infer_params_meta_{pair_name}" |
There was a problem hiding this comment.
Use the existing reader/writer handshake
This path introduces pair-suffixed metadata keys, then reads colocate_weights_rank* and writes colocate_done_rank*. The existing WeightsReader/NCCLWeightsWriter protocol uses infer_params_meta, training_params_meta, training_serialized_weights_{ip}_{device}_{step}, weights_update_finished*, and write_finished*. The two sides never rendezvous on the same keys, so this path still times out after the initialization errors are fixed. Do not duplicate the transfer protocol in the plugin; invoke the existing NCCLWorkerWeightsReader through model_context.
| world_size = tp_size * pp_size | ||
| global_rank = int(getattr(self._scheduler, "tp_rank", 0)) | ||
|
|
||
| return { |
There was a problem hiding this comment.
Provide one complete model_context
_build_rank_info() immediately passes this dictionary to get_sglang_rank_info(), which unconditionally reads attn_dp_rank and local_rank. Neither key is present, so the first metadata call raises KeyError on attn_dp_rank. Do not assemble a reduced context inside the conversion adapter. The SGLang child-process interface should build the complete model_context once from the real process groups and provide it to every AWEX reader task.
| def main(): | ||
| from awex.engine.sglang.plugin import register_awex_sglang_plugin | ||
|
|
||
| register_awex_sglang_plugin() |
There was a problem hiding this comment.
Register the advertised HTTP routes
The launcher patches run_scheduler_process and then starts SGLang, but nothing in the repository calls register_awex_endpoints(). As a result, none of the documented /awex/* routes are attached to the FastAPI application and this control plane is unreachable. The primary integration should use the generic child-process task protocol; if the HTTP layer remains, attach it explicitly during app construction and add an integration test that asserts the routes exist.
| infer_meta_key = f"infer_params_meta_{pair_name}" | ||
| train_meta_key = f"training_params_meta_{pair_name}" | ||
|
|
||
| meta_client = MetaServerClient(kv_store_url) |
There was a problem hiding this comment.
Construct MetaServerClient with address and port
MetaServerClient.__init__ requires (address, port), but this call and the second call in the execution path pass only kv_store_url. The initialization RPC therefore raises TypeError before reading any metadata. The better fix is to remove duplicate meta-client ownership from the plugin and reuse the reader. If this path remains temporarily, parse and pass the host and port separately and cover initialization with a CPU-only test.
| step_id=version, | ||
| ) | ||
|
|
||
| done_key = f"colocate_done_rank{paired_train_rank}_{version}" |
There was a problem hiding this comment.
Close IPC mappings before acknowledging completion
This writes colocate_done before deleting deserialized_weights and group_shared, synchronizing, and running garbage collection. Once the training side observes the acknowledgement, it may release and reallocate the exported storage while inference still holds CUDA IPC mappings. That reintroduces the stale-mapping and illegal-memory-access failure that this PR fixes in the canonical NCCLWorkerWeightsReader path. Release every IPC reference and synchronize before publishing completion, or route the update through the existing worker reader.
| gc.collect() | ||
| torch.cuda.empty_cache() | ||
|
|
||
| logger.info( |
There was a problem hiding this comment.
Flush SGLang caches after an in-place weight update
The parameters have been modified in place, but this path completes without calling scheduler.flush_cache(). WorkerWeightsReader.update_weights() flushes the cache after every update and verifies that it succeeded. Bypassing that step allows prefix or KV cache entries generated with the old weights to be reused under the new model version. Route the update through the worker reader instead of completing it inside this adapter.
| ) | ||
|
|
||
| os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False" | ||
| from awex.transfer.nccl_comm import init_weights_update_group |
There was a problem hiding this comment.
Import the process-group helper from its owning module
init_weights_update_group is defined in awex.util.process_group; awex.transfer.nccl_comm does not export it. After the meta-client call is fixed, this import still raises ImportError and prevents the colocate group from being created. Prefer dispatching NCCLWorkerWeightsReader.initialize() through the generic worker-task protocol so the existing reader owns the process-group lifecycle instead of duplicating initialization here.
|
Pushed Problem: the colocate adapter's inline Fix: route adapter weight naming through the per-model converter registry instead of ad-hoc unfusing:
Added CPU name-contract tests ( |
|
Follow-up The strategy's expert branch describes the fused layout (experts stacked along dim0, EP splitting that dim). Once experts are expanded to per-expert HF names, EP placement is already encoded in the global expert id inside the name, so re-applying Rules now applied for expanded
Fused stacked names and |
|
E2E validation complete — the expert-expansion fix ( Setup: Qwen3-30B-A3B (MoE, 128 experts), GSM8K PPO, 1 node, train Two runs, both finished cleanly (exit code 0, no errors):
Key points:
This closes the E2E gap noted in my previous comment. |
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.
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.
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.
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.)
1858790 to
2d3b774
Compare
2 participants
What does this PR do?
Two self-contained additions on top of #112 (only the last 2 commits are new; the diff collapses automatically once #112 merges).
feat: register
Qwen3MoeForCausalLMmcore converterAWEX had no name mapping for Qwen3-MoE, so the model could not go through colocate transfer. Differences from the bailing converters:
qwen3_moeserves canonical attention names (self_attn.{q,k,v,o}_proj+q/k_norm), so names are kept verbatim instead of the bailing-flavored renames.linear_qkvis split with GQA-aware group strides: the base equal-thirds split only holds when q/k/v head counts match, which is not the case under GQA (Qwen3-30B-A3B: 32 q-heads / 4 kv-heads). Row-count validation raises on any mismatch instead of mis-splitting silently.Registered via the standard module-level
CONFIGdict (same discovery pattern asling.py); no existing file is modified.fix: group IPC payload tensors by dtype only
group_tensors_by_shape_and_dtypekeyed buckets by(shape, dtype). When payload entries have unique lengths this degenerates to nearly one group (= one CUDA IPC handle) per parameter. Observed on a 30B MoE model: 4888 params -> 4636 groups, and the receiver'scudaIpcOpenMemHandlefails with a driver OOM despite >50 GB free GPU memory.Reconstruction slices each entry by element offset/size and reshapes from per-entry metadata, so same-shape packing was never a requirement. This change keys groups by dtype alone and flattens entries (
reshape(-1)) before concatenation; group count drops to O(dtypes x size-cap) and the handle-exhaustion failure disappears. Dense payloads benefit identically: fewer, larger IPC handles.Testing
reconstruct_tensors_from_groups) is offset/metadata-driven and unchanged.Related issues
None.
Does this PR introduce any user-facing change?
group_tensors_by_shape_and_dtypekeeps its signature; only the internal grouping key changes. Producer and consumer must run the same version, as with any change to the group metadata layout.