feat: implement megatron weight update with awex

areal/experimental/training_service/worker/awex.py

@@ -135,11 +135,18 @@ def action():
 
 def _create_training_adapter(engine):
     from areal.engine.fsdp_engine import FSDPEngine
+    from areal.engine.megatron_engine import MegatronEngine
     from areal.experimental.weight_update.awex.fsdp_adapter import AwexFSDPAdapter
+    from areal.experimental.weight_update.awex.megatron_adapter import (
+        AwexMegatronAdapter,
+    )
 
     if isinstance(engine, FSDPEngine):
         return AwexFSDPAdapter(engine)
 
+    if isinstance(engine, MegatronEngine):
+        return AwexMegatronAdapter(engine)
+
     raise TypeError(
         f"Unsupported engine type for weight update: {type(engine).__name__}"
     )

areal/experimental/weight_update/awex/megatron_adapter.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+from awex.meta.weight_meta import (
+    ParameterMeta,
+    ParameterReplicaMeta,
+    ParameterShardMeta,
+)
+from awex.sharding.param_sharding import ShardingType
+from awex.sharding.rank_info import RankInfo
+from awex.transfer.nccl_comm import batch_send_recv, nccl_build_send_ops
+from awex.transfer.transfer_plan import TransferPlan, TransferPlanBuilder
+
+from areal.experimental.weight_update.awex import fetch_kv_metadata
+from areal.experimental.weight_update.nccl_group import (
+    init_weights_update_group,
+    setup_batch_isend_irecv,
+)
+from areal.experimental.weight_update.training_adapter import (
+    WeightUpdateTrainingAdapter,
+)
+from areal.utils import logging
+
+if TYPE_CHECKING:
+    from areal.engine.megatron_engine import MegatronEngine
+
+logger = logging.getLogger("AwexMegatronAdapter")
+
+
+class AwexMegatronAdapter(WeightUpdateTrainingAdapter):
+    """Awex training adapter wrapping MegatronEngine for DP-only NCCL P2P weight updates.
+
+    Scope: DP-only (tp=1, pp=1). Each rank holds a full replica of every
+    parameter. Parameters are converted to HF naming via convert_to_hf before
+    being handed to the awex transfer planner, matching what the inference
+    engine (SGLang) expects.
+    """
+
+    def __init__(self, engine: MegatronEngine):
+        self._engine = engine
+        self._transfer_plan: TransferPlan | None = None
+        self._weights_update_group = None
+        self._transfer_rank: int | None = None
+
+    @property
+    def parallelism_strategy(self) -> dict:
+        from megatron.core import parallel_state as mpu
+
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+        return {
+            "world_size": self._engine.world_size,
+            "tp_size": tp_size,
+            "pp_size": mpu.get_pipeline_model_parallel_world_size(),
+            "dp_size": self._engine.data_parallel_world_size,
+            "ep_size": mpu.get_expert_model_parallel_world_size(),
+            "dp_replicated": tp_size > 1,
+        }
+
+    def get_weight_metadata(self) -> list[ParameterMeta]:
+        rank_info = self._build_rank_info()
+        metadata: list[ParameterMeta] = []
+
+        for hf_name, tensor in self._iter_hf_params():
+            shape = tuple(tensor.shape)
+            numel = int(tensor.numel())
+            shard_meta = ParameterShardMeta(
+                tp_rank=rank_info.tp_rank,
+                attn_tp_rank=rank_info.attn_tp_rank,
+                pp_rank=rank_info.pp_rank,
+                ep_rank=rank_info.ep_rank,
+                ep_tp_rank=rank_info.ep_tp_rank,
+                global_rank=rank_info.global_rank,
+                world_size=rank_info.world_size,
+                engine_rank=rank_info.engine_rank,
+                cp_rank=rank_info.cp_rank,
+                cp_size=rank_info.cp_size,
+                cp_mode=rank_info.cp_mode,
+                name=hf_name,
+                shape=shape,
+                numel=numel,
+                dtype=tensor.dtype,
+                global_offset=tuple([0] * len(shape)),
+                sharding_type=ShardingType.NO_SHARDING,
+                num_shards=1,
+                sharding_dim=0,
+            )
+            replica = ParameterReplicaMeta(shards=[shard_meta])
+            metadata.append(
+                ParameterMeta(
+                    name=hf_name,
+                    global_numel=numel,
+                    global_shape=shape,
+                    dtype=tensor.dtype,
+                    shards=[shard_meta],
+                    replicas=[replica],
+                )
+            )
+
+        return metadata
+
+    def get_local_shard_parameters(
+        self, required_names: list[str] | None = None
+    ) -> dict[str, torch.Tensor]:
+        required = set(required_names) if required_names else None
+        result: dict[str, torch.Tensor] = {}
+        for hf_name, tensor in self._iter_hf_params():
+            if required is not None and hf_name not in required:
+                continue
+            result[hf_name] = tensor
+        return result
+
+    def save_parameters(self, save_path: str, names: list[str] | None = None) -> None:
+        params = self.get_local_shard_parameters(names)
+        cpu_params = {k: v.detach().cpu().clone() for k, v in params.items()}
+        torch.save(cpu_params, save_path)
+
+    def init_weight_update_group(
+        self,
+        pair_name: str,
+        master_addr: str,
+        master_port: int,
+        transfer_rank: int,
+        world_size: int,
+        kv_store_url: str,
+        infer_world_size: int,
+        train_world_size: int,
+        num_engines: int,
+    ) -> None:
+        self._transfer_rank = transfer_rank
+
+        infer_meta, train_meta = fetch_kv_metadata(kv_store_url, pair_name)
+
+        builder = TransferPlanBuilder(
+            infer_world_size=infer_world_size,
+            train_world_size=train_world_size,
+            num_infer_engines=num_engines,
+        )
+        self._transfer_plan = builder.build_local_transfer_plan(
+            infer_meta, train_meta, global_transfer_rank=transfer_rank
+        )
+
+        os.environ["TORCHELASTIC_USE_AGENT_STORE"] = str(False)
+        self._weights_update_group = init_weights_update_group(
+            master_address=master_addr,
+            master_port=master_port,
+            rank=transfer_rank,
+            world_size=world_size,
+            group_name=f"awex_{pair_name}",
+            role="training",
+        )
+
+    def execute_weight_update(self, version: int) -> None:
+        del version
+        if self._transfer_plan is None:
+            raise RuntimeError("Transfer plan is not initialized")
+        if self._weights_update_group is None:
+            raise RuntimeError("Weight update group is not initialized")
+        if self._transfer_rank is None:
+            raise RuntimeError("Transfer rank is not initialized")
+
+        params = self.get_local_shard_parameters()
+        send_ops, _, _ = nccl_build_send_ops(
+            params,
+            self._transfer_plan,
+            self._weights_update_group,
+            copy_rank=self._transfer_rank,
+        )
+        batch_send_recv(send_ops=send_ops, recv_ops=[], blocking=True)
+        dist.barrier(group=self._weights_update_group)
+
+    def batch_isend_irecv(self, **kwargs) -> None:
+        setup_kwargs = {k: v for k, v in kwargs.items() if k != "world_size"}
+        setup_batch_isend_irecv(
+            self._weights_update_group,
+            self._transfer_rank,
+            kwargs.get("world_size", 0),
+            **setup_kwargs,
+        )
+
+    def teardown_weight_update_group(self) -> None:
+        if self._weights_update_group is not None and dist.is_initialized():
+            dist.destroy_process_group(self._weights_update_group)
+        self._weights_update_group = None
+        self._transfer_plan = None
+        self._transfer_rank = None
+
+    def _build_rank_info(self) -> RankInfo:
+        from megatron.core import parallel_state as mpu
+
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+        tp_rank = mpu.get_tensor_model_parallel_rank()
+        pp_size = mpu.get_pipeline_model_parallel_world_size()
+        pp_rank = mpu.get_pipeline_model_parallel_rank()
+        ep_size = mpu.get_expert_model_parallel_world_size()
+        ep_rank = mpu.get_expert_model_parallel_rank()
+        local_rank = int(os.environ.get("LOCAL_RANK", self._engine.rank))
+
+        return RankInfo(
+            tp_rank=tp_rank,
+            tp_size=tp_size,
+            pp_rank=pp_rank,
+            pp_size=pp_size,
+            dp_size=self._engine.data_parallel_world_size,
+            dp_rank=self._engine.data_parallel_rank,
+            ep_rank=ep_rank,
+            ep_size=ep_size,
+            ep_tp_rank=0,
+            ep_tp_size=1,
+            attn_tp_rank=tp_rank,
+            attn_tp_size=tp_size,
+            attn_dp_rank=self._engine.data_parallel_rank,
+            world_size=self._engine.world_size,
+            global_rank=self._engine.rank,
+            local_rank=local_rank,
+            engine_rank=0,
+            is_infer=False,
+            cp_rank=0,
+            cp_size=1,
+            cp_mode="none",
+        )
+
+    def _iter_hf_params(self):
+        """Yield (hf_name, tensor) for every parameter on this rank.
+
+        Uses the same get_named_parameters + all_gather_param + convert_to_hf
+        pipeline that MegatronEngine._collect_param uses for weight broadcast,
+        so the names and shapes are guaranteed to match what SGLang expects.
+
+        For DP-only (tp=1, pp=1): all_gather_param is a no-op (returns
+        param.data directly) and convert_to_hf remaps mcore names to HF names.
+        """
+        from areal.engine.megatron_utils.megatron import (
+            all_gather_param,
+            convert_to_hf,
+            get_named_parameters,
+        )
+
+        num_moe_experts = getattr(self._engine.tf_config, "num_moe_experts", None)
+        model_name = self._engine.hf_config.model_type
+
+        for mcore_name, param in get_named_parameters(
+            self._engine.model, num_moe_experts
+        ):
+            gathered = all_gather_param(
+                mcore_name,
+                param,
+                fp8_direct_convert=False,
+                quantization_config=None,
+                duplicated_param_names=self._engine._duplicated_param_names,
+            )
+            if not isinstance(gathered, torch.Tensor):
+                gathered = gathered.data
+
+            for hf_name, tensor in convert_to_hf(
+                self._engine.tf_config,
+                model_name,
+                mcore_name,
+                gathered,
+            ):
+                yield hf_name, tensor.detach()