Add Megatron FA4 CP/GDN CP, streaming offload, and optimized LoRA publishing

pyproject.toml

@@ -24,7 +24,7 @@ backend = [
     "bitsandbytes>=0.45.2",
     "unsloth==2026.3.3",
     "unsloth-zoo==2026.3.1",
-    "torch==2.10.0",
+    "torch>=2.11.0",
     "torchao==0.16.0",
     "accelerate==1.7.0",
     "awscli>=1.38.1",
@@ -43,8 +43,10 @@ backend = [
 ]
 megatron = [
     "numpy<2",
-    "torch==2.10.0",
-    "quack-kernels==0.2.5",
+    "torch>=2.11.0",
+    "flash-attn-4==4.0.0b5",
+    "ninja>=1.11.1",
+    "quack-kernels==0.3.7",
     "apex",
     "transformer-engine==2.11.0",
     "transformer-engine-cu12==2.11.0",
@@ -53,6 +55,7 @@ megatron = [
     "pybind11>=2.13.6",
     "megatron-bridge==0.4.0rc0",
     "deep-ep==1.2.1 ; 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",
@@ -76,7 +79,7 @@ tinker = [
     "protobuf>=6.31.1",
     "tinker-cookbook>=0.4.1,<0.5",
     "tinker>=0.21.0,<0.22",
-    "torch==2.10.0",
+    "torch>=2.11.0",
     "transformers==5.2.0",
     "uvicorn>=0.35.0",
     "datrie>=0.8.3",
@@ -152,17 +155,19 @@ override-dependencies = [
     "megatron-core==0.17.0",
     "numpy<2",
     "nvidia-resiliency-ext<0.5",
-    "quack-kernels==0.2.5",
+    "quack-kernels==0.3.7",
     "transformer-engine==2.11.0",
+    "transformers==5.2.0",
+    "torch==2.11.0",
 ]
 exclude-dependencies = ["pynvml", "emerging-optimizers"]
 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"]
-nv-grouped-gemm = ["torch>=2.8.0"]
-transformer-engine-torch = ["torch>=2.8.0"]
+apex = ["torch>=2.11.0"]
+deep-ep = ["torch>=2.11.0"]
+nv-grouped-gemm = ["torch>=2.11.0"]
+transformer-engine-torch = ["torch>=2.11.0"]
 
 [tool.uv.extra-build-variables]
 apex = { APEX_CPP_EXT = "1", APEX_CUDA_EXT = "1", APEX_FAST_LAYER_NORM = "1", APEX_PARALLEL_BUILD = "16", NVCC_APPEND_FLAGS = "--threads 4" }
@@ -180,7 +185,7 @@ requires-dist = []
 
 [[tool.uv.dependency-metadata]]
 name = "transformer-engine-torch"
-version = "0.5.18"
+version = "2.11.0"
 requires-dist = [
     "einops",
     "onnx",
@@ -266,8 +271,15 @@ dev = [
 ]
 
 [tool.uv.sources]
+torch = { index = "pytorch-cu128" }
 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 = "e049cc00c24d03e2ae45d2608c7a44e2d2364e3d" }
 panza = { git = "https://github.com/corbt/panza.git" }
 transformer-engine-torch = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "v2.11", subdirectory = "transformer_engine/pytorch" }
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true

scripts/bump_version.py

@@ -13,15 +13,19 @@
 import subprocess
 import sys
 
+PROJECT_VERSION_RE = re.compile(
+    r'(?ms)^(\[project\]\s+.*?^version = ")(\d+\.\d+\.\d+)(")'
+)
+
 
 def get_current_version():
     """Extract current version from pyproject.toml."""
     pyproject_path = Path(__file__).parent.parent / "pyproject.toml"
     content = pyproject_path.read_text()
-    match = re.search(r'version = "(\d+\.\d+\.\d+)"', content)
+    match = PROJECT_VERSION_RE.search(content)
     if not match:
-        raise ValueError("Could not find version in pyproject.toml")
-    return match.group(1)
+        raise ValueError("Could not find [project] version in pyproject.toml")
+    return match.group(2)
 
 
 def bump_version(current_version, bump_type):
@@ -43,10 +47,11 @@ def update_version(new_version):
     pyproject_path = Path(__file__).parent.parent / "pyproject.toml"
     content = pyproject_path.read_text()
 
-    # Update version
-    new_content = re.sub(
-        r'version = "\d+\.\d+\.\d+"', f'version = "{new_version}"', content
+    new_content, count = PROJECT_VERSION_RE.subn(
+        rf"\g<1>{new_version}\3", content, count=1
     )
+    if count != 1:
+        raise ValueError("Could not update [project] version in pyproject.toml")
 
     pyproject_path.write_text(new_content)
 

src/art/__init__.py

@@ -45,8 +45,12 @@
     import transformers
 
     try:
-        from .transformers.patches import patch_preprocess_mask_arguments
+        from .transformers.patches import (
+            disable_broken_torchvision_for_transformers,
+            patch_preprocess_mask_arguments,
+        )
 
+        disable_broken_torchvision_for_transformers()
         patch_preprocess_mask_arguments()
     except Exception:
         pass
@@ -65,6 +69,7 @@
 from .trajectories import Trajectory, TrajectoryGroup
 from .types import (
     LocalTrainResult,
+    MegatronTopologyConfig,
     Messages,
     MessagesAndChoices,
     ServerlessTrainResult,
@@ -87,6 +92,7 @@
     "LocalBackend",
     "LocalTrainResult",
     "LoRAConfig",
+    "MegatronTopologyConfig",
     "ServerlessBackend",
     "ServerlessTrainResult",
     "Messages",

src/art/_backend_training.py

@@ -9,7 +9,7 @@
     summarize_trajectory_groups,
 )
 from .trajectories import TrajectoryGroup
-from .types import TrainConfig
+from .types import MegatronTopologyConfig, TrainConfig
 
 
 def build_rl_train_configs(
@@ -34,6 +34,7 @@ def build_rl_train_configs(
     scale_learning_rate_by_reward_std_dev: bool | None = None,
     logprob_calculation_chunk_size: int | None = None,
     packed_sequence_length: int | None = None,
+    megatron_topology: MegatronTopologyConfig | dict[str, int | None] | None = None,
     num_trajectories_learning_rate_multiplier_power: float | None = None,
     kl_ref_adapter_path: str | None = None,
 ) -> tuple[TrainConfig, dev.TrainConfig]:
@@ -65,6 +66,10 @@ def build_rl_train_configs(
         dev_config["logprob_calculation_chunk_size"] = logprob_calculation_chunk_size
     if packed_sequence_length is not None:
         dev_config["packed_sequence_length"] = packed_sequence_length
+    if megatron_topology is not None:
+        dev_config["megatron_topology"] = MegatronTopologyConfig.model_validate(
+            megatron_topology
+        ).model_dump(mode="json")
     if num_trajectories_learning_rate_multiplier_power is not None:
         dev_config["num_trajectories_learning_rate_multiplier_power"] = (
             num_trajectories_learning_rate_multiplier_power

src/art/dev/get_model_config.py

@@ -95,4 +95,6 @@ def get_model_config(
         result["trainer_gpu_ids"] = config["trainer_gpu_ids"]
     if "inference_gpu_ids" in config:
         result["inference_gpu_ids"] = config["inference_gpu_ids"]
+    if "megatron_topology" in config:
+        result["megatron_topology"] = config["megatron_topology"]
     return result

src/art/dev/model.py

@@ -1,10 +1,13 @@
 from enum import Enum
-from typing import Literal, NoReturn
+from typing import TYPE_CHECKING, Literal, NoReturn
 
 from typing_extensions import Required, TypedDict
 
 from .engine import EngineArgs
 
+if TYPE_CHECKING:
+    from ..types import MegatronTopologyConfig
+
 RolloutWeightsMode = Literal["lora", "merged"]
 
 
@@ -135,6 +138,7 @@ class InternalModelConfig(TypedDict, total=False):
         chat_template_content_format: vLLM chat template content format.
         chat_template_tool_schema_format: Tool schema rendering format used for
             local training tokenization.
+        megatron_topology: Fixed Megatron parallel topology for this model.
         allow_unvalidated_arch: Permit model-support validation workflows to run
             architectures that are not yet in the supported-model registry.
     """
@@ -152,6 +156,7 @@ class InternalModelConfig(TypedDict, total=False):
     chat_template_path: str
     chat_template_content_format: str
     chat_template_tool_schema_format: Literal["default", "vllm_openai"]
+    megatron_topology: "MegatronTopologyConfig | dict[str, int | None]"
     allow_unvalidated_arch: bool
 
 

src/art/dev/train.py

@@ -25,6 +25,10 @@ class TrainConfig(TypedDict, total=False):
     logprob_calculation_chunk_size: int
     mask_prob_ratio: bool
     max_negative_advantage_importance_sampling_weight: float
+    megatron_topology: dict[
+        Literal["tp", "cp", "ep", "pp", "vpp", "etp"],
+        int | None,
+    ]
     moe_routing_replay_bundle: "MoeRoutingReplayBundle | None"
     moe_routing_replay_path: str | None
     moe_routing_replay_strict: bool

src/art/local/backend.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from contextlib import asynccontextmanager
 import gc
 import hashlib
@@ -9,7 +11,7 @@
 import socket
 import time
 from types import TracebackType
-from typing import Any, AsyncIterator, Iterable, Literal, cast
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterable, Literal, cast
 import warnings
 
 logger = logging.getLogger(__name__)
@@ -22,11 +24,13 @@
 import polars as pl
 import torch
 from tqdm import auto as tqdm
-from transformers import AutoImageProcessor, AutoTokenizer
-from transformers.image_processing_utils import BaseImageProcessor
+from transformers import AutoTokenizer
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 from typing_extensions import Self
 
+if TYPE_CHECKING:
+    from transformers.image_processing_utils import BaseImageProcessor
+
 from art.utils.output_dirs import (
     get_default_art_path,
     get_model_dir,
@@ -66,7 +70,13 @@
     tokenize_trajectory_groups,
 )
 from ..trajectories import Trajectory, TrajectoryGroup
-from ..types import LocalTrainResult, Message, TrainConfig, TrainSFTConfig
+from ..types import (
+    LocalTrainResult,
+    MegatronTopologyConfig,
+    Message,
+    TrainConfig,
+    TrainSFTConfig,
+)
 from ..utils import format_message, get_model_step
 from .adapter_leases import (
     AdapterLeaseManager,
@@ -410,6 +420,16 @@ async def adapter_lease(
         async with pin_inference_step(model.name, step), manager.lease(step):
             yield
 
+    @asynccontextmanager
+    async def adapter_retention_lease(
+        self,
+        model: AnyTrainableModel,
+        step: int,
+    ) -> AsyncIterator[None]:
+        manager = self._adapter_lease_manager(model.name)
+        async with manager.lease(step):
+            yield
+
     async def prune_model_adapters(
         self,
         model: AnyTrainableModel,
@@ -491,6 +511,8 @@ def _get_packed_tensors(
             self._tokenizers[tokenizer_key] = tokenizer
         if model.base_model not in self._image_processors:
             try:
+                from transformers import AutoImageProcessor
+
                 self._image_processors[model.base_model] = (
                     AutoImageProcessor.from_pretrained(model.base_model, use_fast=True)
                 )
@@ -704,6 +726,7 @@ async def train(  # type: ignore[override]
         scale_learning_rate_by_reward_std_dev: bool = False,
         logprob_calculation_chunk_size: int = 1024,
         packed_sequence_length: int | None = None,
+        megatron_topology: MegatronTopologyConfig | None = None,
         num_trajectories_learning_rate_multiplier_power: float = 0.0,
         # Checkpoint behavior
         save_checkpoint: bool = True,
@@ -764,6 +787,9 @@ async def train(  # type: ignore[override]
             packed_sequence_length: Packed sequence length to use for training.
                 When unset, Unsloth keeps the current max-length-rounded-to-2048
                 behavior. Required for Megatron.
+            megatron_topology: Parallel topology for Megatron training. When
+                provided, ART uses it to configure Megatron TP/CP/EP/PP/VPP/ETP
+                before launching the Megatron runtime.
             num_trajectories_learning_rate_multiplier_power: Power for learning
                 rate multiplier based on number of trajectories.
             save_checkpoint: Whether to save a checkpoint after training.
@@ -824,6 +850,7 @@ async def train(  # type: ignore[override]
             scale_learning_rate_by_reward_std_dev=scale_learning_rate_by_reward_std_dev,
             logprob_calculation_chunk_size=logprob_calculation_chunk_size,
             packed_sequence_length=packed_sequence_length,
+            megatron_topology=megatron_topology,
             num_trajectories_learning_rate_multiplier_power=num_trajectories_learning_rate_multiplier_power,
             kl_ref_adapter_path=resolved_kl_ref_adapter_path,
         )