feat(engine): add CPU-staged Megatron AdamW

.github/workflows/test-areal.yml

@@ -307,6 +307,17 @@ jobs:
     steps:
       - uses: actions/checkout@v6
 
+      - name: Sync AWEX dependency
+        run: |
+          AWEX_REQUIREMENT=$("$VIRTUAL_ENV/bin/python" - <<'PY'
+          import tomllib
+          with open("pyproject.toml", "rb") as f:
+              dependencies = tomllib.load(f)["project"]["dependencies"]
+          print(next(dep for dep in dependencies if dep.startswith("awex==")))
+          PY
+          )
+          uv pip install --python "$VIRTUAL_ENV/bin/python" --no-deps "$AWEX_REQUIREMENT"
+
       - name: Validate Docker installation
         run: |
           export PATH="/opt/.venv/bin:$PATH"

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(
@@ -1174,6 +1213,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

@@ -39,6 +39,77 @@
 logger = getLogger("AwexColocate")
 
 
+@torch.no_grad()
+def _pack_tensors_for_cuda_ipc(
+    tensors: list[torch.Tensor],
+    max_group_bytes: int = 5 * 1024**3,
+) -> tuple[list[torch.Tensor], list[dict]]:
+    """Pack tensors into bounded, single-allocation CUDA IPC groups.
+
+    AWEX's packer uses ``torch.cat(...).clone()`` and only checks the size
+    limit after adding the next tensor.  The redundant allocation doubles the
+    transient pressure and makes it much more likely that the eventual export
+    occupies a split training segment.  CUDA IPC exports the whole segment, so
+    one live payload can pin tens of GiB of inactive splits.
+
+    Plan groups first, allocate each flat destination exactly once, and copy
+    shaped source views directly into it.  Grouping by dtype preserves AWEX's
+    low-handle-count format and remains compatible with
+    ``reconstruct_tensors_from_groups``.
+    """
+    if not tensors:
+        return [], []
+
+    by_dtype: dict[torch.dtype, list[tuple[int, torch.Tensor]]] = {}
+    for index, tensor in enumerate(tensors):
+        by_dtype.setdefault(tensor.dtype, []).append((index, tensor))
+
+    planned_groups: list[list[tuple[int, torch.Tensor]]] = []
+    for dtype_group in by_dtype.values():
+        current: list[tuple[int, torch.Tensor]] = []
+        current_bytes = 0
+        for index, tensor in dtype_group:
+            tensor_bytes = tensor.numel() * tensor.element_size()
+            if current and current_bytes + tensor_bytes > max_group_bytes:
+                planned_groups.append(current)
+                current = []
+                current_bytes = 0
+            current.append((index, tensor))
+            current_bytes += tensor_bytes
+            if current_bytes >= max_group_bytes:
+                planned_groups.append(current)
+                current = []
+                current_bytes = 0
+        if current:
+            planned_groups.append(current)
+
+    packed_groups: list[torch.Tensor] = []
+    metadata: list[dict] = []
+    for group_index, group in enumerate(planned_groups):
+        dtype = group[0][1].dtype
+        device = group[0][1].device
+        total_numel = sum(tensor.numel() for _, tensor in group)
+        packed = torch.empty(total_numel, dtype=dtype, device=device)
+        offset = 0
+        for original_index, tensor in group:
+            size = tensor.numel()
+            packed[offset : offset + size].view(tensor.shape).copy_(tensor)
+            metadata.append(
+                {
+                    "original_index": original_index,
+                    "shape": tensor.shape,
+                    "dtype": tensor.dtype,
+                    "group_index": group_index,
+                    "offset": offset,
+                    "size": size,
+                }
+            )
+            offset += size
+        packed_groups.append(packed)
+
+    return packed_groups, metadata
+
+
 def resolve_physical_gpu_id(relative_gpu_id: int) -> int:
     """Map a CUDA-masked relative device index to its physical GPU id.
 
@@ -96,6 +167,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 +422,20 @@ 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
@@ -392,7 +468,7 @@ def execute_colocate_weight_update(self, version: int) -> None:
         )
 
         with cuda_ipc_allocation():
-            group_tensors, metadata = group_tensors_by_shape_and_dtype(tensors)
+            group_tensors, metadata = _pack_tensors_for_cuda_ipc(tensors)
         torch.cuda.synchronize()
         logger.info(
             "Grouped into %d tensor groups for IPC serialization", len(group_tensors)
@@ -430,6 +506,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 +524,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 +574,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 +668,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)