[WIP] [Refactor] Add AReaLite API and examples.

[garrett4wade]

Motivation

AReaL is too heavy for AI researchers to use, understand, and develop with for several reasons. The most important issue is that its code architecture is system-centric rather than AI-centric — the RL algorithm workflow consists of multiple workers that run consecutive model function calls, neither of which are well-known concepts for AI researchers. As a result, users must first understand these concepts before they can develop workflows and algorithms for their own use cases.

Additionally, due to historical reasons, AReaL's code is not clean. There are large pieces of code inherited from previous projects that are not useful but significantly increase the burden on users and developers. Sometimes debugging is difficult even for core developers like myself.

Since the tools for building RL workflows are becoming increasingly mature, implementing a framework that achieves comparable efficiency requires much fewer lines of code. Now is the proper time to revisit the API design and distill the giant codebase into a neat and clean one. The distilled codebase does not need to be ultra-efficient. Instead, we want to deliver 90% functionality of the original AReaL while minimizing the lines of code and the burden on potential users. Our aim is to build an RL training framework that is fast to use, fast to read, and fast to execute. Here comes the lite version of AReaL — AReaLite.

AReaLite is the first step in AReaL's refactoring process. It is not only a standalone training library with shallow interfaces, but will also provide the core API definitions to be used by AReaL in the future. AReaL will essentially transform its current worker-based architecture into an AI-centric architecture like AReaLite. AReaL will extend AReaLite's APIs and implementations to support more backends for efficient large-scale training.

Expectations of AReaLite

Highlights

• Fully asynchronous training with decoupled inference and training.
• Elastic inference device scaling — users can shut down or launch more inference processes independently during training.
• Full SFT/RL algorithmic functionality matching AReaL.
• Arbitrary agentic rollout workflow customization in a single file.
• Easy navigation to implementation references via Ctrl+click in VSCode.
• Support for distributed launching with Ray/SLURM/torchrun.

AReaLite's Scope

• Not bound to Ray.
• Only supports SGLang and PyTorch FSDP2 with SPMD launching.
• No customized data structures like SequenceSample. All data are PyTorch tensors.
• Uses HuggingFace (models, datasets) and PyTorch (FSDP, data structures) as much as possible.

Architecture

Core Components

arealite/
├── api/           # Abstract interfaces and data structures
├── impl/          # Concrete implementations
├── cli/           # Command-line interfaces
├── config/        # Configuration templates
└── tests/         # Standalone test scripts

Data Flow Architecture

AReaLite uses an async producer-consumer pattern:

┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│   LLM Servers   │◄──►│ Rollout Workers  │───►│   Data Buffer   │
│   (SGLang)      │    │  (Async Batch)   │    │                 │
└─────────────────┘    └──────────────────┘    └─────────────────┘
        ▲                                               │
        │                                               ▼
┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│   Checkpoints   │◄───│  FSDP Trainer    │◄───│ Training Loop   │
│                 │    │  (Sync Batch)    │    │                 │
└─────────────────┘    └──────────────────┘    └─────────────────┘

Key Design Principles

1. AI-Centric API Design

Unlike the original AReaL's system-centric approach with workers and model functions, AReaLite uses familiar ML concepts:

• Agent and Environment (from RL literature)
• RolloutWorkflow (combines multiple agents and the environment to generate rollout data)
• Trainer (from HuggingFace/PyTorch, fetches rollout data and updates model parameters)

2. Factory Pattern for Extensibility

Each major component uses a factory pattern for easy customization:

• EngineFactory creates training backends
• TrainerFactory creates training algorithms
• RolloutWorkflowFactory creates rollout workflows

3. Configuration-Driven Architecture

All components are configured through dataclasses defined in cli_args.py, enabling:

• Type-safe configuration
• Easy CLI argument generation
• Clear documentation of available options

Roadmap

• Finalize API design. (In-progress)
• Implement standalone SGLang server (impl/sglang_server.py).
• Implement SGLang client generation (impl/sglang_client.py).
• Rollout pipeline (tests/test_rollout.py).
• SGLang rollout interruption.
• Asynchronous RL system-wide utilities (e.g., RolloutController).
• Various launching scripts: ray, torchrun, slurm.
• FSDP2 engine with transformers models. (In-progress)
• SFT trainer. (In-progress)
• SGLang update weights. (In-progress)
• PPO trainer. (In-progress)
• Add benchmarking against original AReaL
• CI and unittests.
• Other RL algorithms (DPO, REINFORCE, etc.)
• Support for multi-modal models
• User guide for transitioning from v0.3.0.
• Advanced agentic workflows (tool use, planning)
• Examples of training GSM8K, TLDR, and a search agent.
• Allow external persistent SGLang servers for debugging purposes.

[tsaoyu]

Add a main entry point annotation somewhere?

[garrett4wade]

Added a data flow arch graph in the code walk-through example: https://github.com/inclusionAI/AReaL/blob/lite/docs/arealite/gsm8k_grpo.md

[tsaoyu]

One thing in production system often lacks is systematic support of reward model, verifier and generative reward model. From the RL point of view it can be encapsulated in the concept of environment but it is curbersome to integrate them. Method a bit here might be good for those who are using RL in production.

[garrett4wade]

Reward models can be implemented as additional TrainEngines, just like the reference model. Generative reward models is encapsulated in the RolloutWorkflow object (aka the previous RolloutCollector), e.g., the agent call use the same LLM but with different prompts to judge whether the previous generated answer is correct.

[tsaoyu]

mp -> tp ?

[garrett4wade]

No "mp" any more.

[tsaoyu]

Is there any data validation similar to PyDantic to check range of parameters here?

[garrett4wade]

Good idea. I'll mark that.

[tsaoyu]

I saw there is uvloop as the async manager. Does it actually used in the project

[garrett4wade]

Changed to uvloop.run wherever using asyncio.run.

[tsaoyu]

The catch seems a bit too wide

[tsaoyu]

**pad_sequences_to_tensors(data) in function? return data=data_dict? TensorDict is not that bad though

[garrett4wade]

Uses TensorDict as the basic data structure

[tsaoyu]

This is good, very clear on prepare_batch and generate_batch

[tsaoyu]

Throw some hints here such as f"Micro batch mismatch, current {mb_stats} and config {self.config.ppo_n_minibatches}. Check your configuration. "

[tsaoyu]

haha, why this is empty

[garrett4wade]

It's not empty in the new lite branch. :)

[garrett4wade]

@tsaoyu Hi Yu, thank you for the thorough review!

We're currently finalizing some internal discussions about the API design, and the final implementation will differ slightly from what's currently proposed. We'd like to defer addressing these specific issues for a few days while we complete those discussions.

Your feedback is valuable and we'll incorporate it into our revised implementation. We'd appreciate having you review again once the final version is ready.

Thanks for your patience!

[garrett4wade]

Closed since #154 has been merged.