- Runtime: Designed for distributed GPU clusters (FSDP/Megatron + distributed communication libraries). Assume containerized execution; no standalone local runs.
- Testing: Integration and performance tests require multi-node hardware. Explain skips explicitly when you cannot access the cluster.
- Tooling:
.pre-commit-config.yamlruns Ruff (lint+format), mdformat, clang-format, nbstripout, and CLI doc generation; install withpre-commit installbefore submitting patches. - Formatting: Ruff + Ruff-format replace Black/isort; autoflake settings remain in
pyproject.toml. Surface any formatting gaps you cannot auto-fix. - Docs: Source lives under
docs/(Jupyter Book). Coordinate doc edits with the docs build pipeline. - Collaboration: Before editing code, outline the proposed plan and confirm it with the user.
When unsure, leave a TODO(agent) comment and note the constraint in your response.
.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, distributed, testing).claude/skills/- Guided development workflows (add-dataset, add-workflow, add-reward, add-archon-model, 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 shared utilities:areal/api/- Contracts for workflows, engines, schedulers, IO structs, and CLI/config dataclasses.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, TORL, etc.) and utilities that feed rollout jobs safely.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 breaking changes.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 wrappers, custom heads).areal/reward/- Built-in reward functions (GSM8K, Geometry3K, CLEVR, etc.), math parsers, and helpers; wrap slow logic withAsyncRewardWrapper.areal/scheduler/- Placement and allocation policies aligned with launcher configs.areal/tests/- Focused unit/integration suites (many require GPUs or mocked distributed backends).areal/tools/- Developer utilities and maintenance scripts tied to the core package.areal/utils/- Cross-cutting helpers for logging, tensor ops, stats tracking, checkpoints, and recovery.areal/workflow/- Concrete rollout agents implementingRolloutWorkflow: multi-turn, RLVR, vision RLVR workflows, plusopenai_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 reference generators.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 experimentation.
- Clusters & containers: Launch configurations live under
areal/launcher/(local, Ray, Slurm). Each entrypoint documents the scheduler expectations; reuse those specs instead of inventing ad-hoc run scripts. - Shared images: Platform-specific container images and startup scripts are defined alongside launcher configs. Reference them or note when they are missing—do not attempt to rebuild CUDA/driver stacks inline.
- Secrets & endpoints: Credentials for remote inference (SGLang, vLLM, Redis, etc.) are managed outside the repo. Flag their absence rather than hardcoding replacements.
- Testing limitations: End-to-end tests (FSDP, Megatron, distributed RPC) require multi-node clusters using distributed communication libraries. If you cannot execute them, state that your validation is limited to static analysis/doc updates.
- Formatting & docs: Pre-commit runs Ruff (lint+format), mdformat, clang-format,
nbstripout, and CLI doc generation. Run
pre-commit run --all-files(or install the hook) before submitting; keep doc edits aligned with the Jupyter Book structure indocs/.
- Typing & dataclasses: Prefer explicit type hints and reuse existing dataclasses in
areal/api/cli_args.pywhen extending configs. When adding new configuration options, extend an existing dataclass if your changes are backward-compatible or the new config is a strict superset of an existing one. Create a new dataclass if the config is conceptually distinct or would introduce breaking changes. Keep new configs dataclass-based so Hydra/CLI integration stays consistent. - Imports: Avoid wildcard imports; keep third-party vs internal groups consistent. Ruff enforces import ordering (isort rules) when hooks run. Place heavy optional deps inside functions to prevent import-time side effects.
- Logging: Use
areal.utils.logging.getLogger(__name__)rather thanprint. Emit structured metrics throughstats_tracker/StatsLoggerinstead of ad-hoc counters. - Async code: Rollout workflows must stay non-blocking—prefer
awaitwithaiofiles, avoid synchronous file I/O insidearun_episode, and guard long-running CPU work with executors if needed. - Tensor shapes: Follow padded batch conventions; validate with
check_trajectory_formatwhile developing. Use helpers inareal.utils.datafor padding/broadcasting. - Config overrides: Keep Hydra-friendly dotted names; don’t hardcode paths—expose options in config dataclasses and wire via YAML.
- Testing: New features should ship with targeted pytest coverage (mark GPU-heavy
suites appropriately). Use
pytest.skipwith a clear reason when hardware isn’t guaranteed. - Docs & comments: Document non-obvious behaviors inline; prefer short module-level docstrings summarizing workflows or engines you touch.
- Trainer (
areal/trainer/) – High-level training orchestrator. UsePPOTrainerfor RL training orSFTTrainerfor supervised fine-tuning. Seeexamples/math/gsm8k_rl.pyfor a complete example:from areal import PPOTrainer with PPOTrainer(config, train_dataset, valid_dataset) as trainer: trainer.train(workflow="areal.workflow.rlvr.RLVRWorkflow", ...)
- Rollout workflows (
areal/workflow/,areal/api/workflow_api.py) – ImplementRolloutWorkflow.arun_episode. Use helpers likeconcat_padded_tensorsand respect shape[batch, seq_len, …]. Pass workflow class path totrainer.train(). - Inference engines (
areal/engine/) – Handle async generation and weight updates. Interact with workflows viaInferenceEngine.agenerate. Includes SGLang/vLLM remote adapters. - Training engines (
areal/engine/) – Consume rollout tensors, run PPO/GRPO updates, broadcast weight versions. Includes FSDP2 and Megatron backends. - Rewards (
areal/api/reward_api.py,areal/reward/) – Wrap blocking reward code inAsyncRewardWrapper. Standard signature:(prompt, completions, prompt_ids, completion_ids, **data). - Configurations – Dataclasses in
areal/api/cli_args.py, YAML examples inexamples/**. Launchers parse CLI overrides (Hydra-style dotted keys).
Reference docs:
- Agents customization guide:
docs/customization/agent.md. - Lite design doc:
areal/README.md. - Algorithm-specific docs:
docs/algorithms/*.md.
- Start from the existing patterns in
areal/workflow/multi_turn.py,rlvr.py, orvision_rlvr.py, then add a sibling module underareal/workflow/that subclassesRolloutWorkflow. - In
__init__, thread throughGenerationHyperparameters, the tokenizer, reward callable, stat scope, and optionaldump_dir; wrap the reward viaAsyncRewardWrapperexactly likeMultiTurnWorkflowdoes. - Keep
arun_episodeasync-only, drive generation throughInferenceEngine.agenerate, and emit tensors usingconcat_padded_tensorsso outputs stay[batch, seq_len, ...]. - Use
areal/utils/data.pyhelpers for padding/broadcasting,areal/utils/loggingfor logger plumbing, andstats_trackerfor reward metrics. - Persist transcripts under
{dump_dir}/{engine.get_version()}/(follow themulti_turnimplementation) when debugging is enabled. - Update whichever entry script or launcher references the workflow (e.g.,
examples/multi_turn_math/gsm8k_rl_mt.py, configs inexamples/**/conf/, or CLI glue) so Hydra can import the new module.
- Create
areal/reward/<name>.pyand implement a callable followingareal/api/reward_api.py(seegeometry3k_reward_fnfor reference). - Accept
(prompt, completions, prompt_ids, completion_ids, **data)and return a scalar; extract answers deterministically (MathVerifyWorker, regex, etc.) and avoid blocking I/O. - Add the identifier to
VALID_REWARD_FNand branch selection logic inareal/reward/__init__.pyso configs likereward.path=...resolve automatically. - When rewards rely on slow models or external services, keep the heavy code inside the
reward module but let workflows wrap it with
AsyncRewardWrapper(as inMultiTurnWorkflow). - Document required dataset fields or endpoints in the module docstring/README so launch scripts can provision secrets or caches.
- Mirror the layout in
areal/dataset/gsm8k.py,geometry3k.py,clevr_count_70k.py,hhrlhf.py, ortorl_data.py: createareal/dataset/<name>.pywithget_<name>_<type>_datasethelpers for SFT/RL variants. - Update
areal/dataset/__init__.pyby appending the dataset toVALID_DATASETSand adding a dispatch branch inside_get_custom_dataset. - Define the sample schema explicitly (
messages,answer,image_path, metadata) and validate it before returning; filter/trim sequences with tokenizer-aware checks whenmax_lengthis provided. - Expose configuration knobs (path, split, type, max_length, processor/tokenizer
requirements) via the
TrainDatasetConfigandValidDatasetConfigdataclasses inareal/api/cli_args.py, then reference them in the relevantexamples/**/confYAML. - If preprocessing or external storage is required, add a short note beside the loader
or under
examples/<recipe>/README.mdso other agents know how to stage data.
- Choose an existing script in
examples/**(math, multi-turn, VLM, etc.) that mirrors your use case, then replicate its launcher pairing (areal/launcher/local.py,ray.py,slurm.py, orsglang_server.py). - Read the example README to collect scheduler requirements, container images, environment variables, and any dataset preparation steps before running.
- Keep rollout actors and inference engines version-aligned by propagating
WeightUpdateMeta(as shown inexamples/multi_turn_math/gsm8k_rl_mt.py) and noting skipped weight updates explicitly if clusters are unavailable. - Capture the Hydra/CLI overrides you used
(
python ... +train_dataset.path=... engine.type=...) inside the PR/test plan so runs are reproducible. - When cluster access is blocked, document which launcher stages were skipped and what validation (unit tests, static checks) you ran instead.
- Place prose in the right section under
docs/(tutorial, algorithms, customization, etc.) and update_toc.ymlso Jupyter Book exposes the new page. - Run
mdformat(ormdformat --check) on edited Markdown plusruff formaton embedded code blocks when needed. - Regenerate CLI docs with
python docs/generate_cli_docs.pywheneverareal/api/cli_args.pyor CLI entrypoints change, then restagedocs/cli_reference.md. - Coordinate a docs build (or explain why it is skipped) and capture the limitation in your PR/testing notes if the hosted pipeline cannot run.
- Emit rollout/training metrics through
areal/utils/stats_tracker.py; grab a scoped tracker (stats_tracker.get("rollout")) and log scalars so downstreamStatsLoggerbackends (W&B/SwanLab) pick them up automatically. - When debugging, pass
dump_dirinto workflows so transcripts persist under{dump_dir}/{engine.get_version()}/likeareal/workflow/multi_turn.py; scrub sensitive data before committing artifacts. - Checkpoint via
areal/utils/saver.pyand resume withareal/utils/recover.py; note the checkpoint path and version in your PR/test notes so others can reproduce the exact state.
- Place new tests under
areal/tests/usingtest_<topic>.pyso Pytest auto-discovers them (e.g., tensor helpers live intest_utils.py, schedulers intest_local_scheduler.py). - Reuse fixtures + helpers: copy the pattern from
test_utils.py(local fixtures feeding parametrized cases) or import shared logic fromareal/tests/utils.py(is_in_ci,get_bool_env_var). Preferpytest.fixture+pytest.mark.parametrizeover ad-hoc loops. - Keep tests hermetic by mocking engines/workflows similar to
test_inference_engines.py; avoid spinning up real clusters unless you are undertorchrun/orexperimental/. - For GPU/distributed requirements, gate with
pytest.mark.skipifor custom env checks (seetest_fsdp_engine_nccl.pyandareal/tests/torchrun/), and document the hardware dependency inside the skip reason. - When tests need sample artifacts (configs, datasets), reuse the examples in
areal/tests/sftorareal/tests/grporather than downloading new assets. Commit only lightweight fixtures.
- Unit suites: Target the file you touched, e.g.,
pytest areal/tests/test_utils.py. If a full run is infeasible, list the exact command you would have executed. - Workflow smoke tests:
areal/tests/grpoexercises rollout loops and expects CUDA; acknowledge when skipped. - Distributed/FSDP suites:
test_fsdp_*,test_inference_engines.py, RPC/torchrun folders require multi-node setups and distributed communication libraries. Call out the limitation explicitly. - Static checks: Pre-commit runs Ruff lint/format, mdformat, clang-format, nbstripout, CLI doc regeneration, and autoflake. Note if hooks were not run locally and why.
Always mention resource requirements in PRs and in agent responses when tests are skipped.
- Branches: Use kebab-case summaries (e.g.,
feature/multi-turn-metrics,bugfix/fsdp-weight-sync) so PR automation and reviewers can parse intent quickly. - Commits: Follow Conventional Commit prefixes (
feat:,fix:,docs:, etc.), keep the subject around 72 characters for readable logs (go longer only when the extra context is essential), write in imperative voice, and put deeper reasoning in the body. Squash noisy WIP commits before opening or updating a PR. - Pre-merge checks: Run the full pre-commit stack (Ruff lint+format, mdformat,
clang-format, nbstripout, CLI docs, autoflake). For doc-only edits, at least run
mdformat --checkon touched files and call out anything you could not run locally. - Surface scope upfront: Tie the PR to a filed issue, summarize acceptance criteria, highlight risk areas (breaking changes, performance regressions), and note any configs, datasets, or launchers impacted.
- Testing evidence: List the exact commands you executed (unit, workflow smoke, docs build). When hardware is unavailable, state the skipped suites, why they were skipped, and what alternative validation (static analysis, mocks) you performed.
- Async + resource safety: When touching workflows/engines, confirm async code
awaits I/O, avoids blocking calls, and preserves weight versioning
(
set_version/update_weights). Document memory/GPU expectations and dataset or checkpoint storage requirements inside the PR/test notes. - Style, config & docs: Ensure Ruff/clang-format/autoflake output is clean. Thread
new options through the right dataclasses/YAMLs, update docs/CLI references, and
verify hyperlinks. Mention any formatting gaps you plan to address later via
TODO(agent). - Observability & cleanup: Keep metrics flowing through
stats_tracker/StatsLogger, expose dump directories when debugging, and remove stray debug prints or commented code. Note migrations or recovery steps when checkpoints or evaluators change so reviewers know what to verify.
- Docs portal (
https://inclusionai.github.io/AReaL/): Hosted Jupyter Book with the full table of contents; use it to cross-check rendered diagrams, formulas, and links. - Tutorials & quickstart (
docs/tutorial/quickstart.md): End-to-end GSM8K GRPO run covering single-node LocalLauncher flows, Ray/Slurm deployment knobs, and SkyPilot recipes. - Lite deep dive (
docs/tutorial/gsm8k_grpo.md): Architecture-level walkthrough of how launchers, RemoteSGLangEngine, workflows, and FSDP PPO actors coordinate during asynchronous GRPO on GSM8K; great for understanding control flow before editing engines or workflows. - Customization guides (
docs/customization/*.md): Step-by-step patterns for adding datasets, authoring newRolloutWorkflowsubclasses, or wiring custom RL algorithms while keeping configs Hydra-friendly. - Algorithm notes (
docs/algorithms/*.md): Reference math + configuration advice for GRPO, DAPO/DAPO-style filters, async RL, GSPO, LitePPO, m2po, rloo, etc., including when to switch between synchronous and asynchronous modes. - Best practices (
docs/best_practices/*.md): Practical debugging playbooks, reward-drift diagnostics, OOM mitigation, and performance profiling checklists you should cite when explaining skipped tests or perf limitations. - CLI & doc tooling (
docs/cli_reference.md): Auto-generated CLI argument catalog plus instructions for regenerating docs/CLI output before landing config changes. - Benchmarks & reproducibility (
docs/references/*.md): Canonical benchmark setups, dataset/model combos, and experiment-log expectations to mention in PR validation notes. - Version history (
docs/version_history.md): Release timeline noting major API moves, deprecations, and migration steps from legacy AReaL to AReaL-lite (the current AReaL).