feat: add SGLang colocate plugin and colocate P2P transfer fixes - #112
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 colocated weight transfer support for SGLang, adding an adapter, a launcher, and a plugin to integrate SGLang with AWEX. It also includes fixes for MLA projection packing, GQA/MHA interleaved format handling, and regression tests for sharding bugs. The review feedback highlights several critical improvements: ensuring cross-platform compatibility by using device_util instead of direct CUDA calls, fixing a potential OOM or TypeError in chunk size calculation when slice bounds are None, re-introducing a contiguity check on received tensors to prevent runtime crashes, and cleaning up dead or redundant 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.
|
We have an |
|
Done in e6bbffe — moved to Also pushed in the same batch: c87d11f (English comments / no internal issue tags), f8a57fd (restore non-contiguous P2P handling with per-chunk staging + unit tests), and 2d94a5d (portability/robustness items from the automated review). |
chaokunyang
left a comment
chaokunyang
left a comment
There was a problem hiding this comment.
This needs to be reworked around the existing AWEX reader boundary. The SGLang integration should expose a scheduler/model-context bridge that satisfies AWEX's reader call protocol; it should not implement a second weight-conversion and colocate-update pipeline inside the SGLang plugin.
| f"No train weight metadata in kv store for key {train_meta_key}" | ||
| ) | ||
|
|
||
| builder = TransferPlanBuilder( |
There was a problem hiding this comment.
Blocking: AwexSGLangAdapter is a second reader, which reverses AWEX's dependency boundary. WeightsReader dispatches arbitrary AWEX callables through execute_task_in_model_worker; InferParamMetaResolver, WorkerWeightsReader, and NCCLWorkerWeightsReader already own metadata conversion and validation, parameter conversion, device-rank mapping, process-group setup, IPC deserialization, transfer planning, and update execution. The SGLang-specific layer should only execute those callables in the scheduler/model-worker subprocess, inject model plus the complete model_context, and return per-worker results. Please remove the adapter-owned metadata/conversion/update endpoints and implement that existing dispatch contract.
There was a problem hiding this comment.
Agreed — this file (together with plugin.py and launcher.py) has been removed from the PR. It was an early experimental integration that predates the current colocate protocol and is not referenced anywhere: the production SGLang integration lives on the training-framework side, which drives the native NCCLWorkerWeightsReader(enable_colocate_mode=True) through the official writer protocol (training_serialized_weights_* / weights_update_finished_*), so no second reader pipeline exists. This PR is now scoped to the transfer/converter/sharding fixes only.
| "init_colocate_weight_update must be called first" | ||
| ) | ||
| paired_train_rank = self._infer_to_train_device_mapping[self._transfer_rank] | ||
| kv_key = f"colocate_weights_rank{paired_train_rank}_{version}" |
There was a problem hiding this comment.
Blocking: this cannot interoperate with the current AWEX writer. NCCLWeightsWriter._write_weights_in_colocate_mode publishes training_serialized_weights_{ip}_{device}_{step} and waits for weights_update_finished_{ip}_{device}_{step}; this adapter waits on colocate_weights_rank... and later writes colocate_done_rank..., which no writer in this PR uses. This endpoint path will time out here and/or leave the writer waiting forever. This is a symptom of the duplicated update protocol; please route through the existing reader protocol rather than introducing separate meta-server keys.
There was a problem hiding this comment.
Removed along with the adapter — the private colocate_weights_rank{rank}_{version} key protocol is gone from this PR. The production path uses the official writer keys, driven from the training-framework side against the unmodified AWEX reader.
|
|
||
| def _build_rank_info(self) -> RankInfo: | ||
| model_context = self._get_model_context() | ||
| return get_sglang_rank_info(model_context, engine_rank=0) |
There was a problem hiding this comment.
Blocking: _get_model_context() does not satisfy AWEX's model-context contract. get_sglang_rank_info unconditionally reads attn_dp_rank and local_rank, while WorkerWeightsReader unconditionally reads infer_engine_config; this context supplies none of them. /awex/report_weight_meta already fails on the missing rank fields, and a generic reader task would fail during worker construction. Build the complete context once in the subprocess bridge and inject it into dispatched AWEX callables instead of reconstructing a partial context inside the adapter.
There was a problem hiding this comment.
Removed — with the adapter gone this incomplete _get_model_context() no longer exists in the PR. The production integration builds the full model context (including attn_dp_rank, local_rank, infer_engine_config) when instantiating the native reader.
| import signal | ||
|
|
||
| import psutil | ||
| from sglang.srt.managers.scheduler import Scheduler, configure_scheduler |
There was a problem hiding this comment.
This cannot be a drop-in replacement for SGLang's scheduler subprocess. The SGLang revision referenced in this repository launches run_scheduler_process with eight positional arguments, while this function requires ten; current upstream has changed the lifecycle again and no longer exports configure_scheduler. Scheduler workers therefore fail before reporting readiness, and copying the private runner also bypasses upstream setup/teardown as it evolves. Leave SGLang's runner in control and register only the generic AWEX callable/model-context hook through the supported plugin mechanism.
There was a problem hiding this comment.
Resolved by removal — plugin.py is no longer part of this PR (see the thread on adapter.py for context).
|
|
||
| register_awex_sglang_plugin() | ||
|
|
||
| from sglang.launch_server import main as sglang_main |
There was a problem hiding this comment.
sglang.launch_server exports run_server, not main, in both the SGLang revision referenced by this repository and current upstream. The documented python -m awex.engine.sglang.launcher command raises ImportError here before the server starts. Please use SGLang's supported CLI/plugin entrypoint instead of importing a nonexistent main.
There was a problem hiding this comment.
Resolved by removal — launcher.py is no longer part of this PR.
| [layer_cache["q_a_proj"], layer_cache["kv_a_proj"]], dim=0 | ||
| ) | ||
| else: | ||
| fused_tensor = pack_fused_qkv_a_proj_for_tp( |
There was a problem hiding this comment.
This branch assumes the fused tensor will be TP-sliced later, but LinearMLAShardingMixin.get_sharding_strategy() always declares attention.fused_qkv_a_proj_with_mqa.weight as NO_SHARDING. ParamMetaResolver therefore treats the rank-0 result as a complete replica and the transfer plan sends it unchanged to every SGLang TP rank. For training attention TP > 1, this interleaves rows as [q0;kv0;q1;kv1;...] although each recipient expects full [q_all;kv_all], corrupting the MLA projection. Keep the canonical concatenation unless the metadata and transfer plan are changed to perform a real split.
There was a problem hiding this comment.
Fixed in 6a897de — the interleave branch is removed and the canonical [q_a; kv_a] concatenation is now emitted unconditionally, matching the NO_SHARDING declaration (both inputs are already gathered to full tensors at this point, so this is correct for any train attn-TP size). The now-unused pack_fused_qkv_a_proj_for_tp helper is dropped as well.
| ) # Q extends to 2, k v use hidden size | ||
|
|
||
| if actual_size == expected_replicated_size: | ||
| if actual_size == expected_compact_size: |
There was a problem hiding this comment.
For MHA, expected_compact_size and expected_replicated_size are both 3 * hidden_size, so this changed ordering always interprets the weight as interleaved. transform_mcore_qkv_bias() still checks the replicated case first and interprets the identically shaped bias as contiguous [Q; K; V]. A fused QKV linear uses the same output-row ordering for weight and bias, so biased models now permute them inconsistently. Use one explicit layout signal/shared transform for both; shape alone cannot distinguish these MHA layouts.
There was a problem hiding this comment.
Fixed in 6a897de — transform_mcore_qkv_bias now checks the compact layout before the replicated one, mirroring transform_mcore_qkv_weight, so in the MHA case (where the two sizes coincide) weight and bias resolve to the same interleaved interpretation. A shared explicit layout signal would indeed be more robust long-term; left as a follow-up since it touches the converter API surface.
| # is ``embed_tokens``. Normalize here so the transfer-plan key sets on | ||
| # both sides align; the base converter has no rule for ``word_embeddings`` | ||
| # and would otherwise pass it through unchanged. | ||
| if name == "model.word_embeddings.weight": |
There was a problem hiding this comment.
This rename also changes the sharding decision. InferParamMetaResolver asks the strategy about the converted name, while ShardingStrategy.get_sharding_strategy() recognizes embeddings only when the name contains embedding. After this becomes model.embed_tokens.weight, Bailing's replicated-embedding override is bypassed and a full SGLang vocabulary tensor is declared TP_SHARDING. With DP attention and TP > 1, normal initialization reports an infer shape of roughly TP×V and fails metadata validation (or leaves later replicas outside the train extent in debug mode). Normalize the key without losing embedding-specific sharding semantics.
There was a problem hiding this comment.
Fixed in 6a897de — the sharding fallback in param_sharding.py now recognizes embed_tokens as an embedding, so the normalized canonical name keeps the replicated-embedding semantics (Bailing's get_embedding_sharding_strategy override included) instead of falling through to generic TP_SHARDING.
| chunk_recv_tensor_pairs = [] | ||
| chunk_clone_bytes = 0 | ||
|
|
||
| for mapped_peer_rank, ops in send_per_peer.items(): |
There was a problem hiding this comment.
AWEX_CHUNK_MB is not a bound on the allocation made here. step_size is computed from one peer's sample, then that many tensors are cloned for every send peer before transfer, so 16 active peers can allocate roughly 16× the requested budget. train_slice_context also lives across all chunks and retains every dense non-contiguous send slice, and the swallowed import of nonexistent dtype_to_size undercounts FP32/cast-to-FP32 sends. This defeats the OOM protection this path is intended to provide. Build chunks from aggregate actual wire bytes and release the slice cache at each chunk boundary, or transfer peer-by-peer without pre-cloning all peers.
There was a problem hiding this comment.
Fixed in 5109a75 — step_size is now derived from the aggregate wire bytes of one op index summed across ALL send/recv peers (computed with the receiver dtype, which the clones are cast to), so N active peers no longer allocate N x the budget. The densified-send slice cache is scoped per chunk and the self-copy cache is released right after the copy phase, so staging buffers no longer accumulate across the transfer. The swallowed dtype_to_size import is gone.
…ransfer Three fixes for colocate P2P at large asymmetric reshard scale (32 ranks, PP4 -> PP1), plus tracing: - Chunked transfer (AWEX_CHUNK_MB): split each peer's ops into size-bounded chunks with cross-rank aligned step_size/n_chunks (MIN/MAX all-reduced) so no rank exits the chunk loop early while its peers still wait. - Per-peer batch_isend_irecv with per-peer GPU drain: submitting a whole recursive-partition half's ops in one batch floods NCCL P2P channels and the GPU drain never completes. Walking peers in ascending order caps in-flight P2P at O(1) peer; phases are single-direction so serialization cannot form a wait cycle. - Wire-size parity: cast the send clone to the receiver's shard dtype. When the train-side shard is bf16 but the receiver posts an fp32-sized irecv (e.g. mlp.gate.weight), the byte counts differ and the receiver blocks forever. - AWEX_P2P_TRACE per-peer progress log for diagnosing peer-level stalls; execute_tensors_to_copy wrapped in no_grad.
- Declare mcore Lightning attention qkv as TP_SHARDING when attn_tp_size > 1: convert_qkv_weight_along_tp_attention hands each rank only its shard, so declaring NO_SHARDING made the transfer plan cover infer tp_rank 0 only and left Lightning qkv on tp_rank > 0 all-zero. Regression test included. - Adapt the mcore converter for BailingMoeV2.5 colocate transfer and normalize SGLang's word_embeddings name to the canonical embed_tokens so both sides' transfer-plan key sets align.
cudaIpcCloseMemHandle runs in the tensor deleter, which only fires when refcounts/GC actually release the IPC-imported tensors. If that close lags past the train side's release + empty_cache + realloc (train acts right after weights_update_finished), the stale IPC mapping overlaps train's fresh allocations and either side hits an illegal memory access at a drifting location. gc.collect + synchronize before the signal.
Skip the step-2 consistency compare when fewer than two history entries exist, instead of crashing on single-entry histories.
torch.distributed P2P ops require dense buffers; KDA/MLA converter outputs are frequently non-contiguous views. Clone sends into contiguous buffers and stage non-contiguous receives through a dense buffer with an explicit copy-back after wait. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Address review: rewrite test_attn_tp_sharding_cp_bug.py docs/comments in English and replace internal issue references in code comments with self-contained failure-mode descriptions.
…yout detection - emit the plain [q_a;kv_a] concat for fused_qkv_a_proj_with_mqa unconditionally: the param is declared NO_SHARDING, so every SGLang TP rank receives the full tensor and any infer-TP interleave scrambles rows; drop the now-unused pack_fused_qkv_a_proj_for_tp helper - check the compact GQA/MHA layout before the replicated one in transform_mcore_qkv_bias, mirroring transform_mcore_qkv_weight, so MHA (where both sizes coincide) resolves weight and bias identically - treat embed_tokens as an embedding in the sharding fallback so the canonical HF name keeps replicated-embedding semantics instead of falling through to generic TP_SHARDING
- derive step_size from the aggregate wire bytes of one op index summed across ALL send/recv peers (using the receiver dtype, which the clones are cast to), instead of sampling a single peer's op — with N active peers the old estimate allocated up to N x the requested budget - scope the densified-send slice cache to each chunk and release the self-copy cache right after the copy phase, so staging buffers no longer accumulate across the whole transfer - drop the try/except around the removed dtype_to_size import that silently fell back to 2-byte elements for fp32 sends
2d94a5d to
5109a75
Compare
|
Thanks for the thorough review. The branch has been reworked (force-pushed, now at 5109a75): SGLang integration layer removed. Review fixes added on top of the retained transfer/converter/sharding commits:
Scope is now 9 files, +839/−85, limited to transfer/converter/sharding fixes. Full test suite passes (147 passed, 4 skipped). |
…#113) ## 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 `Qwen3MoeForCausalLM` mcore converter AWEX had no name mapping for Qwen3-MoE, so the model could not go through colocate transfer. Differences from the bailing converters: - Stock HF/SGLang `qwen3_moe` serves canonical attention names (`self_attn.{q,k,v,o}_proj` + `q/k_norm`), so names are kept verbatim instead of the bailing-flavored renames. - Megatron's fused `linear_qkv` is 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 `CONFIG` dict (same discovery pattern as `ling.py`); no existing file is modified. ### fix: group IPC payload tensors by dtype only `group_tensors_by_shape_and_dtype` keyed 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's `cudaIpcOpenMemHandle` fails 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 - End to end on Qwen3-30B-A3B colocate RL training: weight-sync correctness verified bitwise against a full-transfer snapshot across multi-step runs. - The reconstruction path (`reconstruct_tensors_from_groups`) is offset/metadata-driven and unchanged. ## Related issues None. ## Does this PR introduce any user-facing change? - [ ] Does this PR introduce any public API change? - [ ] Does this PR introduce any binary protocol compatibility change? `group_tensors_by_shape_and_dtype` keeps 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.
2 participants
What does this PR do?
Production-validated features and fixes for AWEX colocate weight transfer,
from months of colocate RL training deployments on 30B-class MoE models
(Megatron PP training + SGLang inference).
feat: SGLang colocate plugin
Wires SGLang inference workers as AWEX weight receivers in colocate mode:
sglang_plugin.pyhooks model-worker task executionsglang_awex_adapter.pymaps AWEX shard metadata onto SGLang's fused parameter layoutsglang_awex_launcher.pybootstraps the transfer group alongside the SGLang serverfix: chunked per-peer P2P schedule and wire-size parity
Three fixes for colocate P2P at large asymmetric reshard scale (32 ranks, PP4 -> PP1):
AWEX_CHUNK_MB, opt-in): splits each peer's ops into size-bounded chunks;step_size/n_chunksare MIN/MAX all-reduced so no rank exits the chunk loop early while peers still wait.batch_isend_irecv: submitting a whole recursive-partition half's ops in one batch floods NCCL P2P channels and the GPU drain never completes. Walking peers in ascending order caps in-flight P2P at O(1) peer; each phase is single-direction, so serialization cannot form a wait cycle.irecv(e.g.mlp.gate.weight), the byte counts differ and the receiver blocks forever.AWEX_P2P_TRACEper-peer progress log for diagnosing peer-level stalls.fix: sharding and converter corrections
TP_SHARDINGwhenattn_tp_size > 1. It was declaredNO_SHARDING, so the transfer plan covered infer tp_rank 0 only and Lightning qkv on tp_rank > 0 stayed all-zero. Regression test included (awex/tests/test_attn_tp_sharding_cp_bug.py).word_embeddingsto canonicalembed_tokensso both sides' transfer-plan key sets align.fix: force IPC close before
weights_update_finishedcudaIpcCloseMemHandleruns in the tensor deleter and can lag past the train side's release + realloc, so the stale IPC mapping overlaps fresh allocations and triggers an illegal memory access at a drifting location. Forcegc.collect()+synchronize()before signalling.fix: defensive guards in update-history compare
Skip the step-2 consistency compare with fewer than two history entries; early-return when
num_updates == 0.Testing
pytest awex/tests/test_attn_tp_sharding_cp_bug.py— 3 passedRisk notes
AWEX_CHUNK_MBis set (the chunked path is opt-in via env).Related issues
None.
Does this PR introduce any user-facing change?
New opt-in env vars only:
AWEX_CHUNK_MB,AWEX_CHUNK_OPS,AWEX_P2P_TRACE. Default behavior is unchanged when they are unset.