Enable default MoE expert replay for Megatron train/inference parity

.python-version

@@ -1 +1 @@
-3.11
+3.12

pyproject.toml

@@ -3,7 +3,7 @@ name = "openpipe-art"
 version = "0.5.18"
 description = "The OpenPipe Agent Reinforcement Training (ART) library"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 dependencies = [
     "openai>=2.14.0",
     "typer>=0.15.2",
@@ -49,13 +49,14 @@ megatron = [
     "transformer-engine==2.11.0",
     "transformer-engine-cu12==2.11.0",
     "transformer-engine-torch==2.11.0",
-    "megatron-core==0.16.0rc0",
+    "megatron-core==0.17.0",
     "pybind11>=2.13.6",
     "megatron-bridge==0.4.0rc0",
     "deep-ep==1.2.1 ; sys_platform == 'linux'",
     "causal-conv1d==1.6.1 ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version < '3.12'",
     "mamba-ssm==2.3.1 ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version < '3.12'",
     "nvidia-ml-py==13.580.82",
+    "nvidia-modelopt>=0.42.0a0 ; sys_platform != 'darwin'",
     "nvidia-resiliency-ext<0.5 ; sys_platform == 'linux'",
     "ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
 ]
@@ -147,18 +148,18 @@ markers = [
 required-version = ">=0.11.7"
 override-dependencies = [
     "flashinfer-python==0.6.1",
+    "megatron-core==0.17.0",
     "numpy<2",
     "nvidia-resiliency-ext<0.5",
     "quack-kernels==0.2.5",
     "transformer-engine==2.11.0",
 ]
 exclude-dependencies = ["pynvml", "emerging-optimizers"]
-no-build-isolation-package = ["apex", "transformer-engine", "transformer-engine-cu12", "transformer-engine-torch", "megatron-core", "megatron-bridge", "deep-ep", "nv-grouped-gemm"]
+no-build-isolation-package = ["apex", "transformer-engine", "transformer-engine-cu12", "transformer-engine-torch", "megatron-bridge", "deep-ep", "nv-grouped-gemm"]
 
 [tool.uv.extra-build-dependencies]
 apex = ["torch>=2.8.0"]
 deep-ep = ["torch>=2.8.0"]
-megatron-core = ["pybind11"]
 nv-grouped-gemm = ["torch>=2.8.0"]
 transformer-engine-torch = ["torch>=2.8.0"]
 

src/art/auto_trajectory.py

@@ -8,6 +8,7 @@
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from .openai import init_chat_completion, update_chat_completion
+from .preprocessing.moe_routing import attach_moe_routing_metadata_to_choice
 from .trajectories import History, Trajectory
 
 logger = logging.getLogger(__name__)
@@ -105,7 +106,13 @@ def handle_httpx_response(self, response: httpx._models.Response) -> None:
                 chat_completion = parse_sse_to_chat_completion(content)
                 choice = chat_completion.choices[0]
             else:
-                choice = Choice(**json.loads(content)["choices"][0])
+                response_payload = json.loads(content)
+                choice = Choice(**response_payload["choices"][0])
+                attach_moe_routing_metadata_to_choice(
+                    choice=choice,
+                    response_payload=response_payload,
+                    choice_index=0,
+                )
         except (json.JSONDecodeError, KeyError, ValueError) as e:
             logger.debug(f"Failed to parse response content: {e}")
             return

src/art/dev/engine.py

@@ -125,6 +125,7 @@ class EngineArgs(TypedDict, total=False):
     override_generation_config: dict[str, Any] | None
     enable_sleep_mode: bool
     enable_expert_parallel: bool
+    enable_return_routed_experts: bool
     model_impl: str
 
     calculate_kv_scales: bool | None

src/art/local/backend.py

@@ -9,7 +9,7 @@
 import socket
 import time
 from types import TracebackType
-from typing import AsyncIterator, Iterable, Literal, cast
+from typing import Any, AsyncIterator, Iterable, Literal, cast
 import warnings
 
 logger = logging.getLogger(__name__)
@@ -147,6 +147,7 @@ def __init__(
         in_process: bool = False,
         path: str | None = None,
         gpu_cost_per_hour_usd: float | None = None,
+        enable_expert_replay: bool = True,
     ) -> None:
         """
         Initializes a local, directory-based Backend interface at the given path.
@@ -162,12 +163,15 @@ def __init__(
                 automatic `costs/gpu` accounting on train steps. When unset,
                 ART auto-detects supported GPU types (H200 at $3/hr today) and
                 skips GPU cost logging for unknown devices instead of guessing.
+            enable_expert_replay: For supported MoE Megatron training, capture
+                vLLM routed experts and replay them in Megatron. Defaults to True.
         """
         self._in_process = in_process
         self._path = path or get_default_art_path()
         self._gpu_cost_per_hour_usd = (
             float(gpu_cost_per_hour_usd) if gpu_cost_per_hour_usd is not None else None
         )
+        self._enable_expert_replay = enable_expert_replay
         os.makedirs(self._path, exist_ok=True)
 
         # Other initialization
@@ -182,6 +186,27 @@ def __init__(
             "default"
         )
 
+    def _model_uses_expert_replay(self, model: AnyTrainableModel) -> bool:
+        if not self._enable_expert_replay or not self._supports_result_packing:
+            return False
+        from ..megatron.model_support.registry import (
+            UnsupportedModelArchitectureError,
+            model_uses_expert_parallel,
+        )
+
+        allow_unvalidated_arch = bool(
+            (model._internal_config or dev.InternalModelConfig()).get(
+                "allow_unvalidated_arch", False
+            )
+        )
+        try:
+            return model_uses_expert_parallel(
+                model.base_model,
+                allow_unvalidated_arch=allow_unvalidated_arch,
+            )
+        except UnsupportedModelArchitectureError:
+            return False
+
     def supports_automatic_train_step_metrics(self) -> bool:
         return True
 
@@ -455,6 +480,7 @@ def _get_packed_tensors(
         plot_tensors: bool,
         packed_sequence_length: int | None,
         logprob_calculation_chunk_size: int,
+        include_moe_routing: bool = False,
     ) -> PackedTensors | None:
         internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
@@ -547,6 +573,7 @@ def _get_packed_tensors(
             truncate_long_results=False,
             advantage_balance=advantage_balance,
             pack_results=self._supports_result_packing,
+            include_moe_routing=include_moe_routing,
         )
         if (
             not allow_training_without_logprobs
@@ -603,6 +630,11 @@ async def _prepare_backend_for_training(
         config_dict: dict = dict(config or {})
         internal_config = cast(dev.InternalModelConfig, model._internal_config or {})
         _apply_configured_chat_template_server_args(config_dict, internal_config)
+        if self._model_uses_expert_replay(model):
+            engine_args = dict(config_dict.get("engine_args", {}))
+            engine_args["enable_return_routed_experts"] = True
+            engine_args["async_scheduling"] = False
+            config_dict["engine_args"] = engine_args
         server_args = dict(config_dict.get("server_args", {}))
 
         # Avoid binding collisions on busy hosts when no explicit port is provided.
@@ -850,7 +882,7 @@ async def _train_model(
             summary,
             include_trainable_groups=True,
         )
-
+        include_moe_routing = self._model_uses_expert_replay(model)
         packed_tensors = self._get_packed_tensors(
             model,
             trajectory_groups,
@@ -864,6 +896,7 @@ async def _train_model(
             logprob_calculation_chunk_size=dev_config.get(
                 "logprob_calculation_chunk_size", 1024
             ),
+            include_moe_routing=include_moe_routing,
         )
         if packed_tensors is None:
             print(
@@ -927,17 +960,34 @@ async def _train_model(
         disk_packed_tensors = packed_tensors_to_dir(
             packed_tensors, f"{get_model_dir(model=model, art_path=self._path)}/tensors"
         )
-        # Note: scale_learning_rate_by_reward_std_dev is now handled by the frontend (Model.train())
-        grad_accumulation_sequences = max(
-            1, int(config.grad_accumulation_sequences or 1)
+        service_dev_config = cast(dev.TrainConfig, {**dev_config})
+        grad_accumulation_sequences = await self._resolve_grad_accumulation_sequences(
+            service,
+            config,
         )
+        if include_moe_routing:
+            from ..megatron.routing_replay import (
+                build_moe_routing_replay_bundle_from_packed_tensors,
+            )
+
+            routing_replay_dir = (
+                f"{get_model_dir(model=model, art_path=self._path)}/tensors/"
+                "moe_routing_replay"
+            )
+            build_moe_routing_replay_bundle_from_packed_tensors(
+                packed_tensors=packed_tensors,
+                global_grad_accumulation_sequences=grad_accumulation_sequences,
+            ).to_dir(routing_replay_dir)
+            service_dev_config["moe_routing_replay_path"] = routing_replay_dir
+            service_dev_config["moe_routing_replay_strict"] = True
+        # Note: scale_learning_rate_by_reward_std_dev is now handled by the frontend (Model.train())
         fallback_gradient_steps = math.ceil(
             disk_packed_tensors["num_sequences"] / grad_accumulation_sequences
         )
         pbar = tqdm.tqdm(total=fallback_gradient_steps, desc="train")
         reported_gradient_steps: int | None = None
         async for result in service.train(
-            disk_packed_tensors, config, dev_config, verbose
+            disk_packed_tensors, config, service_dev_config, verbose
         ):
             raw_num_gradient_steps = result.pop(TRAIN_GRADIENT_STEPS_KEY, None)
             if raw_num_gradient_steps is not None:
@@ -965,6 +1015,20 @@ async def _train_model(
         if verbose:
             print("_train_model complete")
 
+    async def _resolve_grad_accumulation_sequences(
+        self,
+        service: ModelService,
+        config: TrainConfig,
+    ) -> int:
+        resolver = getattr(
+            cast(Any, service),
+            "resolve_global_grad_accumulation_sequences",
+            None,
+        )
+        if callable(resolver):
+            return max(1, int(await resolver(config)))
+        return max(1, int(config.grad_accumulation_sequences or 1))
+
     # Note: _get_reward_std_dev_learning_rate_multiplier and _log_metrics
     # have been moved to the Model class (frontend)
 

src/art/megatron/lora.py

@@ -962,6 +962,14 @@ def __init__(
             b_parallel_spec=b_parallel_spec,
             allreduce=False,
         )
+        component_size = (
+            linear_fc1.out_features * _get_shard_world_size("expert_tp")
+        ) // 2
+        _set_lora_shard_strategy_metadata(
+            self.lora.B_T,
+            strategy="componentwise",
+            component_sizes=(component_size, component_size),
+        )
 
     def forward(
         self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor

src/art/megatron/provider.py

@@ -103,13 +103,6 @@ def _apply_default_parallel_topology(provider: GPTModelProvider) -> None:
     provider.expert_tensor_parallel_size = 1
 
 
-def _etp_ep_parallel_domain_size(provider: GPTModelProvider) -> int:
-    return (
-        cast(int, provider.expert_tensor_parallel_size)
-        * provider.expert_model_parallel_size
-    )
-
-
 def _apply_art_training_runtime_prepare_defaults(provider: GPTModelProvider) -> None:
     provider.recompute_granularity = "full"
     provider.recompute_method = "uniform"
@@ -119,7 +112,7 @@ def _apply_art_training_runtime_prepare_defaults(provider: GPTModelProvider) ->
 
 
 def _apply_art_training_runtime_finalize_defaults(provider: GPTModelProvider) -> None:
-    if _etp_ep_parallel_domain_size(provider) <= 1:
+    if provider.expert_model_parallel_size <= 1:
         return
     # use DeepEP for MoE expert comm. comm can be the same amount of time as actual MLP
     # compute, so these are very beneficial