AxisRL is an agentic reinforcement learning post-training framework that connects SGLang-based rollout and Megatron-based distributed training for large-scale, multi-turn agent workflows.

Stars

445

7-day growth

No data

Forks

20

Open issues

0

License

Apache-2.0

Last updated

2026-07-24

AI repository intelligence
FR-AI / ANALYSIS

Why it is worth attention

It addresses the unique challenges of agentic RL—long trajectories, tool calls, and weight synchronization—by providing built-in partial rollout, handle-based data movement, and mismatch analysis for rollout-trainer consistency, all supported by production-grade serving and training engines.

Who it is for

  • Researchers and engineers working on post-training of large language models
  • ML infrastructure teams managing distributed RL systems
  • Developers building agentic AI applications with tool-use and multi-turn environments

Use cases

  • Training math reasoning models with GRPO or PPO on datasets like GSM8K
  • Training retrieval-augmented multi-turn search agents (Search-R1)
  • Complex tool-use agent training via black-box harness integration (e.g., OpenHands/E2B)
  • Debugging training stability issues through spike replay and mismatch analysis

Strengths

  • Built on SGLang for high-throughput rollout and Megatron for scalable distributed training
  • Supports long-context (300+ turn) trajectories and hundreds-of-billions-parameter training
  • Improves efficiency with partial rollout, handle-based data plane, and prefix-tree merge with MagiAttention
  • Provides comprehensive rollout-trainer consistency checks (mismatch analysis, routing replay, spike replay)

Considerations

  • Tightly coupled with SGLang and Megatron, requiring specific Docker images and engine expertise
  • Black-box RL recipe is still a work in progress with config and interfaces subject to change
  • Requires a GPU cluster with substantial resources for the default recipe scale

README quick start

English · 简体中文

XYZ AI Lab ·
Technical Blog

AxisRL

AxisRL is an agentic RL post-training framework built on SGLang rollout, Megatron training, and real-world agent workflows.

AxisRL connects high-throughput rollout, large-scale training, weight synchronization, data movement, resource scheduling, and reproducible debugging inside one coherent framework. SGLang and Megatron remain the core serving and training engines; AxisRL handles the system layer around agentic post-training.

✨ Highlights

  • Built on SGLang for high-throughput rollout and Megatron for large-scale distributed training.
  • Used in agent RL workflows with 300+ turn trajectories and training runs at hundreds-of-billions-parameter scale.
  • Provides configurable policy optimization objectives including PPO, GRPO/GRPO2, GSPO, TOPR, TIS, and related variants.
  • Supports both white-box agent environments and black-box harness capture through an OpenAI-compatible proxy.
  • Reduces rollout/training idle time with partial rollout and a lightweight control plane.
  • Provides handle-based data movement, context packing, routing replay, mismatch analysis, and spike replay for rollout-trainer consistency.

🧭 Why AxisRL?

LLM post-training workloads are moving beyond single-turn question answering. In agentic RL, a model may interact with a long-lived environment, call tools, observe tool results, update context, and receive a reward only after several turns.

That changes the job of a post-training framework. It has to coordinate multi-turn rollout, environment state, tool calls, verifiers, reward collection, training sample construction, and weight synchronization. It also has to make training behavior observable, because small differences in tokenization, chat templates, logprobs, routing, packing, or weight sync can appear later as loss spikes, reward instability, or rollout-trainer mismatch.

AxisRL is designed for this setting: real agent workflows, SGLang rollout, Megatron training, and the system contracts between them.

🏗️ Architecture

At a high level, an AxisRL run follows this loop:

  1. Rollout actors execute task-specific agent workflows.
  2. SGLang workers serve model generation.
  3. Environments, tools, verifiers, or external harnesses produce interaction records and rewards.
  4. Megatron workers consume training sam

Description

AxisRL is an agentic RL post-training framework built on SGLang rollout, Megatron training, and real-world agent workflows.

Related repositories

Similar projects matched by category, topics, and programming language.

lopopolo
Featured
lopopolo GitHub avatar

harness-engineering

Harness Engineering is a methodology for improving coding agent outputs by carefully crafting the environment around them—providing curated context, tools, and executable constraints that encode an organization’s nonfunctional requirements and cumulative lessons.

AI & Machine LearningAI Agents
2,390
slvDev
Featured
slvDev GitHub avatar

esp32-ai

A 28.9 million parameter language model runs on an $8 ESP32-S3 microcontroller entirely on-device, generating simple stories at about 9.5 tokens per second.

AI & Machine LearningLarge Language Models
1,960
littledivy
Featured
littledivy GitHub avatar

mimic

mimic captures traffic from any iOS or web app and automatically generates a Python client library that lets you call the app's API like a regular library.

AI & Machine Learning
1,482