Optimize end-to-end Pipeline RL throughput

.gitattributes

@@ -0,0 +1 @@
+src/art/megatron/_hybrid_ep/** -whitespace

.github/workflows/package-install.yml

@@ -29,6 +29,23 @@ jobs:
       - name: Build wheel
         run: python scripts/build_package.py --wheel
 
+      - name: Smoke test base wheel import
+        run: |
+          wheel_path="$(python - <<'PY'
+          from pathlib import Path
+
+          print(next(Path("dist").glob("openpipe_art-*.whl")).resolve())
+          PY
+          )"
+
+          project_dir="$(mktemp -d)"
+          cd "$project_dir"
+          uv init --name art-base-install-smoke --python 3.12 --bare
+          uv add "openpipe-art @ file://${wheel_path}"
+          uv run python -c "import importlib.util; assert importlib.util.find_spec('numpy') is None; assert importlib.util.find_spec('torch') is None; import art; print(art.__name__)"
+          uv add "weave==0.52.41"
+          uv run python -c "from weave.trace.settings import override_settings; import art"
+
       - name: Smoke test uv add + sync for backend extra
         run: |
           wheel_path="$(python - <<'PY'

.github/workflows/trainer-rank-gpu.yml

@@ -28,7 +28,7 @@ jobs:
           HEAD_SHA: ${{ github.event.pull_request.head.sha }}
         run: |
           set -euo pipefail
-          pattern='^src/art/megatron/(prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py)'
+          pattern='^src/art/megatron/(_hybrid_ep/|hybrid_ep_setup\.py|setup\.sh|prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py)'
           changed="$(git diff --name-only "${BASE_SHA}...${HEAD_SHA}")"
           critical="$(printf '%s\n' "${changed}" | grep -E "${pattern}" || true)"
           if [ -n "${critical}" ]; then

THIRD-PARTY-NOTICES

@@ -41,3 +41,13 @@ BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 DAMAGE.
+
+This project vendors a modified HybridEP runtime derived from DeepEP:
+
+- Repository: https://github.com/deepseek-ai/DeepEP
+- License: MIT License
+- Vendored path: src/art/megatron/_hybrid_ep
+
+The applicable license text is retained at:
+
+- src/art/megatron/_hybrid_ep/LICENSE

dev/yes-no-maybe-fork-pipeline.py

@@ -24,7 +24,7 @@
 
 import art
 from art.local import LocalBackend
-from art.pipeline_trainer import PipelineTrainer
+from art.pipeline_trainer import PipelineRuntimeConfig, PipelineTrainer
 from art.utils.deployment.wandb import deploy_wandb
 from art.utils.output_dirs import get_model_dir, get_step_checkpoint_dir
 
@@ -298,9 +298,11 @@ def on_eval(step: int, rate: float) -> None:
                 rollout_fn=make_rollout_fn(),
                 scenarios=scenario_iter(),
                 config=None,
-                num_rollout_workers=len(PROMPTS),
-                min_batch_size=len(PROMPTS),
-                max_batch_size=len(PROMPTS),
+                pipeline=PipelineRuntimeConfig(
+                    num_rollout_workers=len(PROMPTS),
+                    min_batch_size=len(PROMPTS),
+                    max_batch_size=len(PROMPTS),
+                ),
                 max_steps=TRAIN_STEPS - start_step,
                 learning_rate=1e-4,
                 loss_fn="cispo",
@@ -368,9 +370,11 @@ def on_forked_eval(step: int, rate: float) -> None:
         rollout_fn=make_rollout_fn(),
         scenarios=scenario_iter(),
         config=None,
-        num_rollout_workers=len(PROMPTS),
-        min_batch_size=len(PROMPTS),
-        max_batch_size=len(PROMPTS),
+        pipeline=PipelineRuntimeConfig(
+            num_rollout_workers=len(PROMPTS),
+            min_batch_size=len(PROMPTS),
+            max_batch_size=len(PROMPTS),
+        ),
         max_steps=2,
         learning_rate=1e-4,
         loss_fn="cispo",

pyproject.toml

@@ -13,7 +13,7 @@ dependencies = [
     "typing-extensions>=4.13",
     "typer>=0.15.2",
     "litellm>=1.71.1,<=1.82.0",
-    "weave>=0.52.24",
+    "weave>=0.52.41",
     "polars>=1.26.0",
     "tblib>=3.0.0",
     "nest-asyncio>=1.6.0",
@@ -60,15 +60,17 @@ megatron = [
     "transformer-engine-torch==2.11.0",
     "megatron-core==0.17.0",
     "pybind11>=2.13.6",
+    "setuptools>=78.1.0",
     "megatron-bridge==0.4.0rc0",
-    "deep-ep==1.2.1 ; sys_platform == 'linux'",
+    "nvidia-cuda-cccl-cu12==12.9.27 ; sys_platform == 'linux'",
     "tilelang==0.1.10 ; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "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'",
     "transformers==5.12.1",
+    "scipy>=1.17.0",
     "ml-dtypes>=0.5.0 ; python_full_version < '3.13'",
 ]
 
@@ -141,6 +143,9 @@ exclude = [
     "**/*.pyc",
 ]
 
+[tool.ruff]
+extend-exclude = ["src/art/megatron/_hybrid_ep"]
+
 [tool.ruff.lint]
 select = ["I"]
 
@@ -178,12 +183,11 @@ override-dependencies = [
     "torch==2.11.0",
 ]
 exclude-dependencies = ["pynvml", "emerging-optimizers", "causal-conv1d", "mamba-ssm"]
-no-build-isolation-package = ["apex", "transformer-engine", "transformer-engine-cu12", "transformer-engine-torch", "megatron-bridge", "deep-ep", "nv-grouped-gemm"]
+no-build-isolation-package = ["apex", "transformer-engine", "transformer-engine-cu12", "transformer-engine-torch", "megatron-bridge", "nv-grouped-gemm"]
 
 [tool.uv.extra-build-dependencies]
 apex = ["torch>=2.11.0"]
-deep-ep = ["torch>=2.11.0"]
-megatron-core = ["pybind11"]
+megatron-core = ["pybind11", "setuptools"]
 nv-grouped-gemm = ["torch>=2.11.0"]
 transformer-engine-torch = ["torch>=2.11.0"]
 
@@ -196,11 +200,6 @@ name = "apex"
 version = "0.1"
 requires-dist = ["packaging"]
 
-[[tool.uv.dependency-metadata]]
-name = "deep-ep"
-version = "1.2.1+9af0e0d"
-requires-dist = []
-
 # Keep Bridge's runtime deps explicit here and let ART's megatron extra own the
 # Transformers pin validated by model-support handlers in this branch.
 [[tool.uv.dependency-metadata]]
@@ -252,6 +251,9 @@ requires-dist = [
     "transformer-engine-cu12",
 ]
 
+[tool.ty.src]
+exclude = ["src/art/megatron/_hybrid_ep/**"]
+
 [tool.ty.environment]
 python-version = "3.12"
 
@@ -333,7 +335,6 @@ dev = [
 [tool.uv.sources]
 torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }]
 apex = { git = "https://github.com/NVIDIA/apex.git", rev = "25.09" }
-deep-ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "v1.2.1" }
 flash-attn-4 = { url = "https://files.pythonhosted.org/packages/24/f7/01ee2576ce41f9884d291ee21861ef194afc0b2b1ce3bd175fc7a6e1b133/flash_attn_4-4.0.0b5-py3-none-any.whl" }
 megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "e1a207ac757e5d0ed94d8ffbe1cbd28e81d8c084" }
 panza = { git = "https://github.com/corbt/panza.git" }

scripts/build_package.py

@@ -122,6 +122,10 @@ def verify_wheel(wheel: Path) -> None:
         "art/_vllm_runtime/manifest.json",
         "art/_vllm_runtime/pyproject.toml",
         "art/_vllm_runtime/uv.lock",
+        "art/megatron/_hybrid_ep/LICENSE",
+        "art/megatron/_hybrid_ep/setup.py",
+        "art/megatron/_hybrid_ep/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh",
+        "art/megatron/_hybrid_ep/deep_ep/hybrid_ep_buffer.py",
     }
     with zipfile.ZipFile(wheel) as archive:
         names = set(archive.namelist())
@@ -158,6 +162,10 @@ def verify_sdist(sdist: Path) -> None:
         "src/art/_vllm_runtime/manifest.json",
         "src/art/_vllm_runtime/pyproject.toml",
         "src/art/_vllm_runtime/uv.lock",
+        "src/art/megatron/_hybrid_ep/LICENSE",
+        "src/art/megatron/_hybrid_ep/setup.py",
+        "src/art/megatron/_hybrid_ep/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh",
+        "src/art/megatron/_hybrid_ep/deep_ep/hybrid_ep_buffer.py",
     }
     with tarfile.open(sdist) as archive:
         names = set(archive.getnames())

scripts/ci/trainer-rank-gpu.sky.yaml

@@ -8,6 +8,7 @@ resources:
 
 setup: |
   uv sync --frozen --extra megatron --group dev
+  uv run --no-sync python -m art.megatron.hybrid_ep_setup
 
 run: |
   timeout --signal=TERM --kill-after=30s 20m \

src/art/__init__.py

@@ -62,12 +62,20 @@
 from .backend import Backend
 from .batches import trajectory_group_batches
 from .dev import LoRAConfig
+from .errors import LocalServingUnavailableError
 from .gather import gather_trajectories, gather_trajectory_groups
 from .megatron.runtime_config import (
     get_megatron_runtime_config,
     init_megatron_runtime_config,
 )
+from .metrics import (
+    PIPELINE_RL_DASHBOARD_DEFAULT_METRICS,
+    PIPELINE_RL_METRIC_DEFINITIONS,
+    PIPELINE_RL_SCORE_METRICS,
+    MetricDefinition,
+)
 from .model import Model, TrainableModel
+from .pipeline_tuner import PipelineAutotuneConfig, PipelineRuntimeConfig
 from .serverless import ServerlessBackend
 from .trajectories import (
     AnthropicMessagesHistory,
@@ -123,8 +131,15 @@
     "Backend",
     "LocalTrainResult",
     "LoRAConfig",
+    "LocalServingUnavailableError",
     "MegatronRuntimeConfig",
     "MegatronTopologyConfig",
+    "MetricDefinition",
+    "PipelineAutotuneConfig",
+    "PipelineRuntimeConfig",
+    "PIPELINE_RL_DASHBOARD_DEFAULT_METRICS",
+    "PIPELINE_RL_METRIC_DEFINITIONS",
+    "PIPELINE_RL_SCORE_METRICS",
     "get_megatron_runtime_config",
     "init_megatron_runtime_config",
     "ServerlessBackend",

src/art/_backend_training.py

@@ -37,11 +37,13 @@ def build_rl_train_configs(
     packed_sequence_length: int | None = None,
     num_trajectories_learning_rate_multiplier_power: float | None = None,
     kl_ref_adapter_path: str | None = None,
+    optimizer_save_interval: int = 5,
 ) -> tuple[TrainConfig, dev.TrainConfig]:
     config = TrainConfig(
         learning_rate=learning_rate,
         kl_penalty_coef=kl_penalty_coef,
         kl_penalty_source=kl_penalty_source,
+        optimizer_save_interval=optimizer_save_interval,
     )
     dev_config: dev.TrainConfig = {
         "advantage_balance": advantage_balance,
@@ -96,8 +98,16 @@ def aggregate_rl_training_metrics(
 ) -> dict[str, float]:
     groups_list = list(trajectory_groups)
     avg_metrics = average_metric_samples(training_metrics)
+    tokens_per_second = avg_metrics.pop("tokens_per_second", None)
+    if (
+        tokens_per_second is not None
+        and "throughput/train_packed_tok_per_s" not in avg_metrics
+    ):
+        avg_metrics["throughput/train_packed_tok_per_s"] = float(tokens_per_second)
     summary = summarize_trajectory_groups(groups_list)
-    avg_metrics.setdefault("time/step_trainer_s", time.monotonic() - trainer_started)
+    avg_metrics.setdefault(
+        "time/step_backend_train_s", time.monotonic() - trainer_started
+    )
     avg_metrics.update(
         {
             key: value

src/art/adapter_leases.py

@@ -2,25 +2,58 @@
 from contextvars import ContextVar
 from typing import AsyncIterator
 
-_pinned_inference_steps: ContextVar[dict[str, int]] = ContextVar(
-    "art_pinned_inference_steps",
+from pydantic import BaseModel, ConfigDict
+
+
+class InferenceTarget(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    step: int
+    model_name: str | None = None
+
+
+_pinned_inference_targets: ContextVar[dict[str, InferenceTarget]] = ContextVar(
+    "art_pinned_inference_targets",
     default={},
 )
 
 
+def in_flight_lora_name(model_name: str) -> str:
+    return f"{model_name}:active"
+
+
 def pinned_inference_step(model_name: str) -> int | None:
-    return _pinned_inference_steps.get().get(model_name)
+    target = _pinned_inference_targets.get().get(model_name)
+    return None if target is None else target.step
+
+
+def pinned_inference_name(model_name: str, step: int | None = None) -> str | None:
+    target = _pinned_inference_targets.get().get(model_name)
+    if target is None or (step is not None and step != target.step):
+        return None
+    return target.model_name
 
 
 @asynccontextmanager
 async def pin_inference_step(
     model_name: str,
     step: int,
 ) -> AsyncIterator[None]:
-    pinned_steps = dict(_pinned_inference_steps.get())
-    pinned_steps[model_name] = step
-    token = _pinned_inference_steps.set(pinned_steps)
+    async with pin_inference_target(model_name, step=step):
+        yield
+
+
+@asynccontextmanager
+async def pin_inference_target(
+    model_name: str,
+    *,
+    step: int,
+    inference_name: str | None = None,
+) -> AsyncIterator[None]:
+    targets = dict(_pinned_inference_targets.get())
+    targets[model_name] = InferenceTarget(step=step, model_name=inference_name)
+    token = _pinned_inference_targets.set(targets)
     try:
         yield
     finally:
-        _pinned_inference_steps.reset(token)
+        _pinned_inference_targets.reset(token)

src/art/backend.py

@@ -1,3 +1,4 @@
+from contextlib import AbstractAsyncContextManager
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -44,6 +45,12 @@ async def _prepare_backend_for_training(
         config: dev.OpenAIServerConfig | None,
     ) -> tuple[str, str]: ...
 
+    def exact_adapter_lease(
+        self,
+        model: AnyTrainableModel,
+        step: int,
+    ) -> AbstractAsyncContextManager[None]: ...
+
     # Backends intentionally expose backend-specific optional training arguments.
     # Callable[..., ...] preserves that extensibility without falsely requiring
     # every implementation to accept every other backend's keyword arguments.

src/art/dev/__init__.py

@@ -5,6 +5,7 @@
     InternalModelConfig,
     LoRAConfig,
     PeftArgs,
+    RolloutWeightUpdateMode,
     TinkerArgs,
     TinkerNativeArgs,
     TinkerTrainingClientArgs,
@@ -26,6 +27,7 @@
     "InitArgs",
     "LoRAConfig",
     "PeftArgs",
+    "RolloutWeightUpdateMode",
     "TinkerArgs",
     "TinkerNativeArgs",
     "TinkerTrainingClientArgs",

src/art/dev/get_model_config.py

@@ -34,6 +34,7 @@ def get_model_config(
 
     dedicated = is_dedicated_mode(config)
     rollout_weights_mode = config.get("rollout_weights_mode", "lora")
+    rollout_weight_update_mode = config.get("rollout_weight_update_mode", "step_lora")
 
     if dedicated:
         enable_sleep_mode = False
@@ -99,6 +100,7 @@ def get_model_config(
         engine_args=engine_args,
         lora_config=merged_lora_config,
         rollout_weights_mode=rollout_weights_mode,
+        rollout_weight_update_mode=rollout_weight_update_mode,
         tinker_args=config.get("tinker_args"),
         trainer_args=trainer_args,
     )