Skip to content

fix(v2): build complete MoE inference metadata for weight transfer #1595

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Le8r0nJames wants to merge 2 commits into
base: main
Choose a base branch
from
+325 −24
Open

fix(v2): build complete MoE inference metadata for weight transfer #1595

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
983bf06
fix(v2): advertise every routed expert when transferring MoE weights
Le8r0nJames Aug 11, 2026
14b0a59
fix(v2): preserve inference instance metadata boundaries
Le8r0nJames Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 31 additions & 8 deletions areal/v2/weight_update/awex/sglang_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,21 @@ def parallelism_strategy(self) -> dict:
"num_engines": 1,
}

def _expert_id_offset(self, num_local_experts: int) -> int:
"""Global id of this rank's first routed expert.

Under expert parallelism SGLang holds only ``num_experts // ep_size``
experts per rank, so the fused tensor's leading index is local. The
training side publishes global HuggingFace names and the transfer plan is
indexed by name, so without the offset every rank claims the same
low-numbered experts and the rest are never advertised at all.
"""
rank_info = self._rank_info or self._build_rank_info()
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

@sitabulaixizawaluduo sitabulaixizawaluduo Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there is a shared expert here, it will be recorded in the current num_local_experts, causing the overall routed ID to shift.


def _unfuse_params(
self, name: str, tensor: torch.Tensor
) -> list[tuple[str, torch.Tensor]]:
Expand Down Expand Up @@ -185,13 +200,15 @@ def _unfuse_params(
prefix = name.replace(".w13_weight", "")
result = []
ffn_hidden = tensor.shape[1] // 2
id_offset = self._expert_id_offset(tensor.shape[0])

@sitabulaixizawaluduo sitabulaixizawaluduo Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SGLang uses fused moe. It is recommended to confirm whether the fused tensor can still meet this retrieval method.

@Le8r0nJames Le8r0nJames Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

https://github.com/sgl-project/sglang/blob/7c35342c10e201899e22fe2972d40e60da19ff3e/python/sglang/srt/layers/quantization/unquant.py#L185-L220

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:

https://github.com/sgl-project/sglang/blob/7c35342c10e201899e22fe2972d40e60da19ff3e/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L268-L285

SGLang maps the global routed-expert ID to that local dimension using the same EP-rank offset:

https://github.com/sgl-project/sglang/blob/7c35342c10e201899e22fe2972d40e60da19ff3e/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L557-L571

start_idx = self.moe_ep_rank * num_local_routed_experts
...
return expert_id - start_idx

Therefore, 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.

for i in range(tensor.shape[0]):
expert_tensor = tensor[i]
if i < num_routed:
expert_prefix = f"{prefix}.{i}"
expert_id = i + id_offset
if expert_id < num_routed:
expert_prefix = f"{prefix}.{expert_id}"
else:
shared_idx = i - num_routed
num_shared = tensor.shape[0] - num_routed
shared_idx = expert_id - num_routed
num_shared = tensor.shape[0] + id_offset - num_routed
if num_shared > 1:
expert_prefix = prefix.replace(
"experts", f"shared_experts.{shared_idx}"
Expand All @@ -211,12 +228,14 @@ def _unfuse_params(
num_routed = getattr(cfg, "num_experts", None) or cfg.n_routed_experts
prefix = name.replace(".w2_weight", "")
result = []
id_offset = self._expert_id_offset(tensor.shape[0])
for i in range(tensor.shape[0]):
if i < num_routed:
expert_prefix = f"{prefix}.{i}"
expert_id = i + id_offset
if expert_id < num_routed:
expert_prefix = f"{prefix}.{expert_id}"
else:
shared_idx = i - num_routed
num_shared = tensor.shape[0] - num_routed
shared_idx = expert_id - num_routed
num_shared = tensor.shape[0] + id_offset - num_routed
if num_shared > 1:
expert_prefix = prefix.replace(
"experts", f"shared_experts.{shared_idx}"
Expand Down Expand Up @@ -330,6 +349,10 @@ def get_local_shard_parameters(
) -> dict[str, torch.Tensor]:
required = set(required_names) if required_names else None
local_params: dict[str, torch.Tensor] = {}
# Expert ids below come from this rank's expert-parallel position, and
# the payload has to use the same names the metadata advertised.
if self._rank_info is None:
self._rank_info = self._build_rank_info()

for name, param in self._get_model().named_parameters():
for hf_name, hf_tensor in self._unfuse_params(name, param.data):
Expand Down
46 changes: 30 additions & 16 deletions areal/v2/weight_update/gateway/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,8 @@ def _get_own_ip() -> str:
return "127.0.0.1"


def _merge_training_meta_by_name(meta_list: list[dict]) -> list[dict]:
"""Merge serialized training ParameterMeta entries by parameter name.
def _merge_meta_by_name(meta_list: list[dict]) -> list[dict]:
"""Merge serialized ParameterMeta entries by parameter name.

Each FSDP worker reports metadata for its own local shard only.
With ``dp_size > 1`` the same parameter name appears once per worker,
Expand Down Expand Up @@ -167,6 +167,32 @@ def _merge_training_meta_by_name(meta_list: list[dict]) -> list[dict]:
return list(by_name.values()) + overflow


# Historical alias; the merge is not training-specific.
_merge_training_meta_by_name = _merge_meta_by_name


def _canonical_inference_meta(meta_responses: list[dict]) -> list[dict]:
"""Return metadata for one inference instance after validating its replicas.

AWEX expands one instance's metadata by ``num_infer_engines`` when it builds
the transfer plan. Merging metadata across inference instances here would
make that expansion count every instance twice.
"""
canonical = None
for instance_idx, result in enumerate(meta_responses):
meta = result.get("result", result.get("meta", result))
instance_meta = meta if isinstance(meta, list) else [meta]
instance_meta = _merge_meta_by_name(instance_meta)
if canonical is None:
canonical = instance_meta
elif instance_meta != canonical:
raise ValueError(
f"Inference instance {instance_idx} reported different weight metadata"
)

return canonical or []


def create_app(config: WeightUpdateConfig | None = None) -> FastAPI:
config = config or WeightUpdateConfig()

Expand Down Expand Up @@ -313,13 +339,7 @@ async def connect(request: Request, body: ConnectRequest) -> ConnectResponse:
training_params_meta.append(meta)
training_params_meta = _merge_training_meta_by_name(training_params_meta)

infer_params_meta = []
for result in infer_meta_resps:
meta = result.get("result", result.get("meta", result))
if isinstance(meta, list):
infer_params_meta.extend(meta)
else:
infer_params_meta.append(meta)
infer_params_meta = _canonical_inference_meta(infer_meta_resps)

kv_store.put(pair_name, "training_params_meta", training_params_meta)
kv_store.put(pair_name, "infer_params_meta", infer_params_meta)
Expand Down Expand Up @@ -449,13 +469,7 @@ async def _connect_colocate(
training_params_meta.append(meta)
training_params_meta = _merge_training_meta_by_name(training_params_meta)

infer_params_meta = []
for result in infer_meta_resps:
meta = result.get("result", result.get("meta", result))
if isinstance(meta, list):
infer_params_meta.extend(meta)
else:
infer_params_meta.append(meta)
infer_params_meta = _canonical_inference_meta(infer_meta_resps)

kv_store.put(pair_name, "training_params_meta", training_params_meta)
kv_store.put(pair_name, "infer_params_meta", infer_params_meta)
Expand Down
Loading
Loading