diff --git a/.claude/rules/testing.md b/.claude/rules/testing.md index c6c5d93c1b..ac768fcb4e 100644 --- a/.claude/rules/testing.md +++ b/.claude/rules/testing.md @@ -33,7 +33,7 @@ def test___(): - Use `torch.distributed.fake_pg` for unit tests - Mock `dist.get_rank()` and `dist.get_world_size()` explicitly -- Don't mock internals of FSDP/DTensor — use integration tests +- Don't mock internals of FSDP/DTensor - use integration tests ## GPU Test Constraints @@ -58,4 +58,4 @@ def test___(): - Use `torch.testing.assert_close()` for tensor comparison - Specify `rtol`/`atol` explicitly for numerical tests -- Avoid bare `assert tensor.equal()` — no useful error message +- Avoid bare `assert tensor.equal()` - no useful error message diff --git a/.claude/skills/add-unit-tests/SKILL.md b/.claude/skills/add-unit-tests/SKILL.md index 07a828a774..2eae4bd2b6 100644 --- a/.claude/skills/add-unit-tests/SKILL.md +++ b/.claude/skills/add-unit-tests/SKILL.md @@ -150,7 +150,7 @@ def test_gpu_function(): - Use `torch.testing.assert_close()` for tensor comparison - Specify `rtol`/`atol` explicitly for numerical tests -- Avoid bare `assert tensor.equal()` — no useful error message +- Avoid bare `assert tensor.equal()` - no useful error message ## Reference Implementations diff --git a/AGENTS.md b/AGENTS.md index 701fba29b0..c92765cb93 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,6 @@ -# AGENTS.md — AReaL Agent Operations Guide +# AGENTS.md - AReaL Agent Operations Guide ## TL;DR for coding agents @@ -22,54 +22,54 @@ When unsure, leave a `TODO(agent)` comment and note the constraint in your respo ## Repository map -- `.claude/` — Configuration directory for AI coding agents: - - `.claude/agents/` — Specialized agent definitions (planner, code-verifier, +- `.claude/` - Configuration directory for AI coding agents: + - `.claude/agents/` - Specialized agent definitions (planner, code-verifier, simple-code-reviewer, engine experts, algorithm-expert, launcher-scheduler-expert) - - `.claude/commands/` — User-invocable commands (create-pr, gen-commit-msg, pr-review) - - `.claude/rules/` — Project-wide code quality standards (api-config, code-style, + - `.claude/commands/` - User-invocable commands (create-pr, gen-commit-msg, pr-review) + - `.claude/rules/` - Project-wide code quality standards (api-config, code-style, distributed, testing) - - `.claude/skills/` — Guided development workflows (add-dataset, add-workflow, + - `.claude/skills/` - Guided development workflows (add-dataset, add-workflow, add-reward, add-unit-tests, debug-distributed) - - `.claude/hooks/` — Event-driven automation hooks - - `.claude/data/` — Agent working data and caches - - `.claude/plans/` — Implementation plans generated by planner agent -- `areal/` — Core Python package housing APIs, controllers, engines, workflows, and + - `.claude/hooks/` - Event-driven automation hooks + - `.claude/data/` - Agent working data and caches + - `.claude/plans/` - Implementation plans generated by planner agent +- `areal/` - Core Python package housing APIs, controllers, engines, workflows, and shared utilities: - - `areal/api/` — Contracts for workflows, engines, schedulers, IO structs, and + - `areal/api/` - Contracts for workflows, engines, schedulers, IO structs, and CLI/config dataclasses. - - `areal/infra/` — Single controller implementation, async orchestration primitives, + - `areal/infra/` - Single controller implementation, async orchestration primitives, and hardware/platform abstractions for CPU/GPU/NPU runtimes. - - `areal/dataset/` — Stateful dataset loaders (GSM8K, Geometry3K, CLEVR, HH-RLHF, + - `areal/dataset/` - Stateful dataset loaders (GSM8K, Geometry3K, CLEVR, HH-RLHF, TORL, etc.) and utilities that feed rollout jobs safely. - - `areal/engine/` — Training backends (FSDP2, Megatron, PPO, SFT, reward modeling) and + - `areal/engine/` - Training backends (FSDP2, Megatron, PPO, SFT, reward modeling) and inference adapters (SGLang, vLLM remote engines). - - `areal/experimental/` — Prototype engines/workflows that evolve quickly; expect + - `areal/experimental/` - Prototype engines/workflows that evolve quickly; expect breaking changes. - - `areal/launcher/` — Launch specs for local, Ray, and Slurm clusters, plus + - `areal/launcher/` - Launch specs for local, Ray, and Slurm clusters, plus SGLang/vLLM inference server launchers and container guidance. - - `areal/models/` — Model-specific adapters (Megatron-Core layers, Transformers + - `areal/models/` - Model-specific adapters (Megatron-Core layers, Transformers wrappers, custom heads). - - `areal/reward/` — Built-in reward functions (GSM8K, Geometry3K, CLEVR, etc.), math + - `areal/reward/` - Built-in reward functions (GSM8K, Geometry3K, CLEVR, etc.), math parsers, and helpers; wrap slow logic with `AsyncRewardWrapper`. - - `areal/scheduler/` — Placement and allocation policies aligned with launcher + - `areal/scheduler/` - Placement and allocation policies aligned with launcher configs. - - `areal/tests/` — Focused unit/integration suites (many require GPUs or mocked + - `areal/tests/` - Focused unit/integration suites (many require GPUs or mocked distributed backends). - - `areal/tools/` — Developer utilities and maintenance scripts tied to the core + - `areal/tools/` - Developer utilities and maintenance scripts tied to the core package. - - `areal/utils/` — Cross-cutting helpers for logging, tensor ops, stats tracking, + - `areal/utils/` - Cross-cutting helpers for logging, tensor ops, stats tracking, checkpoints, and recovery. - - `areal/workflow/` — Concrete rollout agents implementing `RolloutWorkflow`: + - `areal/workflow/` - Concrete rollout agents implementing `RolloutWorkflow`: multi-turn, RLVR, vision RLVR workflows, plus `openai_agent/` for OpenAI Agent-style implementations. -- `assets/` — Figures and other static assets referenced across docs and blogs. -- `blog/` — Release notes and update write-ups documenting project progress. -- `docs/` — Jupyter Book source for https://inclusionai.github.io/AReaL/ plus CLI +- `assets/` - Figures and other static assets referenced across docs and blogs. +- `blog/` - Release notes and update write-ups documenting project progress. +- `docs/` - Jupyter Book source for https://inclusionai.github.io/AReaL/ plus CLI reference generators. -- `examples/` — End-to-end training scripts and launcher recipes for math reasoning, +- `examples/` - End-to-end training scripts and launcher recipes for math reasoning, multi-turn conversations, VLM, RLHF alignment, agent-based workflows, LoRA fine-tuning, SkyPilot deployment, and more. -- `notebook/` — Reference notebooks (outputs stripped via pre-commit) for quick +- `notebook/` - Reference notebooks (outputs stripped via pre-commit) for quick experimentation. ## Distributed operations & tooling diff --git a/CLAUDE.md b/CLAUDE.md index e8735d41a8..34130318f0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,4 +1,4 @@ -# CLAUDE.md — AReaL +# CLAUDE.md - AReaL ## WHAT: Project Overview @@ -9,15 +9,15 @@ learning. **Core Directories**: -- `areal/` — Core package - - `api/` — Config dataclasses, workflow/engine contracts - - `engine/` — FSDP2, Megatron, SGLang/vLLM adapters - - `workflow/` — RolloutWorkflow implementations - - `reward/` — Reward functions - - `dataset/` — Dataset loaders - - `utils/` — Logging, tensor ops, checkpoints -- `examples/` — Training scripts and configs -- `docs/` — Jupyter Book source +- `areal/` - Core package + - `api/` - Config dataclasses, workflow/engine contracts + - `engine/` - FSDP2, Megatron, SGLang/vLLM adapters + - `workflow/` - RolloutWorkflow implementations + - `reward/` - Reward functions + - `dataset/` - Dataset loaders + - `utils/` - Logging, tensor ops, checkpoints +- `examples/` - Training scripts and configs +- `docs/` - Jupyter Book source ## WHY: Purpose diff --git a/areal/README.md b/areal/README.md index 2616748fdf..474ac5ed57 100644 --- a/areal/README.md +++ b/areal/README.md @@ -23,17 +23,17 @@ principles: To achieve an *algorithm-first* and *lightweight* design while maintaining efficiency, AReaL is guided by seven core principles: -1. **Native asynchronous RL training** — Built from the ground up for decoupled +1. **Native asynchronous RL training** - Built from the ground up for decoupled generation and training -1. **System-abstracted design** — Minimize exposure to low-level system concepts like +1. **System-abstracted design** - Minimize exposure to low-level system concepts like "PlacementGroup" -1. **PyTorch-centric approach** — Use native PyTorch types without unnecessary +1. **PyTorch-centric approach** - Use native PyTorch types without unnecessary abstractions -1. **Transparent orchestration** — Make the flow of operations clear and understandable -1. **Developer-friendly navigation** — Enable easy access to implementation details via +1. **Transparent orchestration** - Make the flow of operations clear and understandable +1. **Developer-friendly navigation** - Enable easy access to implementation details via IDE features (Ctrl+click) -1. **Ecosystem compatibility** — Integrate seamlessly with existing ML/RL tools -1. **Single-file customization** — Support RL pipeline modifications within a single +1. **Ecosystem compatibility** - Integrate seamlessly with existing ML/RL tools +1. **Single-file customization** - Support RL pipeline modifications within a single file ## Architecture diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 4f0f89c0a3..fc068985b3 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1538,17 +1538,24 @@ class RecoverConfig(_Timer): metadata={ "help": "Recovery mode for the launcher. " "Options: " - "'disabled': Never recover from previous runs. " - "'auto': Automatically recover from previous runs if recover info and checkpoints are available. " - "'fault': Only recover from previous runs if the new run fails. " - "'resume': Force to resume, raise an error if no recover info was found. Never resume if failed again." + "'on' or 'auto': Automatically recover from previous runs if recover info and checkpoints are available. " + "'off' or 'disabled': Never recover from previous runs." }, ) retries: int = field( default=3, - metadata={"help": "Number of recovery retries (auto/fault modes only)."}, + metadata={"help": "Number of recovery retries when recovery is enabled."}, ) + def __post_init__(self): + valid_modes = {"on", "off", "auto", "disabled"} + if self.mode not in valid_modes: + raise ValueError( + f"Invalid recover mode '{self.mode}'. " + f"Valid options: {valid_modes}. " + f"Note: 'fault' and 'resume' modes have been removed." + ) + @dataclass class WandBConfig: diff --git a/areal/launcher/local.py b/areal/launcher/local.py index 91042778d0..a401f09d4e 100644 --- a/areal/launcher/local.py +++ b/areal/launcher/local.py @@ -367,7 +367,6 @@ def local_main(config, run_id: int = 0): gpu = nprocs = alloc_mode.train.world_size _env_vars = dict( AREAL_LLM_SERVER_ADDRS=",".join(server_addrs), - AREAL_RECOVER_RUN=str(int(is_recover_run)), ) if alloc_mode.gen_backend == "sglang": # Required by NCCL weight update group. @@ -414,7 +413,7 @@ def local_main(config, run_id: int = 0): recover_this = ( e.reason in recover_states and run_id < config.recover.retries - and config.recover.mode in ["auto", "fault"] + and config.recover.mode in ("on", "auto") ) if recover_this: time.sleep(RECOVER_TIME_INTERVAL) diff --git a/areal/launcher/ray.py b/areal/launcher/ray.py index a95872ea73..568a216e1c 100644 --- a/areal/launcher/ray.py +++ b/areal/launcher/ray.py @@ -564,7 +564,6 @@ def torch_env_hook(n_tasks: int, placement_group: PlacementGroup) -> list[dict]: _env_vars = dict( AREAL_LLM_SERVER_ADDRS=",".join(llm_addrs), - AREAL_RECOVER_RUN=str(int(is_recover_run)), ) if allocation_mode.gen_backend == "sglang": # Required by NCCL weight update group. @@ -617,7 +616,7 @@ def torch_env_hook(n_tasks: int, placement_group: PlacementGroup) -> list[dict]: recover_this = ( e.reason in recover_states and run_id < config.recover.retries - and config.recover.mode in ["auto", "fault"] + and config.recover.mode in ("on", "auto") ) if recover_this: time.sleep(RECOVER_TIME_INTERVAL) diff --git a/areal/launcher/slurm.py b/areal/launcher/slurm.py index 4f68bdc20e..c3378eb1e4 100644 --- a/areal/launcher/slurm.py +++ b/areal/launcher/slurm.py @@ -609,7 +609,6 @@ def _build_trainer_cmds( # launch trainers _env_vars = dict( AREAL_LLM_SERVER_ADDRS=",".join(llm_addrs), - AREAL_RECOVER_RUN=str(int(is_recover_run)), ) if allocation_mode.gen_backend == "sglang": # Required by NCCL weight update group. @@ -666,7 +665,7 @@ def _build_trainer_cmds( recover_this = ( e.reason in recover_states and run_id < config.recover.retries - and config.recover.mode in ["auto", "fault"] + and config.recover.mode in ("on", "auto") ) if recover_this: time.sleep(RECOVER_TIME_INTERVAL) diff --git a/areal/scheduler/slurm.py b/areal/scheduler/slurm.py index ccc2c67b4a..6faf745dac 100644 --- a/areal/scheduler/slurm.py +++ b/areal/scheduler/slurm.py @@ -1,6 +1,5 @@ import asyncio import getpass -import os import re import shlex import subprocess @@ -401,7 +400,6 @@ def _prepare_worker_specs( # Amend environment variables for sch in schedulings: # AReaL env var forwarding - sch.env_vars["AREAL_RECOVER_RUN"] = os.getenv("AREAL_RECOVER_RUN", str(0)) if self.enable_tms_offload: sch.env_vars.update(get_tms_env_vars()) sch.env_vars.update(get_env_vars()) diff --git a/areal/tests/test_recover.py b/areal/tests/test_recover.py new file mode 100644 index 0000000000..5155fd4033 --- /dev/null +++ b/areal/tests/test_recover.py @@ -0,0 +1,169 @@ +"""Tests for the recovery configuration and functionality.""" + +import tempfile + +import pytest + +from areal.api.cli_args import RecoverConfig +from areal.utils.recover import check_if_auto_recover, check_if_recover + + +class TestRecoverConfig: + """Tests for RecoverConfig dataclass validation.""" + + def test_default_values(self): + """Test that default values are set correctly.""" + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot="/tmp", + ) + assert config.mode == "disabled" + assert config.retries == 3 + + @pytest.mark.parametrize("mode", ["on", "off", "auto", "disabled"]) + def test_valid_modes(self, mode): + """Test that all valid modes are accepted.""" + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot="/tmp", + mode=mode, + ) + assert config.mode == mode + + @pytest.mark.parametrize("mode", ["fault", "resume", "invalid", "ON", "OFF", ""]) + def test_invalid_modes(self, mode): + """Test that invalid modes raise ValueError with helpful message.""" + with pytest.raises(ValueError) as exc_info: + RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot="/tmp", + mode=mode, + ) + error_msg = str(exc_info.value) + assert f"Invalid recover mode '{mode}'" in error_msg + assert "fault" in error_msg and "resume" in error_msg # Migration hint + + +class TestCheckIfRecover: + """Tests for the check_if_recover function.""" + + @pytest.mark.parametrize("mode", ["disabled", "off"]) + def test_disabled_modes_return_false(self, mode): + """Test that disabled modes always return False.""" + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot="/tmp", + mode=mode, + ) + # Should return False regardless of run_id + assert check_if_recover(config, 0) is False + assert check_if_recover(config, 1) is False + assert check_if_recover(config, 10) is False + + @pytest.mark.parametrize("mode", ["on", "auto"]) + def test_enabled_modes_check_for_checkpoint(self, mode): + """Test that enabled modes check for existing checkpoints.""" + with tempfile.TemporaryDirectory() as tmpdir: + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode=mode, + ) + # No checkpoint exists, should return False + assert check_if_recover(config, 0) is False + + @pytest.mark.parametrize("run_id", [0, 1, 5, 100]) + def test_run_id_parameter_unused(self, run_id): + """Test that run_id parameter doesn't affect the result.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Test with disabled mode + config_disabled = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="disabled", + ) + assert check_if_recover(config_disabled, run_id) is False + + # Test with enabled mode (no checkpoint) + config_enabled = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="on", + ) + # Result should be the same regardless of run_id + result = check_if_recover(config_enabled, run_id) + assert result == check_if_recover(config_enabled, 0) + + +class TestCheckIfAutoRecover: + """Tests for the check_if_auto_recover function.""" + + def test_no_checkpoint_returns_false(self): + """Test that missing checkpoint returns False.""" + with tempfile.TemporaryDirectory() as tmpdir: + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="on", + ) + assert check_if_auto_recover(config) is False + + def test_empty_directory_returns_false(self): + """Test that empty directory (no checkpoint) returns False.""" + with tempfile.TemporaryDirectory() as tmpdir: + config = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="on", + ) + assert check_if_auto_recover(config) is False + + +class TestModeEquivalence: + """Tests to verify mode equivalences (on=auto, off=disabled).""" + + def test_on_equals_auto(self): + """Test that 'on' and 'auto' modes behave identically.""" + with tempfile.TemporaryDirectory() as tmpdir: + config_on = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="on", + ) + config_auto = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="auto", + ) + # Both should return the same result + assert check_if_recover(config_on, 0) == check_if_recover(config_auto, 0) + + def test_off_equals_disabled(self): + """Test that 'off' and 'disabled' modes behave identically.""" + with tempfile.TemporaryDirectory() as tmpdir: + config_off = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="off", + ) + config_disabled = RecoverConfig( + experiment_name="test_exp", + trial_name="test_trial", + fileroot=tmpdir, + mode="disabled", + ) + # Both should return False + assert check_if_recover(config_off, 0) is False + assert check_if_recover(config_disabled, 0) is False diff --git a/areal/utils/recover.py b/areal/utils/recover.py index e4de736608..c3fed58359 100644 --- a/areal/utils/recover.py +++ b/areal/utils/recover.py @@ -176,7 +176,7 @@ def dump( processor: AutoProcessor | None = None, base_model_path: str | None = None, ): - if self.config.mode == "disabled": + if self.config.mode in ("disabled", "off"): return # currently only support recover on one engine if not self.freq_ctl.check( @@ -223,9 +223,7 @@ def load( weight_update_meta: WeightUpdateMeta | None = None, inference_engine_update_from: str = "default", ) -> RecoverInfo | None: - if self.config.mode == "disabled": - return - if os.environ.get("AREAL_RECOVER_RUN", "0") != "1": + if self.config.mode in ("disabled", "off"): return if inference_engine is not None and weight_update_meta is None: raise ValueError("Weight update meta is required for recovery.") @@ -324,8 +322,8 @@ def _load_checkpoint( def check_if_auto_recover(config: RecoverConfig) -> bool: - # This method is called only by launchers to check if the experiment should be a recover run - # when "recover_mode" is auto. + # This method is called by check_if_recover to check if the experiment should + # recover from a previous run when recovery is enabled ("on" or "auto" mode). experiment_name = config.experiment_name trial_name = config.trial_name fileroot = config.fileroot @@ -363,16 +361,20 @@ def check_if_auto_recover(config: RecoverConfig) -> bool: return False -def check_if_recover(config: RecoverConfig, run_id: int) -> bool: - # This method is called by the launcher to check if the experiment should be a recover run - # when "recover_mode" is not disabled. - if config.mode == "disabled": +def check_if_recover(config: RecoverConfig, _run_id: int) -> bool: + """Check if the experiment should be a recover run. + + When recovery is enabled ('on' or 'auto'), this checks if valid recover + info and checkpoints are available for automatic recovery. + + Args: + config: Recovery configuration. + _run_id: Unused. Kept for API compatibility. + + Returns: + True if the experiment should recover from a previous run. + """ + if config.mode in ("disabled", "off"): return False - elif config.mode == "auto": - return check_if_auto_recover(config) - elif config.mode == "fault": - return run_id > 0 - elif config.mode == "resume": - return True - else: - raise ValueError(f"Unknown recover mode: {config.mode}") + # Both "on" and "auto" use auto-recovery behavior + return check_if_auto_recover(config) diff --git a/blog/AReaL_v0_1.md b/blog/AReaL_v0_1.md index 89641f7dd0..476f528b01 100644 --- a/blog/AReaL_v0_1.md +++ b/blog/AReaL_v0_1.md @@ -132,8 +132,8 @@ highlight some of the major research topics we are working on. enhancements? - How to properly curate the distillation dataset for cold-start to achieve the best RL results? -- **Boosting Reasoning with Auxiliary Training Signals.** Can auxiliary signals — such - as Process Reward Models (PRMs) and expert-curated labels — be adopted in RL training +- **Boosting Reasoning with Auxiliary Training Signals.** Can auxiliary signals - such + as Process Reward Models (PRMs) and expert-curated labels - be adopted in RL training to refine the reasoning processes of LRMs? - **Adaptive Response Lengths: Quality Over Quantity.** How can RL training schemes teach LRMs to dynamically adjust response complexity based on task difficulty? The diff --git a/docs/_toc.yml b/docs/_toc.yml index 866f6e5e85..d83b658043 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -7,22 +7,22 @@ parts: - caption: Version History chapters: - file: version_history - - caption: Getting Started - chapters: - - file: tutorial/gsm8k_grpo - caption: Tutorial chapters: - file: tutorial/installation - file: tutorial/installation_npu - file: tutorial/quickstart - - file: tutorial/openai_workflows - file: tutorial/agentic_rl - file: tutorial/eval - file: tutorial/megatron - file: cli_reference + - caption: Code Walkthrough + chapters: + - file: tutorial/gsm8k_grpo - caption: Best Practices chapters: - file: best_practices/algo_perf + - file: best_practices/workflow - file: best_practices/debugging - file: best_practices/handling_oom - file: best_practices/perf_profiling @@ -37,3 +37,9 @@ parts: - file: algorithms/grpo_series - file: algorithms/m2po - file: algorithms/prox_approx + - caption: Reference + chapters: + - file: reference/checkpointing + - file: reference/metrics_tracking + - file: reference/alloc_mode + - file: reference/rollout_workflow diff --git a/docs/best_practices/debugging.md b/docs/best_practices/debugging.md index b42fa11d45..31916a79b8 100644 --- a/docs/best_practices/debugging.md +++ b/docs/best_practices/debugging.md @@ -14,9 +14,9 @@ logic, enabling repeated testing without server restarts. **Benefits:** -- **Lightweight** — Your debug program only requires CPU while inference runs on GPU -- **IDE-friendly** — Works seamlessly with VS Code's Python debugger and other IDEs -- **Fast iterations** — No server restarts needed between debugging sessions +- **Lightweight** - Your debug program only requires CPU while inference runs on GPU +- **IDE-friendly** - Works seamlessly with VS Code's Python debugger and other IDEs +- **Fast iterations** - No server restarts needed between debugging sessions ### 1. Launch the Standalone SGLang Server @@ -89,10 +89,10 @@ SFT) by using pre-generated data instead of running live inference. **Benefits:** -- **No inference servers** — Eliminate server management overhead -- **Faster iterations** — Skip the expensive data collection step -- **Reproducible** — Use identical data across debugging sessions -- **Isolated testing** — Focus exclusively on your RL logic +- **No inference servers** - Eliminate server management overhead +- **Faster iterations** - Skip the expensive data collection step +- **Reproducible** - Use identical data across debugging sessions +- **Isolated testing** - Focus exclusively on your RL logic ### 1. Configure Allocation Mode diff --git a/docs/best_practices/workflow.md b/docs/best_practices/workflow.md new file mode 100644 index 0000000000..444ab6c0f4 --- /dev/null +++ b/docs/best_practices/workflow.md @@ -0,0 +1,100 @@ +# Writing Agent Workflows + +This guide covers best practices for implementing efficient and robust `RolloutWorkflow` +classes and agent workflows in AReaL. + +For the difference between `RolloutWorkflow` and agent workflows, see the +[Agentic RL Guide](../tutorial/agentic_rl.md). + +## Best Practices + +### Use Async/Await Throughout + +All workflow methods should be async and use `await` for I/O-bound operations. This +enables concurrent execution across multiple rollouts. + +```python +# Or the `run` method in agent workflows +async def arun_episode(self, engine, data): + # Correct: await the engine call + resp = await engine.agenerate(req) + + # Correct: await other LLM calls + async with AsyncOpenAI() as client: + resp = await client.chat.completions.create(...) + + # Incorrect: blocking calls stall other rollouts + # resp = engine.generate(req) # Don't do this + # resp = OpenAI().chat.completions.create(...) # Don't do this + + # Await HTTP requests with reused client + session = await workflow_context.get_aiohttp_session() + async with session.get(url) as response: + result = await response.json() + + # Await file operations (use aiofiles) + async with aiofiles.open(path, "r") as f: + content = await f.read() +``` + +### Wrap Expensive Reward Functions + +Use `AsyncRewardWrapper` for reward functions involving CPU-intensive computation, +external API calls, or any blocking operation. `AsyncRewardWrapper` dispatches reward +computation to a dedicated process pool. + +```python +from areal.api.reward_api import AsyncRewardWrapper + +class MyWorkflow(RolloutWorkflow): + def __init__(self, reward_fn, ...): + # Wrap the reward function once during initialization + self.async_reward_fn = AsyncRewardWrapper(reward_fn) + + async def arun_episode(self, engine, data): + resp = await engine.agenerate(req) + # Await the wrapped reward function + reward = await self.async_reward_fn( + prompt_str, + completion_str, + **data, + ) +``` + +### Avoid Heavy Initialization + +Place expensive setup logic in `__init__`, not in `arun_episode`. The `arun_episode` +method runs for every rollout, so repeated initialization wastes resources. + +### Reuse HTTP Clients via Workflow Context + +Reuse HTTP clients across requests instead of creating new ones. AReaL provides shared +clients through `workflow_context` with automatic lifecycle management. + +When using OpenAI, Anthropic, or other SDK clients, pass the shared HTTP client: + +```python +from openai import AsyncOpenAI +from areal.infra import workflow_context + +class MyAgentWorkflow: + async def run(self, data, **extra_kwargs): + # Get pre-configured client from extra_kwargs + http_client = extra_kwargs.get("http_client") + base_url = extra_kwargs.get("base_url") + + # Pass to SDK constructor + client = AsyncOpenAI( + base_url=base_url, + http_client=http_client, + max_retries=0, + ) + + response = await client.chat.completions.create(...) +``` + +## See Also + +- [Debugging Guide](../best_practices/debugging.md) - Debugging customized workflows +- [RolloutWorkflow Reference](../reference/rollout_workflow.md) - API documentation +- [Agentic RL Guide](../tutorial/agentic_rl.md) - Training with agent frameworks diff --git a/docs/cli_reference.md b/docs/cli_reference.md index 1dbe6997bb..450cdec7f7 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -691,16 +691,16 @@ Configuration for model evaluation scheduling and timing. Configuration for experiment recovery and fault tolerance. -| Parameter | Type | Default | Description | -| ----------------- | --------------- | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `experiment_name` | string | **Required** | - | -| `trial_name` | string | **Required** | - | -| `fileroot` | string | **Required** | - | -| `freq_epochs` | integer \| None | `None` | Trigger frequency in epochs. None disables epoch-based saving. | -| `freq_steps` | integer \| None | `None` | Trigger frequency in steps. None disables step-based saving. | -| `freq_secs` | integer \| None | `None` | Trigger frequency in seconds. None disables time-based saving. | -| `mode` | string | `"disabled"` | Recovery mode for the launcher. Options: 'disabled': Never recover from previous runs. 'auto': Automatically recover from previous runs if recover info and checkpoints are available. 'fault': Only recover from previous runs if the new run fails. 'resume': Force to resume, raise an error if no recover info was found. Never resume if failed again. | -| `retries` | integer | `3` | Number of recovery retries (auto/fault modes only). | +| Parameter | Type | Default | Description | +| ----------------- | --------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `experiment_name` | string | **Required** | - | +| `trial_name` | string | **Required** | - | +| `fileroot` | string | **Required** | - | +| `freq_epochs` | integer \| None | `None` | Trigger frequency in epochs. None disables epoch-based saving. | +| `freq_steps` | integer \| None | `None` | Trigger frequency in steps. None disables step-based saving. | +| `freq_secs` | integer \| None | `None` | Trigger frequency in seconds. None disables time-based saving. | +| `mode` | string | `"disabled"` | Recovery mode for the launcher. Options: 'on' or 'auto': Automatically recover from previous runs if recover info and checkpoints are available. 'off' or 'disabled': Never recover from previous runs. | +| `retries` | integer | `3` | Number of recovery retries when recovery is enabled. | (section-saver)= diff --git a/docs/reference/alloc_mode.md b/docs/reference/alloc_mode.md new file mode 100644 index 0000000000..bb97946760 --- /dev/null +++ b/docs/reference/alloc_mode.md @@ -0,0 +1,142 @@ +# Allocation Mode + +This document describes AReaL's allocation mode system, which controls how GPUs are +distributed between inference and training backends during distributed RL training. + +## Overview + +The `allocation_mode` configuration option is a pattern-based string that specifies: + +- Which backends to use for inference (SGLang, vLLM) and training (FSDP, Megatron, + Archon) +- The parallelization strategy for each backend +- The total number of GPUs required + +AReaL parses this string into an `AllocationMode` object that orchestrates resource +allocation across the cluster. + +## Syntax + +### Basic Format + +``` +: +``` + +### Two-Component Format (Inference + Training) + +``` +: + : +``` + +The `+` operator separates components that run on **separate GPU pools**. + +### Parallelism Dimensions + +| Dimension | Abbreviation | Description | Valid For | +| --------- | ------------ | ---------------------------------- | ---------------- | +| Data | `d` | Number of model replicas | All backends | +| Tensor | `t` | Split operations across GPUs | All backends | +| Pipeline | `p` | Split layers across GPUs in stages | Megatron, Archon | +| Context | `c` | Split sequence length across GPUs | All backends | +| Expert | `e` | Split MoE experts across GPUs | Megatron, Archon | + +Dimensions are specified as ``, e.g., `d4t2` means data parallel size 4 +and tensor parallel size 2. + +## Calculating GPU Requirements + +The total GPUs for a component is computed as: + +``` +world_size = dp × tp × pp × cp +``` + +Expert parallelism (`e`) does not increase world size—it redistributes how experts are +placed within the existing GPU mesh. + +### Examples + +| Allocation Mode | Inference GPUs | Training GPUs | Total | +| --------------------------------- | -------------- | ------------- | ----- | +| `d8` | - | 8 | 8 | +| `sglang:d2t4` | 8 | - | 8 | +| `sglang:d2t4 + fsdp:d4t2` | 8 | 8 | 16 | +| `sglang:d4t4 + megatron:d2p2t4e4` | 16 | 16 | 32 | + +## Backend Selection + +### Inference Backends + +| Backend | Supported Dimensions | +| -------- | -------------------- | +| `sglang` | `d`, `t` | +| `vllm` | `d`, `t`, `p` | + +For inference, `d` represents the number of independent server instances, and each +instance uses `t × p` GPUs. + +Note that the internal backed configurations do not affect how AReaL allocate GPUs. +Given allocation mode `sglang:d4t4`, you can also config `sglang.dp_size=4`, +`sglang.ep_size=4`, and `sglang.enable_dp_attention=True`. In this case, we launch 4 +model replicas each with 4 GPUs. Within each instance, SGLang will still use DP +attention and expert parallelism to distribute computations in attention and expert +layers. + +### Training Backends + +| Backend | Supported Dimensions | Use Case | +| ---------- | ----------------------- | ---------------------------------------- | +| `fsdp` | `d`, `t`, `c` | Default for simple parallelism | +| `megatron` | `d`, `t`, `p`, `c`, `e` | Required for pipeline or expert parallel | +| `archon` | `d`, `t`, `p`, `c`, `e` | Alternative to Megatron (experimental) | + +When the backend is omitted, AReaL auto-selects based on the parallelism configuration: + +- **FSDP**: Used when only `d`, `t`, `c` are specified +- **Megatron**: Used when `p > 1` or `e > 1` + +``` +# Equivalent forms +d4t2 # Auto-selects FSDP +fsdp:d4t2 # Explicit FSDP + +d2p2t4 # Auto-selects Megatron (pp > 1) +megatron:d2p2t4 # Explicit Megatron +``` + +## MoE Hybrid Parallelism + +For Mixture-of-Experts models, Megatron/Archon supports different parallelism strategies +for attention and FFN (expert) modules using the hybrid syntax: + +``` +megatron:(attn:|ffn:) +``` + +This enables +[MoE Parallel Folding](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/moe#moe-parallel-folding), +which reduces the minimum GPU requirement for combined context and expert parallelism. + +### Constraints + +- Pipeline parallel size (`p`) must be identical for `attn` and `ffn` +- World size must match (if `d` is omitted in `ffn`, it is derived automatically) +- Expert parallel (`e`) is only valid in the `ffn` section + +### Example + +``` +megatron:(attn:d4p2t2c2|ffn:d2p2t4e2) +``` + +| Module | dp | pp | tp | cp | ep | World Size | +| ------ | --- | --- | --- | --- | --- | ---------- | +| attn | 4 | 2 | 2 | 2 | - | 32 | +| ffn | 2 | 2 | 4 | - | 2 | 32 | + +## See Also + +- [Fine-tuning Large MoE Models](../tutorial/megatron.md) - Tutorial for Megatron + backend +- [Megatron Performance Best Practice](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/moe#performance-best-practice) diff --git a/docs/reference/checkpointing.md b/docs/reference/checkpointing.md new file mode 100644 index 0000000000..264a3a877d --- /dev/null +++ b/docs/reference/checkpointing.md @@ -0,0 +1,225 @@ +# Checkpointing + +This document describes AReaL's checkpointing system, which handles model saving for +evaluation and fault-tolerant recovery during distributed RL training. + +## Overview + +AReaL provides two complementary checkpointing mechanisms: + +| Mechanism | Purpose | Format | Includes Optimizer/DataLoader State | +| ------------------ | ------------------------------------------ | ---------------------------- | ----------------------------------- | +| **Saver** | Export models for evaluation or publishing | HuggingFace | No | +| **RecoverHandler** | Resume training after failures | DCP (Distributed Checkpoint) | Yes | + +Both mechanisms are invoked automatically during training and can be configured via +`config.saver` and `config.recover` respectively. + +## Checkpoint Formats + +### HuggingFace Format + +Used by `Saver` for model export: + +- Standard HuggingFace model format (safetensors + config.json) +- Compatible with `transformers.AutoModel.from_pretrained()` +- Can be uploaded to HuggingFace Hub +- Does not include optimizer state + +### DCP Format (Distributed Checkpoint) + +Used by `RecoverHandler` for fault tolerance: + +- Backend's native distributed checkpoint format (`torch.distributed.checkpoint` or + Megatron distributed checkpoint) +- Sharded across all ranks for efficient parallel I/O +- Includes model weights, optimizer state, RNG state, etc +- Backend-specific: checkpoints are only compatible with the same parallelism + configuration +- Overwrites previous checkpoint to save disk space + +## Architecture + +``` +PPOTrainer.train() +│ +├── Training loop +│ ├── Rollout, compute values, PPO update... +│ │ +│ ├── _save_hf() # HuggingFace export +│ │ └── Saver.save() +│ │ └── engine.save(weight_format="hf") +│ │ +│ └── _save_recover_checkpoint() # Fault tolerance +│ └── RecoverHandler.dump() +│ └── engine.save(weight_format="dcp", with_optim=True) +│ +└── On restart + └── RecoverHandler.load() + ├── Restore dataloader, saver, evaluator states + └── engine.load(weight_format="dcp", with_optim=True) +``` + +## Saver: HuggingFace Model Export + +The [`Saver`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/saver.py) +periodically exports model weights in HuggingFace format for evaluation or deployment. + +### Configuration + +Configure via `config.saver`: + +| Parameter | Type | Default | Description | +| ------------- | ----------- | ------- | ------------------------------------ | +| `freq_epochs` | int \| None | None | Save every N epochs. None disables. | +| `freq_steps` | int \| None | None | Save every N steps. None disables. | +| `freq_secs` | int \| None | None | Save every N seconds. None disables. | + +Example configuration: + +```yaml +saver: + freq_epochs: 1 # Save at end of each epoch + freq_steps: null # Disabled + freq_secs: null # Disabled +``` + +Saving is triggered when any of epoch/step/time condition is met. + +### Output Location + +Checkpoints are saved to: + +``` +{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/default/ +└── epoch{E}epochstep{S}globalstep{G}/ + ├── config.json + ├── model.safetensors (or model-00001-of-00002.safetensors, etc.) + ├── tokenizer.json + └── ... +``` + +### Usage + +Load saved checkpoints with standard HuggingFace APIs: + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model = AutoModelForCausalLM.from_pretrained( + "/path/to/checkpoint/epoch0epochstep99globalstep99" +) +tokenizer = AutoTokenizer.from_pretrained( + "/path/to/checkpoint/epoch0epochstep99globalstep99" +) +``` + +## RecoverHandler: Fault Tolerance + +The +[`RecoverHandler`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/recover.py) +enables resuming training after failures by saving complete training state. + +### Configuration + +Configure via `config.recover`: + +| Parameter | Type | Default | Description | +| ------------- | ----------- | ---------- | ------------------------------------------------ | +| `mode` | str | "disabled" | Recovery mode: "on"/"auto" or "off"/"disabled" | +| `freq_epochs` | int \| None | None | Checkpoint every N epochs | +| `freq_steps` | int \| None | None | Checkpoint every N steps | +| `freq_secs` | int \| None | None | Checkpoint every N seconds | +| `retries` | int | 3 | Number of recovery retries when recovery enabled | + +#### Recovery Modes + +| Mode | Behavior | +| ------------------- | ----------------------------------------------- | +| `on` or `auto` | Automatically resume if valid checkpoint exists | +| `off` or `disabled` | No checkpointing or recovery | + +When recovery is enabled (`on`/`auto`), the system will: + +1. Periodically save recovery checkpoints (model weights, optimizer state, dataloader + position) +1. Automatically resume from the last valid checkpoint on restart +1. Retry up to `retries` times on failure + +Example configuration: + +```yaml +recover: + mode: on # or "auto" for backward compatibility + freq_steps: 100 # Checkpoint every 100 steps + retries: 3 +``` + +### What Gets Saved + +RecoverHandler saves complete training state: + +| Component | Contents | +| ----------------- | -------------------------------------------------- | +| Model weights | DCP format, sharded across ranks | +| Optimizer state | Momentum, variance (Adam), learning rate scheduler | +| RNG state | Python, NumPy, PyTorch, CUDA random states | +| Dataloader state | Current position in dataset | +| Training progress | Epoch, step, global_step counters | +| Auxiliary states | Saver, Evaluator, StatsLogger states | + +### Output Location + +Recovery checkpoints are saved to: + +``` +{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/ +├── default/ +│ └── recover_checkpoint/ # Model + optimizer (DCP format) +│ ├── __0_0.distcp +│ ├── __1_0.distcp +│ └── ... +├── critic/ # If using critic +│ └── recover_checkpoint/ +└── recover_info/ # Metadata + ├── step_info.json + ├── saver_info.json + ├── evaluator_info.json + ├── stats_logger_info.json + ├── checkpoint_info.json + └── dataloader_info.pkl +``` + +### Recovery Process + +When training resumes: + +1. `RecoverHandler.load()` restores all saved state (if any) +1. Training continues from `last_step_info.next().global_step` +1. Inference engine weights are synchronized to match recovered state + +## Best Practices + +### Frequency Guidelines + +| Scenario | Recommended Setting | +| ------------------ | ----------------------------------------- | +| Long training runs | `freq_epochs: 1` or `freq_steps: 1000` | +| Unpredictable time | `freq_secs: 7200` | +| Unstable clusters | `freq_steps: 100` with `recover.mode: on` | +| Limited disk space | Lower frequency, rely on final checkpoint | +| Debugging | `freq_steps: 1` for quick iteration | + +### Disk Space Considerations + +- **Saver**: Each save creates a new directory. High frequency consumes significant + space. +- **RecoverHandler**: Overwrites previous checkpoint. Only one copy exists at a time. + +### Recovery Tips + +1. **Verify checkpoint validity**: Check `recover_info/step_info.json` for the last + saved step +1. **Same config required**: DCP checkpoints require identical parallelism + configuration, experiment name, and trial name +1. **Clean restart**: Delete `recover_info/` directory to start fresh diff --git a/docs/reference/metrics_tracking.md b/docs/reference/metrics_tracking.md new file mode 100644 index 0000000000..87c75b7460 --- /dev/null +++ b/docs/reference/metrics_tracking.md @@ -0,0 +1,326 @@ +# Metrics Tracking + +AReaL provides a unified metrics tracking system that handles statistics collection +across distributed training and rollout workers. The system supports two distinct +paradigms optimized for their respective use cases: **streaming metrics** for +asynchronous rollout workflows and **batch metrics** for synchronous training updates. + +## Core Components + +The metrics system is built around `areal.utils.stats_tracker`, which provides: + +- **Named trackers**: Isolated metric namespaces for different components +- **Hierarchical scoping**: Organize metrics into logical groups +- **Distributed aggregation**: Automatic reduction across workers +- **Multiple reduce types**: Support for averages, sums, min/max, and scalars + +```python +from areal.utils import stats_tracker + +# Default tracker (training metrics) +stats_tracker.scalar(learning_rate=0.001) + +# Named tracker (rollout metrics) +stats_tracker.get("rollout").scalar(reward=0.5) +``` + +## Two Logging Paradigms + +### Streaming Metrics (Rollout Workers) + +Rollout workers execute workflows asynchronously, with each workflow logging metrics +independently. This streaming approach handles variable completion times naturally. + +**Characteristics:** + +- Each workflow logs scalars individually as they complete +- Metrics accumulate in a list within the worker process +- No synchronization between workers during logging +- Aggregation happens at export time via the controller + +**Example from `RLVRWorkflow`:** + +```python +# areal/workflow/rlvr.py +async def _collect_samples(self, engine, req, prompt_str, task_data): + resp = await engine.agenerate(req) + reward = await self._compute_rewards(resp, prompt_str, task_data) + + # Log single scalar - appends to internal list + # `workflow_context.stat_scope()` automatically differentiates evaluation/training scopes + stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) + + return resp, reward +``` + +You can log any other scalars in your customized workflow, e.g., + +```python +async def run(self, data, **extra_kwargs): + # `workflow_context.stat_scope()` automatically differentiates evaluation/training scopes + stats_tracker.get(workflow_context.stat_scope()).scalar(num_turns=num_turns, max_tokens=max_tokens, reward=reward) + return reward +``` + +**Controller aggregation:** + +The `RolloutController` collects stats from all workers and computes weighted averages: + +```python +# areal/infra/controller/rollout_controller.py +def export_stats(self) -> dict[str, float]: + all_raw_stats = self._collective_rpc(method="export_stats") + + # Aggregate using counts as weights + stats, counts = defaultdict(float), defaultdict(int) + for raw_stats in all_raw_stats: + for k, v in raw_stats.items(): + if k.endswith("__count"): + counts[k] += v + else: + stats[k] += v * raw_stats.get(k + "__count", 0) + + # Compute weighted averages + return {k: v / counts[k + "__count"] for k, v in stats.items() + if counts.get(k + "__count", 0) > 0} +``` + +### Batch Metrics (Training Engines) + +Training engines process data in synchronized batches across data-parallel ranks. +Metrics are logged as tensors with boolean masks, then reduced across all ranks at +export time. + +**Characteristics:** + +- Log entire batch tensors with denominator masks +- Support for per-token and per-sequence statistics +- All-reduce synchronization ensures consistent stats across ranks +- Multiple reduce types: `AVG_MIN_MAX`, `AVG`, `SUM`, `MIN`, `MAX` + +**Example from `PPOActor`:** + +```python +# areal/engine/ppo/actor.py +def ppo_update(self, data): + loss_mask = data["loss_mask"].bool() + reward_score = data["rewards"] + + # Define denominators (boolean masks) + stats_tracker.denominator( + n_seqs=torch.ones_like(reward_score, dtype=torch.bool), + n_valid_tokens=loss_mask, + ) + + # Log tensor metrics with denominator reference + stats_tracker.stat( + advantages=data["advantages"], # [batch, seq_len] + kl_rewards=data["kl_rewards"], # [batch, seq_len] + denominator="n_valid_tokens" + ) + + stats_tracker.stat( + task_reward=reward_score.float(), # [batch] + seq_len=seqlens.float(), # [batch] + denominator="n_seqs" + ) +``` + +**Export behavior:** + +```python +# areal/engine/fsdp_engine.py +def export_stats(self) -> dict[str, float]: + # All-reduce across data-parallel group + return stats_tracker.export_all(reduce_group=self.data_parallel_group) + # All DP ranks receive identical results +``` + +## API Reference + +### Recording Methods + +| Method | Use Case | Example | +| ----------------------------- | --------------------------- | ---------------------------------------- | +| `scalar(**kwargs)` | Single float values | `scalar(lr=0.001, eps=0.2)` | +| `denominator(**kwargs)` | Define boolean masks | `denominator(valid=mask.bool())` | +| `stat(denominator, **kwargs)` | Tensor metrics with masking | `stat(loss=tensor, denominator="valid")` | + +### Reduce Types + +When using `stat()`, metrics default to `AVG_MIN_MAX`, which creates three output keys: + +```python +stats_tracker.stat(loss=tensor, denominator="valid") +# Exports: {"loss/avg": 0.5, "loss/min": 0.1, "loss/max": 0.9} +``` + +Available reduce types: + +| Type | Output | Description | +| ------------- | ------------------------------- | ------------------------ | +| `AVG_MIN_MAX` | `key/avg`, `key/min`, `key/max` | Default for tensor stats | +| `AVG` | `key` | Weighted average only | +| `SUM` | `key` | Sum across all elements | +| `MIN` | `key` | Minimum value | +| `MAX` | `key` | Maximum value | +| `SCALAR` | `key`, `key__count` | For scalar values | + +### Scoping + +Organize related metrics using hierarchical scopes: + +```python +with stats_tracker.scope("ppo_actor"): + with stats_tracker.scope("update"): + stats_tracker.stat(loss=loss_tensor, denominator="valid") + # Key: "ppo_actor/update/loss/avg" +``` + +### Timing + +Measure execution time with automatic scoping under `timeperf/`: + +```python +with stats_tracker.record_timing("rollout"): + batch = actor.prepare_batch(dataloader, workflow) +# Key: "timeperf/rollout" +``` + +### Named Trackers + +Isolate metrics for different components: + +```python +# Training metrics (default tracker) +stats_tracker.scalar(grad_norm=1.5) + +# Rollout metrics +stats_tracker.get("rollout").scalar(reward=0.8) + +# Evaluation metrics +stats_tracker.get("eval-rollout").scalar(reward=0.9) + +# Export from all trackers +all_stats = stats_tracker.export_all(reduce_group=group) +``` + +## Data Flow + +The complete metrics flow from collection to logging: + +``` +Rollout Workers Training Workers +─────────────── ──────────────── +workflow.arun_episode() actor.ppo_update(batch) + │ │ + ▼ ▼ +get("rollout").scalar(r=0.5) stat(tensor, denom=mask) + │ │ + ▼ ▼ +export_stats(reduce_group=None) export_stats(reduce_group=dp_group) +{reward: 0.5, reward__count: 1} → all_reduce across DP ranks + │ │ + ▼ │ +RolloutController.export_stats() │ +→ weighted avg across workers │ + │ │ + └────────────────┬───────────────────────┘ + ▼ + PPOTrainer._export_and_commit_stats() + │ + ▼ + StatsLogger.commit(stats) + │ + ┌────────────┼────────────┐ + ▼ ▼ ▼ + wandb tensorboard swanlab +``` + +## StatsLogger: Logging Backends + +The +[`StatsLogger`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/stats_logger.py) +sends aggregated metrics to external logging backends. It is automatically managed by +`PPOTrainer` and runs only on rank 0 to avoid duplicate logging. + +### Supported Backends + +| Backend | Configuration | Description | +| -------------------- | --------------------------------- | ------------------------------- | +| **Weights & Biases** | `config.stats_logger.wandb` | Cloud-based experiment tracking | +| **SwanLab** | `config.stats_logger.swanlab` | Alternative experiment tracking | +| **TensorBoard** | `config.stats_logger.tensorboard` | Local visualization | + +### Integration with PPOTrainer + +The trainer calls `StatsLogger.commit()` at the end of each training step: + +```python +# areal/experimental/trainer/rl.py +def _export_and_commit_stats(self, epoch, epoch_step, global_step): + # 1. Collect metrics from all components + stats = self.actor.export_stats() # Training metrics (all-reduced) + stats.update(self.rollout.export_stats()) # Rollout metrics (controller-aggregated) + stats.update(self.eval_rollout.export_stats()) # Eval metrics + + # 2. Send to logging backends (rank 0 only) + self.stats_logger.commit(epoch, epoch_step, global_step, stats) +``` + +### StatsLogger.commit() + +The `commit()` method filters out internal count keys and logs to all configured +backends: + +```python +# areal/utils/stats_logger.py +def commit(self, epoch, step, global_step, data): + if dist.is_initialized() and dist.get_rank() != 0: + return # Only rank 0 logs + + # Filter out __count keys (used internally for weighted averaging) + data = {k: v for k, v in data.items() if not k.endswith("__count")} + + # Log to all backends + wandb.log(data, step=global_step) + swanlab.log(data, step=global_step) + if self.summary_writer: + for key, val in data.items(): + self.summary_writer.add_scalar(key, val, global_step) +``` + +### Configuration + +Configure logging backends in your experiment config: + +```yaml +stats_logger: + experiment_name: "gsm8k_grpo" + trial_name: "run_001" + fileroot: "/path/to/logs" + + wandb: + mode: "online" # "online", "offline", or "disabled" + project: "my-project" + entity: "my-team" + + swanlab: + mode: "online" # "online", "local", or "disabled" + project: "my-project" + + tensorboard: + path: "/path/to/tensorboard/logs" # null to disable +``` + +## Best Practices + +1. **Choose the right paradigm**: Use `scalar()` for scalars, `stat()` with denominators + for batched pytorch tensors (usually training metrics). + +1. **Define denominators first**: Always call `denominator()` before `stat()` to + establish the masking relationship. + +1. **Use named trackers**: Use + `stats_tracker.get(workflow_context.stat_scope()).scalar(...)` to isolate rollout + (`"rollout"`) and evaluation (`"eval-rollout"`) metrics from training metrics. diff --git a/docs/reference/rollout_workflow.md b/docs/reference/rollout_workflow.md new file mode 100644 index 0000000000..5adb2a0109 --- /dev/null +++ b/docs/reference/rollout_workflow.md @@ -0,0 +1,212 @@ +# RolloutWorkflow Reference + +This document describes the `RolloutWorkflow` abstraction, the core interface for +implementing rollout generation in AReaL's reinforcement learning pipeline. + +**Note**: This page targets developers seeking a deep understanding of the codebase. For +agentic RL training, use the high-level API described in the +[Agentic RL Guide](../tutorial/agentic_rl.md). + +## Overview + +A `RolloutWorkflow` defines how to generate training trajectories from input data. It +encapsulates the logic for: + +- Tokenizing prompts and preparing model inputs +- Calling the inference engine to generate completions +- Computing rewards for generated outputs +- Packaging results into tensor dictionaries for training + +## Interface + +```python +from areal.api.workflow_api import RolloutWorkflow + +class RolloutWorkflow(ABC): + @abstractmethod + async def arun_episode( + self, engine: InferenceEngine, data: dict[str, Any] + ) -> dict[str, Any] | None | dict[str, InteractionWithTokenLogpReward]: + """Run a single episode of the workflow.""" + ... +``` + +### Parameters + +| Parameter | Type | Description | +| --------- | ----------------- | ----------------------------------------------- | +| `engine` | `InferenceEngine` | Inference engine for generating model responses | +| `data` | `dict[str, Any]` | A single sample from the dataloader | + +### Return Types + +The `arun_episode` method supports three return types: + +| Return Type | Description | +| ------------------------------------------- | -------------------------------------------------------------------------------------------------- | +| `dict[str, torch.Tensor]` | Standard tensor format for training | +| `dict[str, InteractionWithTokenLogpReward]` | Token-level interactions (auto-converted to tensors); produced by the high-level `ArealOpenAI` API | +| `None` | Rejected trajectory, excluded from training | + +## Tensor Dictionary Format + +When returning a tensor dictionary, the following fields are expected: + +| Field | Shape | Type | Required | Description | +| ---------------- | ----------------------- | ------- | -------- | ----------------------------------- | +| `input_ids` | `[batch_size, seq_len]` | int32 | Yes | Token IDs (prompt + completion) | +| `attention_mask` | `[batch_size, seq_len]` | bool | Yes | Valid token mask | +| `loss_mask` | `[batch_size, seq_len]` | int32 | No | Completion token mask (1 = train) | +| `logprobs` | `[batch_size, seq_len]` | float32 | No | Log probabilities per token | +| `rewards` | `[batch_size]` | float32 | No | Per-sequence rewards | +| `versions` | `[batch_size, seq_len]` | int32 | No | Weight version when token generated | + +Example return value: + +```python +return { + "input_ids": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32), + "attention_mask": torch.ones(1, 5, dtype=torch.bool), + "loss_mask": torch.tensor([[0, 0, 1, 1, 1]], dtype=torch.int32), + "logprobs": torch.tensor([[0.0, 0.0, -0.5, -0.3, -0.2]], dtype=torch.float32), + "rewards": torch.tensor([1.0], dtype=torch.float32), + "versions": torch.tensor([[0, 0, 1, 1, 1]], dtype=torch.int32), +} +``` + +## Workflow Context + +Inside `arun_episode`, access the execution context via the `workflow_context` module. +Each workflow instance has its own isolated context: + +```python +from areal.infra import workflow_context + +async def arun_episode(self, engine, data): + # Get current execution context + ctx = workflow_context.get() + + # Check if running in evaluation mode + if ctx.is_eval: + # Use different parameters for evaluation + ... + + # Get task ID for logging + task_id = ctx.task_id + + # Get stats scope based on mode ("rollout" or "eval-rollout") + scope = workflow_context.stat_scope() +``` + +## Trajectory Dumping + +When `InferenceEngineConfig.dump_to_file=True`, trajectories are automatically saved to +disk for debugging and analysis. + +### Configuration + +```yaml +rollout: + dump_to_file: true + fileroot: "/path/to/logs" + tokenizer_path: "model/tokenizer" # Required for text decoding +``` + +### Output Location + +Trajectories are saved to: + +``` +{fileroot}/{experiment_name}/{trial_name}/[rollout|eval-rollout]/{version}/{task_id}.jsonl +``` + +Example: + +``` +/tmp/areal/my_exp/trial1/rollout/5/42.jsonl +``` + +### Output Format + +Each line in the JSONL file contains: + +```json +{ + "task_id": 42, + "sample_idx": 0, + "seqlen": 256, + "prompt_len": 128, + "head_version": 5, + "tail_version": 5, + "reward": 1.0, + "prompt": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n", + "completion": "The answer is 4.<|im_end|>" +} +``` + +## Implementing Custom Workflows + +To create a custom workflow: + +1. **Subclass `RolloutWorkflow`**: + +```python +from areal.api.workflow_api import RolloutWorkflow + +class MyWorkflow(RolloutWorkflow): + def __init__(self, tokenizer, gconfig, **kwargs): + self.tokenizer = tokenizer + self.gconfig = gconfig + + async def arun_episode(self, engine, data): + # 1. Prepare input + input_ids = self.tokenizer.encode(data["prompt"]) + + # 2. Generate completion + req = ModelRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig, + tokenizer=self.tokenizer, + ) + resp = await engine.agenerate(req) + + # 3. Compute reward + reward = self.compute_reward(resp, data) + + # 4. Return tensor dict (or None to reject) + if reward < 0: + return None + + return self.build_tensor_dict(resp, reward) +``` + +2. **Register with trainer**: + +```python +trainer.train( + workflow=MyWorkflow, + workflow_kwargs={ + "tokenizer": tokenizer, + "gconfig": config.gconfig, + }, +) +``` + +## Workflow Resolution + +Workflows can be specified in multiple ways: + +| Format | Example | Description | +| -------------- | -------------------------------- | -------------------------- | +| Instance | `MyWorkflow(...)` | Pre-instantiated workflow | +| Class | `MyWorkflow` | Class (requires kwargs) | +| String path | `"my_module.MyWorkflow"` | Dynamic import | +| Agent workflow | Any class with `async def run()` | Wrapped with proxy support | + +The training system automatically resolves these to `RolloutWorkflow` instances. + +## See Also + +- [Agentic RL Tutorial](../tutorial/agentic_rl.md) - Training with agent frameworks +- [Adding Custom Workflows](../customization/agent.md) - Step-by-step guide diff --git a/docs/tutorial/agentic_rl.md b/docs/tutorial/agentic_rl.md index aa381f331e..2bca0043f3 100644 --- a/docs/tutorial/agentic_rl.md +++ b/docs/tutorial/agentic_rl.md @@ -49,7 +49,7 @@ AReaL addresses these limitations by providing: diversity and improves policy gradient estimates. We demonstrate several concrete examples below. More examples can be found in the -[`workflow/` directory](../../areal/workflow/). +[`workflow/` directory](https://github.com/inclusionAI/AReaL/tree/main/areal/workflow). ## Examples diff --git a/docs/tutorial/eval.md b/docs/tutorial/eval.md index 292eeebb12..af9424e40d 100644 --- a/docs/tutorial/eval.md +++ b/docs/tutorial/eval.md @@ -1,8 +1,13 @@ # Evaluation -AReaL provides distributed inference evaluation using the same RolloutController -infrastructure as training. This allows you to leverage existing workflows and -schedulers to scale evaluation across multiple GPUs and nodes. +AReaL supports distributed inference using the same controller infrastructure as +training. This allows you to leverage existing workflows and schedulers to scale +evaluation across multiple GPUs and nodes. + +**Note:** AReaL provides distributed inference for your trained model, not a complete +evaluation pipeline with dataset retrieval and metrics computation. You can use +third-party evaluation frameworks with AReaL checkpoints directly --- no conversion +required since AReaL saves HuggingFace-compatible checkpoints. ## Quick Start @@ -33,6 +38,61 @@ python3 examples/math/gsm8k_eval.py \ cluster.n_nodes=12 ``` +## Evaluation Metrics + +Select an appropriate dataset and metrics for your task, then integrate the evaluation +logic as a workflow. See the [Agentic RL guide](./agentic_rl.md) for details. + +Example with an agentic math evaluator (the evaluation code is independent with AReaL): + +```python +from agents import Agent, OpenAIProvider, RunConfig, SQLiteSession, function_tool +from agents import Runner as OpenAIRunner +from math_verify import parse, verify +from openai import AsyncOpenAI + + +@function_tool +def add(a: float, b: float) -> float: + """Add two numbers.""" + return a + b + + +@function_tool +def multiply(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b + + +def math_reward_fn(completions: str, answer: str) -> float: + return float(verify(parse(completions), parse(answer))) + + +class MathAgent: + async def run(self, data, **extra_kwargs): + http_client = extra_kwargs.get("http_client") + base_url = extra_kwargs.get("base_url") + client = AsyncOpenAI(base_url=base_url, http_client=http_client, max_retries=0) + + run_config = RunConfig( + model_provider=OpenAIProvider(openai_client=client), + model="default", + tracing_disabled=True, + ) + agent = Agent( + name="RLVR Math with Calculator", + instructions="Answer math questions using the calculator tools.", + tools=[add, multiply], + ) + result = await OpenAIRunner.run( + agent, + input=data["messages"][-1]["content"], + session=SQLiteSession("math"), + run_config=run_config, + ) + return math_reward_fn(result.final_output, data["answer"]) +``` + ## Architecture Evaluation uses a single-controller architecture without training workers: @@ -40,9 +100,9 @@ Evaluation uses a single-controller architecture without training workers: ``` Controller Process │ - └─> RolloutController - ├─> Scheduler creates inference workers (SGLang/vLLM) - ├─> BatchTaskDispatcher submits eval tasks + └─> Inference Engine Controller (SGLang/vLLM) + ├─> Scheduler creates inference workers + ├─> Submits evaluation tasks with workflow └─> Collects results and computes metrics ``` @@ -55,18 +115,41 @@ See [`examples/math/gsm8k_eval.py`](../../examples/math/gsm8k_eval.py) for a com example. The key pattern: ```python -# Parse allocation mode and initialize scheduler +from areal.api.alloc_mode import AllocationMode +from areal.api.cli_args import GRPOConfig, SGLangConfig, load_expr_config, vLLMConfig +from areal.engine.sglang_remote import RemoteSGLangEngine +from areal.engine.vllm_remote import RemotevLLMEngine +from areal.scheduler import LocalScheduler, RayScheduler, SlurmScheduler + +# Load config and parse allocation mode +config, _ = load_expr_config(args, GRPOConfig) allocation_mode = AllocationMode.from_str(config.allocation_mode) -scheduler = LocalScheduler(exp_config=config) # or Ray/Slurm -# Create RolloutController +# Initialize scheduler based on config +if config.scheduler.type == "local": + scheduler = LocalScheduler(exp_config=config) +elif config.scheduler.type == "ray": + scheduler = RayScheduler(exp_config=config) +elif config.scheduler.type == "slurm": + scheduler = SlurmScheduler(exp_config=config) + +# Select inference engine and build server args if allocation_mode.gen_backend == "sglang": engine_cls = RemoteSGLangEngine - server_args = SGLangConfig.build_args(...) + server_args = SGLangConfig.build_args( + sglang_config=config.sglang, + tp_size=allocation_mode.gen.tp_size, + base_gpu_id=0, + ) elif allocation_mode.gen_backend == "vllm": engine_cls = RemotevLLMEngine - server_args = vLLMConfig.build_args(...) + server_args = vLLMConfig.build_args( + vllm_config=config.vllm, + tp_size=allocation_mode.gen.tp_size, + pp_size=allocation_mode.gen.pp_size, + ) +# Create controller and initialize eval_rollout = engine_cls.as_controller(config.rollout, scheduler) eval_rollout.initialize( role="eval-rollout", @@ -74,40 +157,52 @@ eval_rollout.initialize( server_args=server_args, ) +# Define workflow and its configuration +workflow = "areal.workflow.rlvr.RLVRWorkflow" +workflow_kwargs = dict( + reward_fn="areal.reward.gsm8k.gsm8k_reward_fn", + gconfig=config.gconfig, + tokenizer=config.tokenizer_path, + enable_thinking=False, +) + # Submit evaluation tasks -for data in dataloader: +cnt = 0 +for data in valid_dataloader: for item in data: - eval_rollout.submit(item, workflow, group_size=config.gconfig.n_samples) + eval_rollout.submit( + item, + workflow=workflow, + workflow_kwargs=workflow_kwargs, + group_size=config.gconfig.n_samples, + ) cnt += 1 -# Wait and collect results +# Wait for completion and collect results eval_rollout.wait(cnt, timeout=None) eval_stats = eval_rollout.export_stats() ``` -This follows the same controller pattern as training (see -[`areal/experimental/trainer/rl.py:540-594`](../../areal/experimental/trainer/rl.py)) -but without training components. +This follows the same controller pattern as training but without training components. ## Configuration -Evaluation configs use only inference components. - -It is also valid to use the previous training config but with the evaluation script. +Evaluation reuses the same config structure as training. You can use an existing +training config directly with the evaluation script. ```yaml experiment_name: gsm8k-eval trial_name: eval0 seed: 1 -allocation_mode: sglang:d4p1t1 # Inference-only +allocation_mode: sglang:d4p1t1 # Inference-only allocation scheduler: type: local # or 'ray', 'slurm' rollout: max_concurrent_rollouts: 256 - max_head_offpolicyness: 1e12 # No staleness control + # max_head_offpolicyness is set to 1e12 internally for eval gconfig: n_samples: 8 @@ -129,33 +224,26 @@ valid_dataset: batch_size: 32 ``` -## Logging Metrics +## Logging Results -Export stats and log to both console and wandb: +Use `tabulate_stats` to format evaluation metrics: ```python -eval_stats = eval_rollout.export_stats() - -if rank == 0: - print(tabulate_stats(eval_stats)) +from areal.utils.printing import tabulate_stats - try: - import wandb - if wandb.run is not None: - wandb.log({"eval": eval_stats}) - except ImportError: - pass +eval_stats = eval_rollout.export_stats() +logger.info(f"Evaluation Results: {tabulate_stats(eval_stats)}") ``` ## Custom Workflows -Reuse training workflows or create custom ones. See -[agentic RL tutorial](../tutorial/agentic_rl.md) and +Reuse training workflows or create custom ones. See the +[Agentic RL tutorial](../tutorial/agentic_rl.md) and [Customization: Rollout Workflows](../customization/agent.md) for complete guides. ## Next Steps -- [Distributed Experiments](quickstart.md#distributed-experiments-with-ray-or-slurm) +- {ref}`Distributed Experiments ` - [Customization: Workflows](../customization/agent.md) - [Agentic RL Tutorial](../tutorial/agentic_rl.md) - [Large MoE Training](../tutorial/megatron.md) diff --git a/docs/tutorial/gsm8k_grpo.md b/docs/tutorial/gsm8k_grpo.md index cfc6defda5..2af5de251c 100644 --- a/docs/tutorial/gsm8k_grpo.md +++ b/docs/tutorial/gsm8k_grpo.md @@ -56,46 +56,46 @@ Controller Process (Your Script) ### Data Flow with RTensor ``` -Rollout Workers (GPUs 0-3) Controller Training Workers (GPUs 4-7) -───────────────────────── ────────── ─────────────────────────── +Rollout Workers (GPUs 0-3) Controller Training Workers (GPUs 4-7) +───────────────────────────── ──────────── ───────────────────────────── Worker 0: Generates 16 samples - ├─> Shard 0 stored ──────┐ -Worker 1: Generates 16 samples │ - ├─> Shard 1 stored ────┐ │ -Worker 2: Generates 16 samples │ │ - ├─> Shard 2 stored ──┐ │ │ -Worker 3: Generates 16 samples │ │ │ - └─> Shard 3 stored ─┐│ │ │ - ││ │ │ - ││ │ │ RTensor metadata - ││ │ └──> Controller ──> data_parallel_dispatch() - ││ └─────────────┼──────────────┬─────────────┐ - │└───────────────┼──────────────┼─────────────┤ - └────────────────┼──────────────┼─────────────┤ - │ │ │ - ▼ ▼ ▼ - Worker 4: Worker 5: Worker 6: - Fetch Fetch Fetch - Shards 0,1 Shards 2 Shards 3 - │ │ │ - ├─> Forward ├─> Forward ├─> Forward - ├─> Backward ├─> Backward ├─> Backward - └─> Gradients └─> Gradients└─> Gradients - │ - NCCL AllReduce - │ - Worker 4: Worker 5: Worker 6: - Returns Returns Returns - RTensor RTensor RTensor - │ │ │ - └──────────────┴─────────────┘ - │ - data_parallel_merge() - │ - ▼ - Controller receives: - • loss (scalar) - • metrics (dict) + ├─> Shard 0 stored ────────────┐ +Worker 1: Generates 16 samples │ + ├─> Shard 1 stored ──────────┐ │ +Worker 2: Generates 16 samples │ │ + ├─> Shard 2 stored ────────┐ │ │ +Worker 3: Generates 16 samples │ │ │ + └─> Shard 3 stored ──────┐ │ │ │ + │ │ │ │ + │ │ │ │ RTensor metadata + │ │ │ └─> Controller ─> data_parallel_dispatch() + │ │ └───────────┼────────────┬────────────┐ + │ └─────────────┼────────────┼────────────┤ + └───────────────┼────────────┼────────────┤ + │ │ │ + ▼ ▼ ▼ + Worker 4: Worker 5: Worker 6: + Fetch Fetch Fetch + Shards 0,1 Shards 2 Shards 3 + │ │ │ + ├─> Forward ├─> Forward ├─> Forward + ├─> Backward ├─> Backward ├─> Backward + └─> Grads └─> Grads └─> Grads + │ + NCCL AllReduce + │ + Worker 4: Worker 5: Worker 6: + Returns Returns Returns + RTensor RTensor RTensor + │ │ │ + └────────────┴────────────┘ + │ + data_parallel_merge() + │ + ▼ + Controller receives: + • loss (scalar) + • metrics (dict) ``` In the following sections, we'll walk through the code to explain each component in @@ -270,11 +270,11 @@ defines how prompts become training samples. Each trajectory goes through these **GSM8K Reward**: Binary reward (1.0 for correct answer, 0.0 otherwise). See [`gsm8k_reward_fn`](https://github.com/inclusionAI/AReaL/blob/main/areal/reward/gsm8k.py). -This workflow adopts the low-level API of inference engines --- the `agenerate` API. It -is perferrable if you want more fine-grained control over token IDs. `agenerate` inputs -token IDs to the inference server and produces output token IDs for user's processing. -We also provide high-level API for convenient agentic workflow orchestration. We refer -to the [agentic RL guide](../tutorial/agentic_rl.md). +**NOTE:** This workflow adopts the low-level API of inference engines --- the +`agenerate` API. It is preferable if you want more fine-grained control over token IDs. +`agenerate` inputs token IDs to the inference server and produces output token IDs for +user's processing. We also provide high-level API for convenient agentic workflow +orchestration. We refer to the [agentic RL guide](../tutorial/agentic_rl.md). ### Asynchronous Rollout Collection @@ -284,28 +284,28 @@ overlap between generation and training. #### Three-Process Architecture ``` -Controller Process Worker Process (RPC Server) GPU Process -────────────────── ─────────────────────────── ─────────── -RolloutController Flask HTTP Server (CPU) SGLang/vLLM - │ │ │ - └─> BatchTaskDispatcher /call endpoint Inference - (background thread) │ Engine - │ └─> Engine Thread │ - ├─ submit task 1 └─> RemoteInfEngine │ - │ (HTTP POST) └─> submit() ────────────>│ - │ Generate - ├─ submit task 2 tokens - │ (HTTP POST) │ - │ │ - ├─ submit task 3 HTTP Callback <─────────────┘ +Controller Process Worker Process (RPC Server) GPU Process +────────────────── ─────────────────────────── ─────────── +RolloutController Flask HTTP Server (CPU) SGLang/vLLM + │ │ │ + └─> BatchTaskDispatcher /call endpoint Inference + (background thread) │ Engine + │ └─> Engine Thread │ + ├─ submit task 1 └─> RemoteInfEngine │ + │ (HTTP POST) └─> submit() ──────────────>│ + │ Generate + ├─ submit task 2 tokens + │ (HTTP POST) │ + │ │ + ├─ submit task 3 HTTP Callback <──────────────┘ │ (trajectory) - │ ┌──────────────┘ + │ ┌─────────────┘ └─ collect <──────┘ Meanwhile (on different GPUs)... TrainController Training Worker │ │ - └─> ppo_update(batch) ──────────>│ Forward/Backward + └─> ppo_update(batch) ──────────> Forward/Backward Key: Generation and training happen SIMULTANEOUSLY on different GPUs ``` @@ -321,6 +321,9 @@ HTTP: - Maintains 2+ batches of inflight requests to hide latency - Non-blocking: returns task_id immediately +As such, **rollout and training happen simultaneously** in AReaL, even though the code +looks like a synchronous orchestration. + **Level 2 - Worker RPC Server**: Each rollout worker runs a Flask HTTP server ([`rpc_server.py`](https://github.com/inclusionAI/AReaL/blob/main/areal/scheduler/rpc/rpc_server.py)) on **CPU**: @@ -368,21 +371,6 @@ results = dispatcher.wait_results(batch_size) return concat_padded_tensors(results) # Shape: [batch_size, seq_len] ``` -#### Overlap with Training - -The key benefit of this architecture is that **rollout and training happen -simultaneously**: - -``` -Timeline: -───────────────────────────────────────────────────────────── -Rollout GPUs: [Generate Batch N+1] [Generate Batch N+2] ... -Training GPUs: [Train on Batch N] [Train Batch N+1] ... - │ - Weight sync happens here - (rollout paused briefly) -``` - **Staleness Control**: [`StalenessManager`](https://github.com/inclusionAI/AReaL/blob/main/areal/infra/staleness_manager.py) limits concurrent inflight requests: @@ -391,16 +379,6 @@ limits concurrent inflight requests: - `max_head_offpolicyness`: Reject samples generated with weights too old - Version tracking: Each token tagged with model version used during generation -**Pause/Resume**: During weight sync, rollout is paused to avoid stale generations: - -```python -# In PPOTrainer.train() loop -rollout.pause() # Pause new submissions -actor.update_weights(...) # Sync weights to inference GPUs -rollout.set_version(step) # Update version tracker -rollout.resume() # Resume submissions -``` - ## Training: Controller-Worker Pattern Training follows a standard controller-worker pattern. The controller dispatches @@ -419,24 +397,24 @@ provides the core RPC dispatch: **Data Flow with RTensor:** ``` -Controller Worker 0 Worker 1 - │ │ │ - ├─ RTensor (metadata) ─────┼───────────────────────┤ - │ • Shard 0,1,2,3 │ │ - │ │ │ - ├─ dispatch() ─────────────>│ │ - │ • Worker 0: Shards 0,1 │ │ - │ • Worker 1: Shards 2,3 │ │ - │ │ │ - │ ├─> Fetch Shards 0,1 │ - │ │ from rollout workers│ - │ │ ├─> Fetch Shards 2,3 - │ │ │ from rollout workers - │ │ │ - │ ├─> compute_logp() ├─> compute_logp() - │ │ │ - │ ├─> RTensor (result) ├─> RTensor (result) - │<─ merge() ───────────────┴───────────────────────┘ +Controller Worker 0 Worker 1 + │ │ │ + ├─ RTensor (metadata) ──────┼─────────────────────────┤ + │ • Shards 0,1,2,3 │ │ + │ │ │ + ├─ dispatch() ────────────> │ │ + │ • Worker 0: Shards 0,1 │ │ + │ • Worker 1: Shards 2,3 │ │ + │ │ │ + │ ├─> Fetch Shards 0,1 │ + │ │ from rollout workers │ + │ │ ├─> Fetch Shards 2,3 + │ │ │ from rollout workers + │ │ │ + │ ├─> compute_logp() ├─> compute_logp() + │ │ │ + │ ├─> RTensor (result) ├─> RTensor (result) + │<─ merge() ────────────────┴─────────────────────────┘ │ • Reconstruct ordering │ • Return unified RTensor └─> batch["logp"] = result @@ -520,9 +498,9 @@ supports two transfer methods: **NCCL-based transfer** (Recommended): -- Direct GPU-to-GPU broadcast +- Direct GPU-to-GPU communication based on NCCL broadcast - Faster but uses more GPU memory -- Requires training and inference GPUs on the same communication backend +- Requires non-overlapped training and inference GPUs on the same communication backend **Disk-based transfer**: @@ -533,10 +511,11 @@ supports two transfer methods: The weight sync process in `PPOTrainer.train()` follows this pattern: -1. Pause rollout to avoid stale generations +1. Pause rollout servers to interrupt all inflight generations back to the rollout + client (e.g., `RemoteSGLangEngine`) 1. Transfer weights via configured method (NCCL or disk) 1. Update version tracking for staleness management -1. Resume rollout with updated weights +1. Resume rollout with updated weights with re-computed KV cache See [`PPOTrainer.train()`](https://github.com/inclusionAI/AReaL/blob/main/areal/experimental/trainer/rl.py) @@ -549,9 +528,22 @@ metrics tracking. These are automatically orchestrated during training. ### Checkpointing -The [`Saver`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/saver.py) -handles periodic checkpoint saving. Configure via `config.saver` (interval, format, -etc.). Called automatically in `trainer.train()`. +AReaL provides two checkpointing mechanisms: + +| Component | Purpose | Format | Configuration | +| ----------------------------------------------------------------------------------------- | -------------------------------- | ------------- | ---------------- | +| [`Saver`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/saver.py) | Export for evaluation/deployment | HuggingFace | `config.saver` | +| [`RecoverHandler`](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/recover.py) | Resume after failures | DCP (sharded) | `config.recover` | + +**Saver** creates HuggingFace-compatible checkpoints that can be loaded with +`transformers` or published to HuggingFace Hub. Each save creates a new directory. + +**RecoverHandler** saves complete training state (model, optimizer, dataloader, RNG) for +fault tolerance. Checkpoints are backend-specific and require the same parallelism +configuration to load. Each save overwrites the previous checkpoint. + +Both are called automatically during `trainer.train()`. For details, see the +[Checkpointing Reference](../reference/checkpointing.md). ### Evaluation @@ -562,21 +554,41 @@ automatically in `trainer.train()`. ### Metrics Tracking +AReaL uses a two-component metrics system: + **`stats_tracker`** ([source](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/stats_tracker.py)): -Collects and aggregates training statistics across ranks. +Collects statistics with two paradigms optimized for different use cases: -- `scalar(key=value)`: Record simple metrics -- `stat(key=tensor, denominator=mask)`: Record tensor statistics with selective - aggregation -- `record_timing(name)`: Context manager for timing -- `scope(name)`: Hierarchical metric keys -- `export()`: Returns aggregated stats across all ranks +- **Streaming metrics** for rollout workers: Each workflow logs scalars individually + (e.g., `reward`), which are aggregated across workers by the controller +- **Batch metrics** for training: Tensor statistics with boolean masks are logged per + batch, then all-reduced across data-parallel ranks + +```python +# Rollout metrics (streaming) - in workflows +stats_tracker.get("rollout").scalar(reward=0.8, num_turns=3) + +# Training metrics (batch) - in PPO actor +stats_tracker.denominator(n_valid_tokens=loss_mask.bool()) +stats_tracker.stat(advantages=tensor, denominator="n_valid_tokens") +``` **`StatsLogger`** ([source](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/stats_logger.py)): -Sends metrics to logging backends (W&B, TensorBoard) from rank 0. Configure via -`config.stats_logger`. +Sends aggregated metrics to logging backends (Weights & Biases, SwanLab, TensorBoard) +from rank 0. At each training step, `PPOTrainer` collects metrics from all components +and commits them: + +```python +# areal/experimental/trainer/rl.py +stats = self.actor.export_stats() # Training metrics +stats.update(self.rollout.export_stats()) # Rollout metrics +self.stats_logger.commit(epoch, step, global_step, stats) # → wandb/tensorboard +``` + +For the complete API reference, see the +[Metrics Tracking Reference](../reference/metrics_tracking.md). ## Next Steps @@ -584,6 +596,7 @@ Now that you understand the basics, explore these advanced topics: **Tutorials**: +- [Evaluation](../tutorial/eval.md) - Evaluate your trained model - [Training Large MoE Models](../tutorial/megatron.md) - Scale to massive models with Megatron integration - [Agentic RL with OpenAI APIs](../tutorial/agentic_rl.md) - Build agents that use tools diff --git a/docs/tutorial/megatron.md b/docs/tutorial/megatron.md index 1dcee84f20..ed44da3d79 100644 --- a/docs/tutorial/megatron.md +++ b/docs/tutorial/megatron.md @@ -10,69 +10,28 @@ large MoE models for your application. Shifting from FSDP to Megatron requires only a single line of change: the `allocation_mode` field from `sglang:d4+fsdp:d4` to `sglang:d4+megatron:d4`. -We already have some internal logic for determining the backend to use if the backend -name is omitted. If neither pipline parallelism nor expert parallelism is enabled, FSDP -will be used as the backend. Otherwise, Megatron will be used. However, we encourage -specifying the backend name explicitly like above. +For a complete guide on allocation mode syntax, parallelism dimensions, and GPU +calculations, see the [Allocation Mode Reference](../reference/alloc_mode.md). -## Understanding `allocation_mode` +## MoE Parallel Strategy -The allocation mode is defined in -[areal/api/alloc_mode.py](https://github.com/inclusionAI/AReaL/blob/main/areal/api/alloc_mode.py). -The allocation mode is a pattern-based string option that tells AReaL how to parallelize -models across GPUs in training and inference backends. When running the experiment, -AReaL converts the string option into an `AllocationMode` object that stores the backend -choice and parallel strategy for each model. For a simple example, -`sglang:d4+megatron:t4` configures AReaL to use the SGLang backend with **data -parallel** size 4 and the Megatron training backend with **tensor parallel** size 4. +For MoE models, Megatron supports separate parallelism for attention and FFN modules +using the hybrid syntax. For example: -### Training Parallel Strategy - -For a dense model, there are only 4 available parallel dimensions: data parallel (DP, -d), tensor parallel (TP, t), pipeline parallel (PP, p), and context parallel (CP, c). -The numbers that follow the single-character abbreviation of parallel dimensions -describe the parallel size. For example, `megatron:d2t4p2c2` describes a 32-GPU parallel -strategy that has DP size 2, TP size 4, PP size 2, and CP size 2. +``` +megatron:(attn:d1p4t2c2|ffn:d1p4t1e4) +``` -For MoE models, the AReaL allocation mode supports separate parallel strategies for -expert modules and attention modules, which is related to the +This 16-GPU configuration uses PP=4, with attention modules using TP=2 and CP=2, while +expert modules use TP=1 and EP=4. See [MoE Parallel Folding](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/moe#moe-parallel-folding) -feature in Megatron. It reduces the minimal number of GPUs required to enable both -context and expert parallelism (EP, e), and enables different TP sizes for attention and -expert modules for better efficiency. The parallel strategies for attention and expert -modules are denoted by `attn:` and `ffn:`, and separated by `|`. For example, -`megatron:(attn:d1p4t2c2|ffn:d1p4t1e4)` describes a 16-GPU parallel strategy with PP -size 4, that has DP size 1, TP size 2, and CP size 2 for attention modules and DP size -1, TP size 1, and EP size 4 for expert modules. +for details on this feature. -**5D parallel strategy Tuning Guides:** +**Tuning Guides:** - [Megatron Performance Best Practice](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/moe#performance-best-practice) - [verl with Megatron Practice](https://github.com/ISEEKYAN/verl_megatron_practice) -### Inference Parallel Strategy - -The optimal parallel strategy is ususally different for training and inference. -Inference parallel strategies only accept DP, TP, and PP, e.g., `vllm:d2t4`. Note that -DP degree is the number of independent instances to deploy. Other parallelism -configurations are passed through the `sglang` and `vllm` field in configurations, e.g., - -```yaml -sglang: - ep_size: 2 - dp_size: 4 - enable_dp_attention: true - ... -``` - -Note that the above configurations controls the internal hybrid parallelism strategy -within each inference instance, e.g., DP attention. These techniques are ususally not an -orthogonal dimension to DP, TP, and PP that determine GPU allocation. We refer to the -large-scale EP delopment guide of -[SGLang](https://lmsys.org/blog/2025-05-05-large-scale-ep/) and -[vLLM](https://docs.vllm.ai/projects/ascend/en/v0.9.1-dev/developer_guide/performance/distributed_dp_server_with_large_ep.html) -for detailed information. - ## Aligning Inference and Training Precision Due to the sparse nature of MoE models, the logits calculated by forward passes during diff --git a/docs/tutorial/troubleshooting.md b/docs/tutorial/troubleshooting.md deleted file mode 100644 index 9031249800..0000000000 --- a/docs/tutorial/troubleshooting.md +++ /dev/null @@ -1,89 +0,0 @@ -# Troubleshooting - -If the following content does not address your issue, feel free to raise a GitHub Issue. - -## Automatic Recovery - -When setting `recover_mode=auto` and the experiment configuration remains unchanged, -AReaL will attempt to discover previous checkpoints and recover the experiment from -them. - -### Recovery Failure Causes - -If automatic recovery fails, check the following possibilities: - -**Configuration Changes:** - -- The `experiment_name` and `trial_name` in the training script differ from the previous - run -- Changes in batch size (`dataset.train_bs_n_seqs` parameter) -- Changes in group size (`group_size` parameter) -- Changes in number of nodes (`n_nodes` parameter) - -**Missing Recovery Checkpoints:** Recovery checkpoints are generated under two -conditions by default: - -- After completion of the second step -- When a step completes and more than 600 seconds have passed since the last recovery - checkpoint (controlled by `exp_ctrl.ckpt_freq_secs=600`) - -### Verify Recovery Checkpoint Creation - -You can confirm if a recovery checkpoint was generated by searching for the following -message in the logs: - -```bash -(master_worker/0 pid=96390, ip=xxx.xxx.xxx.xxx) 20250222-11:52:02.760 master worker INFO: Dumped recover info to file. -``` - -## Memory Issues - -### torch.cuda.CudaOutOfMemoryError - -The key to resolving this issue is identifying the phase where the error occurs: - -#### During Initialization - -- Check for idle processes on the GPU -- **Distributed scenarios**: Restart the Ray cluster -- **Single-machine scenarios**: Use `pkill` to terminate processes - -#### During SGLang Generation - -- Decrease the `actor.sglang.mem_fraction_static` parameter -- Increase the tensor parallelism degree -- Decrease the `max_concurrent_rollouts` parameter for asynchronous RL - -#### During `actor_inf` or `actor_train` - -- **Adjust microbatch size**: Decrease the parameter - `{actor_train|actor_inf}.mb_spec.max_tokens_per_mb=20480`. This parameter limits - tokens per forward/backward pass and can be set as low as the maximum sequence length - (including prompt) -- **Modify parallelism strategy**: Adjust `allocation_mode` by: - - Reducing data parallelism - - Increasing tensor or pipeline parallelism - - Preferring pipeline parallelism over tensor parallelism - -### CUDA Error: Out of Memory - -This issue may occur during data transfer. Try increasing `mem_per_model_worker` in the -CLI arguments. - -## Permission Issues - -### Cache Directory Permission Errors - -If you encounter permission errors related to `/tmp/areal`, this typically occurs in -multi-user environments where another user has already created the cache directory with -restrictive permissions. - -**Solution:** Set the `AREAL_CACHE_DIR` environment variable to a user-specific path: - -```bash -export AREAL_CACHE_DIR=/tmp/areal-$USER -# or use your home directory -export AREAL_CACHE_DIR=$HOME/.cache/areal -``` - -You can add this to your `~/.bashrc` or `~/.profile` to make it persistent.