diff --git a/areal/experimental/training_service/worker/awex.py b/areal/experimental/training_service/worker/awex.py
index a2b73e94ad..f1d13cbb18 100644
--- a/areal/experimental/training_service/worker/awex.py
+++ b/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__}"
     )
diff --git a/areal/experimental/weight_update/awex/megatron_adapter.py b/areal/experimental/weight_update/awex/megatron_adapter.py
new file mode 100644
index 0000000000..7c61a6a06c
--- /dev/null
+++ b/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()
diff --git a/tests/experimental/weight_update/test_nccl_integration.py b/tests/experimental/weight_update/test_nccl_integration.py
index 2cfee17a30..1107282587 100644
--- a/tests/experimental/weight_update/test_nccl_integration.py
+++ b/tests/experimental/weight_update/test_nccl_integration.py
@@ -202,11 +202,86 @@ def _validate_weight_update_correctness(
     )
 
 
+def _validate_weight_update_correctness_megatron(
+    train_worker_urls: list[str],
+    inf_worker_url: str,
+    param_dir,
+) -> None:
+    import concurrent.futures
+
+    n_train = len(train_worker_urls)
+    print(
+        f"\n[weight-validation] Fetching parameters from {n_train} Megatron "
+        f"worker(s) and 1 inference worker …"
+    )
+
+    train_paths = [
+        str(param_dir / f"megatron_train_params_rank{i}.pt") for i in range(n_train)
+    ]
+
+    def _fetch_train(args):
+        i, url, p = args
+        resp = httpx.post(
+            f"{url}/awex/debug/get_parameters",
+            json={"save_path": p, "names": _VALIDATE_PARAM_NAMES},
+            timeout=120.0,
+        )
+        assert resp.status_code == 200, (
+            f"get_parameters failed on training worker {i}: {resp.text}"
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=n_train) as pool:
+        list(
+            pool.map(
+                _fetch_train,
+                [
+                    (i, url, p)
+                    for i, (url, p) in enumerate(zip(train_worker_urls, train_paths))
+                ],
+            )
+        )
+
+    inf_path = str(param_dir / "megatron_infer_params.pt")
+    resp = httpx.post(
+        f"{inf_worker_url}/awex/debug/get_parameters",
+        json={"save_path": inf_path, "names": _VALIDATE_PARAM_NAMES},
+        timeout=120.0,
+    )
+    assert resp.status_code == 200, (
+        f"get_parameters failed on inference worker: {resp.text}"
+    )
+
+    infer_params = torch.load(inf_path, map_location="cpu", weights_only=True)
+    train_params = torch.load(train_paths[0], map_location="cpu", weights_only=True)
+
+    print(f"[weight-validation] Comparing {len(_VALIDATE_PARAM_NAMES)} parameters …")
+    for name in _VALIDATE_PARAM_NAMES:
+        assert name in infer_params, f"Inference missing param: {name}"
+        assert name in train_params, f"Training rank 0 missing param: {name}"
+
+        torch.testing.assert_close(
+            train_params[name],
+            infer_params[name],
+            rtol=0,
+            atol=0,
+            msg=f"Parameter mismatch after weight update: {name}",
+        )
+        print(
+            f"[weight-validation]   {name}: OK "
+            f"(shape={list(train_params[name].shape)}, dtype={train_params[name].dtype})"
+        )
+
+    print(
+        f"[weight-validation] All {len(_VALIDATE_PARAM_NAMES)} parameters "
+        f"match between Megatron training and inference ✓"
+    )
+
+
 @pytest.mark.multi_gpu
 @pytest.mark.slow
 @pytest.mark.sglang
 @pytest.mark.parametrize("n_gpus", [2, 4, 8], ids=["2gpu", "4gpu", "8gpu"])
-def test_awex_e2e_weight_update(n_gpus, tmp_path_factory):
+def test_awex_fsdp_e2e_weight_update(n_gpus, tmp_path_factory):
     """Full round trip: FSDPEngine (pure DP) → weight-update gateway → SGLang.
 
     A single :class:`LocalScheduler` owns all *n_gpus* devices.  Inference
@@ -360,3 +435,311 @@ def test_awex_e2e_weight_update(n_gpus, tmp_path_factory):
         train_ctrl.destroy()
         inf_ctrl.destroy()
         scheduler.delete_workers(None)
+
+
+@pytest.mark.multi_gpu
+@pytest.mark.slow
+@pytest.mark.sglang
+@pytest.mark.parametrize("n_gpus", [2, 4, 8], ids=["2gpu", "4gpu", "8gpu"])
+def test_awex_megatron_e2e_weight_update(n_gpus, tmp_path_factory):
+    """Full round trip: MegatronEngine (pure DP) → weight-update gateway → SGLang.
+
+    Mirrors test_awex_e2e_weight_update but uses MegatronLMEngine as the
+    training backend.  DP-only (tp=1, pp=1): every training rank holds a full
+    copy of every parameter, so validation compares rank-0's params directly
+    against the inference server without any concatenation.
+    """
+    if current_platform.device_count() < n_gpus:
+        pytest.skip(f"This test requires {n_gpus} GPUs")
+
+    from areal.api import FinetuneSpec
+    from areal.api.cli_args import (
+        OptimizerConfig,
+        SchedulingSpec,
+        TrainEngineConfig,
+    )
+    from areal.experimental.inference_service.controller.config import (
+        GatewayControllerConfig,
+    )
+    from areal.experimental.inference_service.controller.controller import (
+        GatewayInferenceController,
+    )
+    from areal.experimental.training_service.controller.controller import (
+        GatewayTrainController,
+    )
+    from areal.experimental.weight_update.controller import (
+        WeightUpdateController,
+        WeightUpdateControllerConfig,
+    )
+
+    n_half = n_gpus // 2
+    tmp = tmp_path_factory.mktemp("awex_megatron_e2e")
+    model_path = _get_test_model_path()
+
+    scheduler = _make_local_scheduler(
+        tmp, "megatron-e2e", gpu_devices=list(range(n_gpus))
+    )
+
+    inf_config = GatewayControllerConfig(
+        tokenizer_path=model_path,
+        model_path=model_path,
+        backend=f"sglang:d{n_half}",
+        scheduling_spec=(
+            SchedulingSpec(
+                gpu=1,
+                cmd="python -m areal.experimental.inference_service.guard",
+            ),
+        ),
+        consumer_batch_size=8,
+        max_head_offpolicyness=1024,
+        setup_timeout=300.0,
+        admin_api_key="test-admin",
+    )
+    inf_ctrl = GatewayInferenceController(config=inf_config, scheduler=scheduler)
+
+    train_config = TrainEngineConfig(
+        backend=f"megatron:d{n_half}",
+        experiment_name="test-awex-megatron-e2e",
+        trial_name="t0",
+        path=model_path,
+        optimizer=OptimizerConfig(),
+        _version="v2",
+        setup_timeout=300.0,
+        scheduling_spec=(
+            SchedulingSpec(
+                gpu=1,
+                cmd="python -m areal.experimental.training_service.guard",
+                env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"),
+            ),
+        ),
+    )
+    train_ctrl = GatewayTrainController(
+        train_engine="areal.engine.megatron_engine.MegatronLMEngine",
+        config=train_config,
+        scheduler=scheduler,
+    )
+
+    wu_ctrl: WeightUpdateController | None = None
+
+    try:
+        inf_ctrl.initialize(
+            role="rollout",
+            server_args={"mem_fraction_static": 0.7},
+        )
+        inf_worker_urls = list(inf_ctrl._inf_addrs)
+
+        for url in inf_worker_urls:
+            resp = httpx.post(f"{url}/awex/debug/randomize_parameters", timeout=120.0)
+            assert resp.status_code == 200, f"randomize_parameters failed: {resp.text}"
+
+        ft_spec = FinetuneSpec(
+            total_train_epochs=1, dataset_size=100, train_batch_size=2
+        )
+        train_ctrl.initialize(role="actor", ft_spec=ft_spec)
+        train_worker_urls = list(train_ctrl._worker_addrs)
+
+        wu_ctrl = WeightUpdateController(
+            config=WeightUpdateControllerConfig(
+                host="127.0.0.1",
+                request_timeout=300.0,
+            )
+        )
+        wu_ctrl.initialize()
+
+        assert wu_ctrl.health_check(), "Weight update gateway health check failed"
+
+        wu_ctrl.connect(
+            pair_name="test_megatron_e2e",
+            train_worker_urls=train_worker_urls,
+            inference_worker_urls=inf_worker_urls,
+        )
+
+        result = wu_ctrl.update_weights(version=1)
+        assert result.status == "ok"
+        assert result.version == 1
+
+        wu_ctrl.disconnect()
+
+        gen_resp = httpx.post(
+            f"{inf_worker_urls[0]}/generate",
+            json={
+                "text": "Hello",
+                "sampling_params": {"max_new_tokens": 5, "temperature": 0},
+            },
+            timeout=30.0,
+        )
+        assert gen_resp.status_code == 200, (
+            f"Generation failed after weight update: {gen_resp.text}"
+        )
+
+        _validate_weight_update_correctness_megatron(
+            train_worker_urls=train_worker_urls,
+            inf_worker_url=inf_worker_urls[0],
+            param_dir=tmp,
+        )
+
+    finally:
+        if wu_ctrl is not None:
+            wu_ctrl.destroy()
+        train_ctrl.destroy()
+        inf_ctrl.destroy()
+        scheduler.delete_workers(None)
+
+
+@pytest.mark.multi_gpu
+@pytest.mark.slow
+@pytest.mark.sglang
+@pytest.mark.parametrize(
+    "n_gpus,tp_size",
+    [(4, 2), (8, 2), (8, 4)],
+    ids=["4gpu-tp2", "8gpu-tp2", "8gpu-tp4"],
+)
+def test_awex_megatron_tp_e2e_weight_update(n_gpus, tp_size, tmp_path_factory):
+    """Full round trip: MegatronEngine (DP+TP) → weight-update gateway → SGLang.
+
+    Uses dp_replicated=True in parallelism_strategy so awex knows TP ranks
+    within a DP group all hold the same full parameter (gathered via
+    all_gather_param) and only one rank per group needs to send.
+
+    Minimum 4 GPUs: 2 for SGLang inference, 2 for Megatron (dp=1, tp=2).
+    """
+    if current_platform.device_count() < n_gpus:
+        pytest.skip(f"This test requires {n_gpus} GPUs")
+
+    from areal.api import FinetuneSpec
+    from areal.api.cli_args import (
+        OptimizerConfig,
+        SchedulingSpec,
+        TrainEngineConfig,
+    )
+    from areal.experimental.inference_service.controller.config import (
+        GatewayControllerConfig,
+    )
+    from areal.experimental.inference_service.controller.controller import (
+        GatewayInferenceController,
+    )
+    from areal.experimental.training_service.controller.controller import (
+        GatewayTrainController,
+    )
+    from areal.experimental.weight_update.controller import (
+        WeightUpdateController,
+        WeightUpdateControllerConfig,
+    )
+
+    n_infer = n_gpus // 2
+    n_train = n_gpus - n_infer
+    dp_size = n_train // tp_size
+    if dp_size < 1:
+        pytest.skip(f"Not enough GPUs for dp={dp_size} tp={tp_size}")
+
+    tmp = tmp_path_factory.mktemp("awex_megatron_tp_e2e")
+    model_path = _get_test_model_path()
+
+    scheduler = _make_local_scheduler(
+        tmp, "megatron-tp-e2e", gpu_devices=list(range(n_gpus))
+    )
+
+    inf_config = GatewayControllerConfig(
+        tokenizer_path=model_path,
+        model_path=model_path,
+        backend=f"sglang:d{n_infer}",
+        scheduling_spec=(
+            SchedulingSpec(
+                gpu=1,
+                cmd="python -m areal.experimental.inference_service.guard",
+            ),
+        ),
+        consumer_batch_size=8,
+        max_head_offpolicyness=1024,
+        setup_timeout=300.0,
+        admin_api_key="test-admin",
+    )
+    inf_ctrl = GatewayInferenceController(config=inf_config, scheduler=scheduler)
+
+    train_config = TrainEngineConfig(
+        backend=f"megatron:d{dp_size}t{tp_size}",
+        experiment_name="test-awex-megatron-tp-e2e",
+        trial_name="t0",
+        path=model_path,
+        optimizer=OptimizerConfig(),
+        _version="v2",
+        setup_timeout=300.0,
+        scheduling_spec=(
+            SchedulingSpec(
+                gpu=1,
+                cmd="python -m areal.experimental.training_service.guard",
+                env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"),
+            ),
+        ),
+    )
+    train_ctrl = GatewayTrainController(
+        train_engine="areal.engine.megatron_engine.MegatronLMEngine",
+        config=train_config,
+        scheduler=scheduler,
+    )
+
+    wu_ctrl: WeightUpdateController | None = None
+
+    try:
+        inf_ctrl.initialize(
+            role="rollout",
+            server_args={"mem_fraction_static": 0.7},
+        )
+        inf_worker_urls = list(inf_ctrl._inf_addrs)
+
+        for url in inf_worker_urls:
+            resp = httpx.post(f"{url}/awex/debug/randomize_parameters", timeout=120.0)
+            assert resp.status_code == 200, f"randomize_parameters failed: {resp.text}"
+
+        ft_spec = FinetuneSpec(
+            total_train_epochs=1, dataset_size=100, train_batch_size=2
+        )
+        train_ctrl.initialize(role="actor", ft_spec=ft_spec)
+        train_worker_urls = list(train_ctrl._worker_addrs)
+
+        wu_ctrl = WeightUpdateController(
+            config=WeightUpdateControllerConfig(
+                host="127.0.0.1",
+                request_timeout=300.0,
+            )
+        )
+        wu_ctrl.initialize()
+
+        assert wu_ctrl.health_check(), "Weight update gateway health check failed"
+
+        wu_ctrl.connect(
+            pair_name="test_megatron_tp_e2e",
+            train_worker_urls=train_worker_urls,
+            inference_worker_urls=inf_worker_urls,
+        )
+
+        result = wu_ctrl.update_weights(version=1)
+        assert result.status == "ok"
+        assert result.version == 1
+
+        wu_ctrl.disconnect()
+
+        gen_resp = httpx.post(
+            f"{inf_worker_urls[0]}/generate",
+            json={
+                "text": "Hello",
+                "sampling_params": {"max_new_tokens": 5, "temperature": 0},
+            },
+            timeout=30.0,
+        )
+        assert gen_resp.status_code == 200, (
+            f"Generation failed after weight update: {gen_resp.text}"
+        )
+
+        _validate_weight_update_correctness_megatron(
+            train_worker_urls=train_worker_urls,
+            inf_worker_url=inf_worker_urls[0],
+            param_dir=tmp,
+        )
+
+    finally:
+        if wu_ctrl is not None:
+            wu_ctrl.destroy()
+        train_ctrl.destroy()
+        inf_ctrl.destroy()
+        scheduler.delete_workers(None)