fix(v2): build complete MoE inference metadata for weight transfer - #1595
fix(v2): build complete MoE inference metadata for weight transfer#1595Le8r0nJames wants to merge 2 commits into
Conversation
Two defects kept most of an MoE model's experts out of the transfer plan, and both are silent: the shapes are right either way, so the run proceeds and only the train/inference logprob divergence shows that inference is holding stale weights. SGLang stores routed experts fused, and under expert parallelism each rank holds num_experts // ep_size of them, so the fused tensor's leading index is a local id. It was used directly as the global HuggingFace name, so with four expert-parallel ranks every rank advertised experts 0..31 of 128 and the rest were never named. get_local_shard_parameters shares the same naming path, so the payload has to be built with the rank info resolved as well; otherwise the metadata and the data disagree about which expert a tensor is. The gateway then collected one metadata entry per (name, rank) and passed the flat list on. The transfer plan is keyed by name alone, so all but one entry per name was shadowed. The training side already merged by name; do the same for inference, which also stops the payload growing with the inference world size. Verified on a 2-node separation run of a 128-expert model with ep_size=4: inference keys match the training side exactly, the aggregated metadata drops to a fraction of its previous size, and the behaviour importance weight stays at ~1.0 across four weight transfers instead of drifting.
| prefix = name.replace(".w13_weight", "") | ||
| result = [] | ||
| ffn_hidden = tensor.shape[1] // 2 | ||
| id_offset = self._expert_id_offset(tensor.shape[0]) |
There was a problem hiding this comment.
SGLang uses fused moe. It is recommended to confirm whether the fused tensor can still meet this retrieval method.
There was a problem hiding this comment.
I checked this against SGLang v0.5.10.post1 (commit 7c35342c).
The fused MoE implementation preserves the local-expert index as tensor dimension 0. The unquantized implementation allocates both tensors as:
w13_weight = torch.nn.Parameter(
torch.empty(num_experts, w13_weight_n, w13_weight_k, dtype=params_dtype),
requires_grad=False,
)
...
w2_weight = torch.nn.Parameter(
torch.empty(num_experts, w2_weight_n, w2_weight_k, dtype=params_dtype),
requires_grad=False,
)FusedMoE passes self.num_local_experts as num_experts:
SGLang maps the global routed-expert ID to that local dimension using the same EP-rank offset:
start_idx = self.moe_ep_rank * num_local_routed_experts
...
return expert_id - start_idxTherefore, for the Qwen3-MoE path covered by this PR, dimension 0 remains the local-expert dimension and the ep_rank * num_local_experts offset matches SGLang's own mapping. The Triton-kernel layout only swaps the two inner dimensions and does not change the expert dimension.
| # (name, rank). The transfer plan indexes by name alone, so without | ||
| # merging only the last entry per name survives and the other ranks' | ||
| # shards are dropped. | ||
| infer_params_meta = _merge_meta_by_name(infer_params_meta) |
There was a problem hiding this comment.
Each inference_worker_url represents a complete inference instance, with the response containing the metadata for the instance's full TP/EP rank. PR first merges the shards of N instances into a single replica; AWEX TransferPlanBuilder then replicates this replica with num_infer_engines=N. As a result, each inference rank receives N identical operations, and the global plan is expanded from N groups to N² groups.
There was a problem hiding this comment.
Confirmed and fixed in 14b0a599.
Each inference URL is now merged independently into one complete per-instance metadata set. The gateway validates that every inference instance reports equivalent metadata and stores only the first canonical instance. num_infer_engines remains the only cross-instance expansion performed by TransferPlanBuilder.
I also added a planner-level regression test with two inference instances and two shards per instance. The previous global merge produces 8 communication operations; the corrected path produces exactly 4 operations, targeting inference ranks 0, 1, 2, and 3 once each. The same canonicalization is applied by both separation and colocate /connect paths.
| ep_size = getattr(rank_info, "ep_size", 1) or 1 | ||
| if ep_size <= 1: | ||
| return 0 | ||
| return getattr(rank_info, "ep_rank", 0) * num_local_experts |
There was a problem hiding this comment.
If there is a shared expert here, it will be recorded in the current num_local_experts, causing the overall routed ID to shift.
2 participants
Description
Fix v2 AWEX weight transfer for MoE models when SGLang uses expert
parallelism.
The inference adapter previously used the fused tensor's local leading index
as the global expert ID. With 128 experts and
ep_size=4, every EP ranktherefore advertised experts 0–31 while experts 32–127 were never named. The
same naming path is also used to build the actual shard payload, so metadata
and transferred tensors must use the same EP-aware global IDs.
The gateway also concatenated inference metadata from every rank without
merging entries that shared a parameter name. Since the transfer plan is keyed
by name, duplicate entries shadowed one another and metadata size grew with the
inference world size. The training side already performs this merge; this PR
applies the same behavior to inference metadata while preserving all rank
shards and replicas.
This PR:
ep_rank * num_local_expertsfor fused gate/upand down-projection weights;
names consistent with advertised metadata;
sites.
Related Issue
N/A
Type of Change
Checklist
Contributing Guide
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prBreaking Change Details (if applicable):
None.
Additional Context
Before the fix, a 128-expert model with
ep_size=4exposed two failure modes:0–31, producing missing expert keys; and
the flat metadata list to 16 copies of each parameter and causing the
connection setup to time out.
In a 2-node separation run after the fix:
entries;
that inference and training remained aligned.
Verification:
tests/v2/weight_update: 54 passed, 33 skipped;the gateway merge call both make the regression tests fail; and
run passed.
Need help? Check the
Contributing Guide
or ask in GitHub Discussions!