feat(engine): add CPU-staged Megatron AdamW

.github/workflows/test-areal.yml

@@ -74,12 +74,12 @@ jobs:
         variant: ${{ fromJson(needs.determine-variants.outputs.matrix) }}
     env:
       CONTAINER_IMAGE: ghcr.io/areal-project/areal-runtime:${{ inputs.image_tag || 'dev' }}-${{ matrix.variant }}
-      RUNNER_LABELS: gcp-a2-highgpu-2g,variant-${{ matrix.variant }}
+      RUNNER_LABELS: gcp-a2-highgpu-2g-run-${{ github.run_id }}-${{ github.run_attempt }},variant-${{ matrix.variant }}
     steps:
       - name: Set instance variables
         id: vars
         run: |
-          echo "instance_name=gcp-runner-${{ github.run_id }}-${{ matrix.variant }}" >> "$GITHUB_OUTPUT"
+          echo "instance_name=gcp-runner-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.variant }}" >> "$GITHUB_OUTPUT"
 
       - name: Authenticate to Google Cloud
         uses: google-github-actions/auth@v3
@@ -297,7 +297,7 @@ jobs:
       contents: read
     runs-on:
       - self-hosted
-      - gcp-a2-highgpu-2g
+      - "gcp-a2-highgpu-2g-run-${{ github.run_id }}-${{ github.run_attempt }}"
       - "variant-${{ matrix.variant }}"
     timeout-minutes: 120
     env:
@@ -347,7 +347,7 @@ jobs:
 
       - name: Delete runner instance
         env:
-          INSTANCE_NAME: gcp-runner-${{ github.run_id }}-${{ matrix.variant }}
+          INSTANCE_NAME: gcp-runner-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.variant }}
         run: |
           # Look up instance zone dynamically since matrix job outputs
           # cannot be consumed per-element by downstream matrix jobs.

areal/api/cli_args.py

@@ -914,6 +914,36 @@ class FP8EngineConfig:
     )
 
 
+@dataclass
+class CPUStagedOffloadConfig:
+    """Bounded GPU staging for CPU-resident Megatron optimizer state."""
+
+    enabled: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Keep optimizer state in pinned CPU memory and stream bounded "
+                "update units through GPU staging buffers. The optimizer algorithm "
+                "is selected independently by optimizer.type. Megatron only."
+            )
+        },
+    )
+    buffer_count: int = field(
+        default=2,
+        metadata={"help": "Number of reusable GPU optimizer staging buffers."},
+    )
+    bucket_size_mb: float = field(
+        default=128.0,
+        metadata={"help": "Maximum size in MiB of one GPU optimizer staging unit."},
+    )
+
+    def __post_init__(self) -> None:
+        if self.buffer_count < 1:
+            raise ValueError("cpu_staged_offload.buffer_count must be at least 1")
+        if self.bucket_size_mb <= 0:
+            raise ValueError("cpu_staged_offload.bucket_size_mb must be positive")
+
+
 @dataclass
 class MegatronEngineConfig:
     """Configuration for Megatron-LM training framework.
@@ -962,6 +992,15 @@ class MegatronEngineConfig:
     main_params_dtype: str = "float32"
     exp_avg_dtype: str = "float32"
     exp_avg_sq_dtype: str = "float32"
+    cpu_staged_offload: CPUStagedOffloadConfig = field(
+        default_factory=CPUStagedOffloadConfig,
+        metadata={
+            "help": (
+                "Keep the selected optimizer's state CPU-resident and use bounded "
+                "GPU staging buffers for optimizer steps."
+            )
+        },
+    )
 
     # Checkpointing Configuration
     async_save: bool = field(
@@ -1180,6 +1219,11 @@ class MegatronEngineConfig:
     )
 
     def __post_init__(self) -> None:
+        if self.cpu_staged_offload.enabled and self.async_save:
+            raise ValueError(
+                "megatron.async_save is not supported with "
+                "megatron.cpu_staged_offload.enabled=true"
+            )
         if self.enable_mtp_training and not self.enable_mtp:
             raise ValueError("enable_mtp_training requires enable_mtp=True")
         if self.mtp_only:

areal/dataset/__init__.py

@@ -30,6 +30,7 @@
     "virl39k",
     "hh-rlhf",
     "torl_data",
+    "dapo-math-17k",
     "swe_sft",
 ]
 
@@ -174,6 +175,16 @@ def _get_custom_dataset(
             max_length=max_length,
             **kwargs,
         )
+    elif "dapo-math-17k" in path.lower() and type == "rl":
+        from .dapo_math_17k import get_dapo_math_17k_rl_dataset
+
+        return get_dapo_math_17k_rl_dataset(
+            path=path,
+            split=split,
+            tokenizer=tokenizer,
+            max_length=max_length,
+            **kwargs,
+        )
     elif (
         _SWE_PATH_PATTERN.search(path.lower()) or _is_swe_sft_artifact(path)
     ) and type == "sft":

areal/dataset/dapo_math_17k.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+from datasets import Dataset, load_dataset
+
+from areal.utils.hf_utils import apply_chat_template
+
+
+def get_dapo_math_17k_rl_dataset(
+    path: str,
+    split: str,
+    tokenizer,
+    max_length: int | None = None,
+) -> Dataset:
+    """Load local DAPO-Math-17k parquet data for RL training."""
+    del split
+    dataset_root = Path(path)
+    parquet_files = sorted((dataset_root / "data").glob("*.parquet"))
+    json_files = sorted(dataset_root.glob("*.jsonl")) + sorted(
+        dataset_root.glob("*.json")
+    )
+    if parquet_files:
+        dataset = load_dataset(
+            "parquet",
+            data_files=[str(data_file) for data_file in parquet_files],
+            split="train",
+        )
+    elif json_files:
+        dataset = load_dataset(
+            "json",
+            data_files=[str(data_file) for data_file in json_files],
+            split="train",
+        )
+    else:
+        raise FileNotFoundError(
+            f"No parquet or JSON data files found under {dataset_root}"
+        )
+
+    def process(sample):
+        answer = sample.get("label")
+        if answer is None:
+            answer = sample["reward_model"]["ground_truth"]
+        return {"messages": sample["prompt"], "answer": str(answer)}
+
+    dataset = dataset.map(process, remove_columns=dataset.column_names)
+
+    if max_length is not None:
+        if tokenizer is None:
+            raise ValueError("tokenizer is required when max_length is set")
+
+        def filter_length(sample):
+            prompt_ids = apply_chat_template(
+                tokenizer, sample["messages"], add_generation_prompt=True, tokenize=True
+            )
+            return len(prompt_ids) <= max_length
+
+        dataset = dataset.filter(filter_length)
+
+    return dataset

areal/engine/awex/colocate_writer.py

@@ -96,6 +96,9 @@ def __init__(
         self._infer_world_size: int | None = None
         self._num_infer_engines: int | None = None
         self._logical_train_rank: int | None = None
+        self._pending_ipc_export: (
+            tuple[list[torch.Tensor], list[torch.Tensor]] | None
+        ) = None
 
     @property
     def residency(self) -> MegatronWeightResidency:
@@ -348,18 +351,21 @@ def execute_colocate_weight_update(self, version: int) -> None:
           4. Signal all_training_offloaded_weights (reader waits for this)
           5. share_memory_() + cuda_ipc_serialize → put to MetaServer
           6. Wait for weights_update_finished (reader done copying)
-          7. Clean up shared tensors → signal write_finished
-          8. Wait for all infer engines to finish (finished_weights_update_engines)
+          7. Release exported tensors, then signal write_finished
+          8. Wait for all infer engines to finish the update lifecycle
         """
         from awex.util.tensor_util import (
             cuda_ipc_serialize,
             group_tensors_by_shape_and_dtype,
             release_tensors,
         )
 
-        # Reclaim any IPC-exported blocks from the previous version whose
-        # peer mappings closed after our last collect (belt-and-braces with the
-        # collect in this method's finally block).
+        if self._pending_ipc_export is not None:
+            raise RuntimeError(
+                "Previous CUDA IPC export is still awaiting reader acknowledgement"
+            )
+
+        # Reclaim any legacy IPC-exported blocks from a previous version.
         torch.cuda.ipc_collect()
 
         # Free optimizer states + grad buffers BEFORE reloading the train
@@ -430,6 +436,8 @@ def execute_colocate_weight_update(self, version: int) -> None:
         release_tensors(owned)
         del tensors, owned
         parameters.clear()
+        gc.collect()
+        torch.cuda.empty_cache()
 
         self._residency.release_memory(tags=["weights"])
 
@@ -446,6 +454,9 @@ def execute_colocate_weight_update(self, version: int) -> None:
         )
 
         group_shared = [t.share_memory_() for t in group_tensors]
+        # Retain producer storage until the reader acknowledges completion.
+        # On failure or timeout the reader may still hold an IPC mapping.
+        self._pending_ipc_export = (group_tensors, group_shared)
         serialized_weights = cuda_ipc_serialize((group_shared, metadata, names))
         torch.cuda.synchronize()
         logger.info("CUDA IPC serialization complete, putting to MetaServer")
@@ -493,20 +504,28 @@ def execute_colocate_weight_update(self, version: int) -> None:
             self._meta_server_client.delete_if_exists(serialized_weights_key)
             logger.info("Got done signal from infer side: %s", update_finished_key)
         finally:
-            release_tensors(group_tensors)
-            release_tensors(group_shared)
-            del group_tensors, group_shared
             torch.cuda.synchronize()
-            gc.collect()
-            # Storages exported via cudaIpcGetMemHandle park in PyTorch's
-            # CudaIPCSentDataLimbo when freed and are NOT returned to the
-            # allocator until ipc_collect() confirms the peer closed its
-            # mapping. GiB-scale group tensors are exported every version;
-            # without collection the train-process residual grows by GBs per
-            # version, eating the colocated rollout's prefill headroom until
-            # its logits all-gather OOMs.
-            torch.cuda.ipc_collect()
-            torch.cuda.empty_cache()
+
+        # AWEX drops imported tensor references and synchronizes before ACK.
+        # Release both retained and local references before write_finished,
+        # freeing packed-weight memory before the reader's remaining lifecycle.
+        allocated_before = torch.cuda.memory_allocated()
+        reserved_before = torch.cuda.memory_reserved()
+        self._pending_ipc_export = None
+        del group_tensors, group_shared
+        gc.collect()
+        torch.cuda.ipc_collect()
+        torch.cuda.empty_cache()
+        logger.info(
+            "CUDA IPC cleanup after reader acknowledgement: rank=%d device=%d "
+            "allocated=%.3f->%.3f GiB reserved=%.3f->%.3f GiB",
+            dist.get_rank(self._engine.cpu_group),
+            torch.cuda.current_device(),
+            allocated_before / 1024**3,
+            torch.cuda.memory_allocated() / 1024**3,
+            reserved_before / 1024**3,
+            torch.cuda.memory_reserved() / 1024**3,
+        )
 
         write_finished_key = f"write_finished{key_suffix}"
         self._meta_server_client.put_object(write_finished_key, True)
@@ -579,6 +598,13 @@ def _convert_parameters(self) -> dict[str, torch.Tensor]:
 
     # ── Memory management (manual offload) ────────────────────────────────
 
+    def checkpoint_residency(self, *, with_model: bool, with_optimizer: bool):
+        """Temporarily restore only resources required by a checkpoint."""
+        return self._residency.checkpoint_residency(
+            with_model=with_model,
+            with_optimizer=with_optimizer,
+        )
+
     def release_memory(self, tags: list[str] | None = None) -> None:
         """Compatibility delegate for callers of the former combined adapter."""
         self._residency.release_memory(tags)

areal/engine/awex/sglang_plugin.py

@@ -168,18 +168,20 @@ def _resolve_transfer_rank(
     node_id: int,
     nnodes: int,
     instance_world_size: int,
+    server_base_gpu_id: int = 0,
 ) -> int:
     """Resolve the inference rank in AWEX's global transfer world.
 
     A one-GPU SGLang server can use the colocated actor's inherited global
     rank when CUDA device isolation remaps its only GPU to device zero. For a
-    multi-GPU server, every TP/PP scheduler inherits the same environment, so
-    its scheduler-local physical GPU identity must be used instead.
+    multi-GPU server, every TP/PP scheduler inherits the same environment.
+    Subtract the server's logical base GPU before adding the dense transfer base.
     """
     explicit_rank = os.environ.get("AWEX_TRANSFER_RANK")
     if explicit_rank is not None:
         transfer_rank = int(explicit_rank)
     else:
+        explicit_base = os.environ.get("AWEX_TRANSFER_RANK_BASE")
         env_rank = os.environ.get("RANK")
         env_world_size = os.environ.get("WORLD_SIZE")
         if (
@@ -189,6 +191,12 @@ def _resolve_transfer_rank(
             and int(env_world_size) == infer_world_size
         ):
             transfer_rank = int(env_rank)
+        elif explicit_base is not None:
+            n_gpus_per_node = max(1, infer_world_size // nnodes)
+            local_gpu_rank = gpu_id - server_base_gpu_id
+            transfer_rank = (
+                node_id * n_gpus_per_node + int(explicit_base) + local_gpu_rank
+            )
         else:
             n_gpus_per_node = max(1, infer_world_size // nnodes)
             transfer_rank = node_id * n_gpus_per_node + gpu_id
@@ -973,10 +981,10 @@ def _init_receiver_from_meta_server(self, meta_server_addr: str):
 
         receiver = self._require_receiver()
 
-        # `physical_gpu_id` is node-local. Multi-node colocate needs a globally
-        # unique transfer rank that stays physically paired with the training
-        # process. SGLang may run with logical gpu_id=0 under CUDA_VISIBLE_DEVICES
-        # isolation, so do not use scheduler.gpu_id for AWEX keys.
+        # Transfer ranks are dense process coordinates, while physical GPU IDs
+        # are only used for device pairing and MetaServer keys. The launcher
+        # supplies the server's dense node-local base when GPU visibility is
+        # isolated or non-contiguous.
         gpu_id = self._logical_gpu_id()
         physical_gpu_id = self._physical_gpu_id()
         node_id = int(os.environ.get("SLURM_NODEID", "0"))
@@ -1010,6 +1018,9 @@ def _init_receiver_from_meta_server(self, meta_server_addr: str):
             node_id=node_id,
             nnodes=nnodes,
             instance_world_size=self._instance_world_size(),
+            server_base_gpu_id=self._int_attr(
+                self._scheduler.server_args, "base_gpu_id", 0
+            ),
         )
 
         logger.info(