diff --git a/areal/v2/weight_update/awex/sglang_adapter.py b/areal/v2/weight_update/awex/sglang_adapter.py index d70b5d2cd0..30977c99f8 100644 --- a/areal/v2/weight_update/awex/sglang_adapter.py +++ b/areal/v2/weight_update/awex/sglang_adapter.py @@ -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 + def _unfuse_params( self, name: str, tensor: torch.Tensor ) -> list[tuple[str, torch.Tensor]]: @@ -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]) 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}" @@ -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}" @@ -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): diff --git a/areal/v2/weight_update/gateway/app.py b/areal/v2/weight_update/gateway/app.py index a381b8211c..64c43ddfe3 100644 --- a/areal/v2/weight_update/gateway/app.py +++ b/areal/v2/weight_update/gateway/app.py @@ -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, @@ -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() @@ -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) @@ -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) diff --git a/tests/v2/weight_update/test_moe_expert_metadata.py b/tests/v2/weight_update/test_moe_expert_metadata.py new file mode 100644 index 0000000000..ae93968699 --- /dev/null +++ b/tests/v2/weight_update/test_moe_expert_metadata.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Expert-name and metadata-merge contracts for MoE weight transfer. + +Both behaviours below decide whether a routed expert is transferred at all, and +both fail silently: the shapes stay correct either way, so a mismatch shows up +as diverging train/inference logprobs rather than an error. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from areal.v2.weight_update.awex.sglang_adapter import AwexSGLangAdapter +from areal.v2.weight_update.gateway.app import _merge_meta_by_name + +NUM_EXPERTS = 8 +HIDDEN = 4 +FFN = 6 + + +def _adapter(ep_size, ep_rank): + adapter = AwexSGLangAdapter.__new__(AwexSGLangAdapter) + adapter._get_model = lambda: SimpleNamespace( + config=SimpleNamespace(num_experts=NUM_EXPERTS, n_routed_experts=NUM_EXPERTS) + ) + adapter._rank_info = SimpleNamespace(ep_size=ep_size, ep_rank=ep_rank) + adapter._build_rank_info = lambda: adapter._rank_info + return adapter + + +def _local(ep_size): + return NUM_EXPERTS // ep_size + + +def _expert_ids(pairs): + return { + int(name.split(".experts.")[1].split(".")[0]) + for name, _ in pairs + if ".experts." in name + } + + +@pytest.mark.parametrize("fused", ["w13_weight", "w2_weight"]) +@pytest.mark.parametrize("ep_rank", range(4)) +def test_expert_ids_are_global_under_expert_parallelism(fused, ep_rank): + ep_size = 4 + adapter = _adapter(ep_size, ep_rank) + local = _local(ep_size) + shape = (local, 2 * FFN, HIDDEN) if fused == "w13_weight" else (local, HIDDEN, FFN) + + pairs = adapter._unfuse_params( + f"model.layers.0.mlp.experts.{fused}", torch.zeros(*shape) + ) + + assert _expert_ids(pairs) == set(range(ep_rank * local, (ep_rank + 1) * local)) + + +def test_ranks_together_advertise_every_expert_exactly_once(): + """The transfer plan indexes by name, so a shadowed id is never sent.""" + ep_size = 4 + seen = [] + for ep_rank in range(ep_size): + adapter = _adapter(ep_size, ep_rank) + seen.extend( + _expert_ids( + adapter._unfuse_params( + "model.layers.0.mlp.experts.w13_weight", + torch.zeros(_local(ep_size), 2 * FFN, HIDDEN), + ) + ) + ) + + assert sorted(seen) == list(range(NUM_EXPERTS)) + + +def test_expert_ids_unchanged_without_expert_parallelism(): + adapter = _adapter(ep_size=1, ep_rank=0) + + pairs = adapter._unfuse_params( + "model.layers.0.mlp.experts.w13_weight", + torch.zeros(NUM_EXPERTS, 2 * FFN, HIDDEN), + ) + + assert _expert_ids(pairs) == set(range(NUM_EXPERTS)) + + +def test_shared_experts_keep_their_own_index_under_expert_parallelism(): + ep_size, ep_rank = 2, 1 + adapter = _adapter(ep_size, ep_rank) + local = _local(ep_size) + + pairs = adapter._unfuse_params( + "model.layers.0.mlp.experts.w13_weight", + torch.zeros(local, 2 * FFN, HIDDEN), + ) + + assert _expert_ids(pairs) == {4, 5, 6, 7} + assert not any("shared_experts" in name for name, _ in pairs) + + +def _entry(name, rank): + return { + "data": { + "name": name, + "shards": [{"rank": rank}], + "replicas": [{"data": {"shards": [{"rank": rank}]}}], + } + } + + +def test_merging_keeps_every_rank_shard_under_one_name(): + name = "model.layers.0.mlp.experts.3.down_proj.weight" + flat = [_entry(name, rank) for rank in range(16)] + + merged = _merge_meta_by_name(flat) + + assert len(merged) == 1 + assert len(merged[0]["data"]["shards"]) == 16 + + +def test_merging_preserves_distinct_names(): + flat = [_entry("a.weight", 0), _entry("b.weight", 0), _entry("a.weight", 1)] + + merged = _merge_meta_by_name(flat) + + assert {e["data"]["name"] for e in merged} == {"a.weight", "b.weight"} + assert len(merged) == 2 + + +def test_inference_instances_keep_one_canonical_replica(): + from areal.v2.weight_update.gateway import app as gateway_app + + responses = [ + { + "meta": [ + _entry("model.layers.0.mlp.experts.3.down_proj.weight", rank) + for rank in range(2) + ] + } + for _ in range(2) + ] + + canonical = gateway_app._canonical_inference_meta(responses) + + assert len(canonical) == 1 + assert len(canonical[0]["data"]["shards"]) == 2 + replicas = canonical[0]["data"]["replicas"] + assert len(replicas) == 1 + assert [shard["rank"] for shard in canonical[0]["data"]["shards"]] == [0, 1] + assert [shard["rank"] for shard in replicas[0]["data"]["shards"]] == [0, 1] + + +def test_inference_instances_with_different_metadata_fail_fast(): + from areal.v2.weight_update.gateway import app as gateway_app + + responses = [ + {"meta": [_entry("a.weight", 0), _entry("a.weight", 1)]}, + {"meta": [_entry("a.weight", 0)]}, + ] + + with pytest.raises( + ValueError, + match="Inference instance 1 reported different weight metadata", + ): + gateway_app._canonical_inference_meta(responses) + + +def test_canonical_metadata_builds_one_transfer_per_engine_shard(): + from awex.meta.meta_resolver import ( + ParameterMeta, + ParameterReplicaMeta, + ParameterShardMeta, + ) + from awex.sharding.param_sharding import ShardingType + from awex.transfer.transfer_plan import TransferPlanBuilder + + from areal.infra.rpc.serialization import deserialize_value, serialize_value + from areal.v2.weight_update.gateway import app as gateway_app + + name = "model.layers.0.mlp.experts.3.down_proj.weight" + + def shard(rank, shape, global_offset, world_size): + return ParameterShardMeta( + tp_rank=rank, + attn_tp_rank=rank, + pp_rank=0, + ep_rank=0, + ep_tp_rank=rank, + global_rank=rank, + world_size=world_size, + engine_rank=0, + name=name, + shape=shape, + numel=shape[0], + dtype=torch.float32, + global_offset=global_offset, + sharding_type=ShardingType.TP_SHARDING, + num_shards=world_size, + sharding_dim=0, + ) + + def per_rank_meta(rank): + infer_shard = shard(rank, shape=(2,), global_offset=(rank * 2,), world_size=2) + return ParameterMeta( + name=name, + global_numel=4, + global_shape=(4,), + dtype=torch.float32, + shards=[infer_shard], + replicas=[ParameterReplicaMeta(shards=[infer_shard])], + ) + + responses = [ + {"meta": serialize_value([per_rank_meta(0), per_rank_meta(1)])} + for _ in range(2) + ] + canonical = deserialize_value(gateway_app._canonical_inference_meta(responses)) + + train_shard = shard(0, shape=(4,), global_offset=(0,), world_size=1) + training = [ + ParameterMeta( + name=name, + global_numel=4, + global_shape=(4,), + dtype=torch.float32, + shards=[train_shard], + replicas=[ParameterReplicaMeta(shards=[train_shard])], + ) + ] + builder = TransferPlanBuilder( + infer_world_size=4, + train_world_size=1, + num_infer_engines=2, + ) + + operations = builder.build_weights_mapping_operations(canonical, training) + + assert len(operations) == 4 + assert [operation.recv_rank for operation in operations] == [0, 1, 2, 3] + + +def test_connect_canonicalizes_the_inference_metadata_it_collected(): + """Guard the call site, not just the helper. + + The helper being correct is useless if /connect forgets to apply it, which + is exactly the state this fix repairs. + """ + import inspect + + from areal.v2.weight_update.gateway import app as gateway_app + + source = inspect.getsource(gateway_app) + collect_sites = source.count("infer_meta_resps = await asyncio.gather") + canonicalize_sites = source.count( + "infer_params_meta = _canonical_inference_meta(infer_meta_resps)" + ) + + assert collect_sites > 0 + assert canonicalize_sites == collect_sites, ( + f"{collect_sites} places collect inference metadata but only " + f"{canonicalize_sites} preserve one canonical inference instance" + )