feat(colocate): support AWEX colocated actor-rollout training by Le8r0nJames · Pull Request #1500 · areal-project/AReaL

Le8r0nJames · 2026-07-10T15:38:18Z

Description

Add shared-GPU colocated RL training: the Megatron actor and the SGLang rollout engine time-share the same devices, coordinated through AWEX weight transfer.

AWEX weight-sync writer with tag-based offload/onload orchestration, and an SGLang engine plugin implementing the AWEX colocate reader protocol (pause/retract generation around the train phase, kv-cache/weights release-resume).
Colocation scheduling support in the Slurm scheduler and controllers, including recover handling (communicator-friendly checkpoint load ordering, pause/offload protocol on the recover path).
Megatron model registration and fast HF checkpoint loading for the Bailing hybrid MoE family, plus a runtime SGLang patch enabling its radix cache (examples/swe/apply_sgl_radix_cache_patch.py).
Rollout/PPO observability: multi-turn prompt length, trained-token log-prob diffs, and rejection-mask token bookkeeping.

Related Issue

N/A

Type of Change

✨ New feature

Checklist

I have read the Contributing Guide
Pre-commit hooks pass
Relevant unit tests pass (tests/test_stats_tracker.py, tests/test_functional.py, tests/test_ppo_stats.py, tests/test_train_controller.py)
Documentation updated (CLI reference regenerated)
Branch is up to date with main
Self-reviewed
This PR was created by a coding agent via /create-pr
This PR is a breaking change

Additional Context

Verified on multi-node colocated RL runs: training statistics align with the separated-GPU baseline, and replaying an identical batch produces bitwise-identical float64 training stats. Integration requires multi-node GPU hardware, so end-to-end colocation runs are not covered by CI; unit coverage focuses on the loss/stats math and controller logic.

Validation

Colocated GSM8K on SGLang 0.5.10.post1, in two shapes.

Two nodes, sixteen GPUs:

model: Qwen3-30B-A3B
allocation_mode: megatron[actor]:(attn:d2p2t4|ffn:d2p2e4)|sglang[rollout]:d4t4p1
n_samples: 8, max_new_tokens: 1024, train batch: 16
sglang: ep_size 4, mem_fraction_static 0.65, memory saver on, CUDA graphs on

One 50-step run on this branch and two 100-step runs on the same colocate path
all finished cleanly. For the 50-step run: 56.2 s/step, every expected
weight-update completion signal present (50 steps x 16 ranks), and no
tracebacks, illegal memory accesses, invalid-argument errors, or fully-idle
assertion failures.

One node, two GPUs, matching the example test added here:

model: Qwen3-0.6B
allocation_mode: megatron[actor]:d2|sglang[rollout]:d1t2p1
n_samples: 2, max_new_tokens: 256, train batch: 2
sglang: mem_fraction_static 0.3, memory saver on, CUDA graphs on

Ten steps at 19 s/step, clean exit, no errors.

Ablation established which changes the SGLang 0.5.10 port actually needs.
Guarding the removed decode-stat hooks and draining before memory release are
both required: without the guard the scheduler dies at startup, and without the
drain the third training step trips SGLang's fully-idle assertion. Explicitly
managing the new cuda_graph memory-saver tag turned out not to be required,
so it is not included here.

Unit tests cover the pause sequence and the physical-GPU-id mapping that keys
the CUDA IPC handoff, and pass against both SGLang 0.5.9 and 0.5.10.post1.

gemini-code-assist

Code Review

This pull request introduces an AWEX colocate adapter and reader to enable zero-copy weight transfer via CUDA IPC and MetaServer between MegatronEngine and SGLang. It adds manual GPU-to-CPU memory offloading/onloading, updates Slurm scheduling to support colocation, and optimizes memory usage during HF weight saving. The review feedback highlights two key issues: a potential deadlock in physical GPU ID detection when CUDA_VISIBLE_DEVICES contains multiple comma-separated GPUs, and a missing initialization guard in get_weight_metadata that could lead to malformed MetaServer keys.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-10T15:40:43Z

+        cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+        if cuda_visible and "," not in cuda_visible:
+            self._physical_gpu_id = int(cuda_visible)
+        else:
+            self._physical_gpu_id = torch.cuda.current_device()


The physical GPU ID detection logic fails when CUDA_VISIBLE_DEVICES contains a comma-separated list of multiple GPUs (e.g., "2,3"). In this scenario, "," not in cuda_visible evaluates to False, causing the code to fall back to torch.cuda.current_device(). This returns the relative device index (e.g., 1) rather than the actual physical GPU ID (e.g., 3), leading to a mismatch with the inference side and potential silent deadlocks during weight transfer.

To fix this, parse the comma-separated list and map the current device index to the corresponding physical GPU ID.

Suggested change

cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")

if cuda_visible and "," not in cuda_visible:

self._physical_gpu_id = int(cuda_visible)

else:

self._physical_gpu_id = torch.cuda.current_device()

cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")

self._physical_gpu_id = torch.cuda.current_device()

if cuda_visible:

try:

gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()]

dev_idx = torch.cuda.current_device()

if dev_idx < len(gpu_ids):

self._physical_gpu_id = gpu_ids[dev_idx]

except ValueError:

pass

Fixed in ca17ad8 per the suggestion: the physical GPU id is now resolved by indexing the parsed CUDA_VISIBLE_DEVICES list with torch.cuda.current_device(), with a fallback to the relative index for UUID-style entries.

# Conflicts: # areal/api/cli_args.py # areal/infra/controller/train_controller.py # areal/utils/stats_tracker.py # docs/en/cli_reference.md # docs/zh/cli_reference.md

gemini-code-assist · 2026-07-23T12:33:04Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

sitabulaixizawaluduo · 2026-07-23T14:46:19Z

+
+        update_finished_key = f"weights_update_finished{key_suffix}"
+        try:
+            self._meta_server_client.get_object(


Can we consider adding an alert message for timeouts here?

Added in ca17ad8: the writer now logs an explicit error (timeout budget + MetaServer key + likely cause) before re-raising when the inference side never consumes the published weights.

sitabulaixizawaluduo · 2026-08-06T02:58:30Z

+import sys
+
+
+def _early_set_alloc_conf() -> None:


AWEX_ACTOR_ALLOC_CONF appears only as a consumer in this diff — I don't see it defined in cli_args.py or documented anywhere, so it is currently an env-only knob that users have to set through SchedulingSpec.env_vars. If it moves into the scheduler it would be natural to promote it to a real config field at the same time. Was leaving it out of the config surface intentional?

Good question, and the answer is that it should never have been a config field in the first place. AWEX_ACTOR_ALLOC_CONF is removed entirely in 835fea7.

It could not become a real config field. PYTORCH_CUDA_ALLOC_CONF is read once when the allocator is first initialised, and importing areal already initialises CUDA, so by the time Hydra has parsed a config the setting is frozen (torch.cuda.memory_snapshot()[...]['is_expandable'] shows this directly). The only place early enough was argv sniffing at the very top of areal/__init__.py, which is exactly the import-time side effect worth avoiding. So 'env-only' was not an oversight in the config surface, it was inherent to doing it inside the process at all.

The knob belongs to whoever starts the process, so it is now a per-role scheduling_spec.env_vars entry. That is also why this PR depends on #1584: without the user-mapping re-apply, an explicit env_vars entry is overwritten by the framework defaults.

Two follow-ups from digging into this:

sglang_plugin.py used to rewrite the allocator config under __main__. That was dead code for the same reason (the module imports areal.utils at line 35, so the config is already frozen). It is now an assertion at the top of the module instead, which fails loudly rather than silently doing nothing.

tests/test_alloc_conf_import_side_effects.py greps areal/ and fails if the name comes back.

The last thing this file still carried was a cosmetic import reformat left over from the removal. Reverted in 51105df, so areal/infra/rpc/rpc_server.py is out of this PR's diff now.

sitabulaixizawaluduo · 2026-08-06T03:02:28Z

 }

+
+def query_terminal_state_sacct(job_id: int) -> JobState | None:
