SkyRL Documentation
repository·main·Indexed 24 days ago
https://github.com/novasky-ai/skyrlA modular, full-stack Reinforcement Learning library designed for training and evaluating Large Language Models (LLMs) on real-world, long-horizon tasks. It provides a unified API for training and inference, featuring support for PPO-style RL loops, deterministic session-based routing, and integrations with Harbor and Modal.
What's inside SkyRL
- The Tinker integration connects SkyRL Agent with Tinker's RL training framework. It provides a pipeline for collecting agent trajectories, performing Reinforcement Learning (RL) training using LoRA (Low-Rank Adaptation) with techniques like PPO or importance sampling, tracking experiments via Weights & Biases, and managing model checkpoints.
Overview of SkyRL components
mainSkyRL is a modular full-stack Reinforcement Learning (RL) library for Large Language Models (LLMs) composed of three primary packages:
skyrl: A unified library for RL on local hardware. It supports the Tinker API and integrates previous work fromskyrl-train(modular training framework) andskyrl-tx(cross-platform backend for the Tinker API).skyrl-agent: An agent layer designed for training long-horizon, real-world agents and optimizing multi-turn tool-use LLM pipelines.skyrl-gym: A collection of tool-use tasks (math, coding, search, and SQL) implemented using the Gymnasium API.
Overview of SkyRL-Train features
mainSkyRL-Train is a modular RL framework designed for post-training LLMs. It supports high-performance training and flexible execution plans.
Core Capabilities:
- Algorithms: PPO, GRPO, RLOO, REINFORCE, GSPO, CISPO, SAPO.
- Training Backends: FSDP and Megatron (with 5D Parallelism support for MoE).
- Inference Backends: vLLM, SGLang, or any OpenAI-compatible endpoint supporting weight sync.
- Parallelism & Scaling: Ulysses sequence parallelism, colocated or disaggregated training/generation, and weight sync via NCCL, gloo, or checkpoint-and-load.
- Execution Modes: Synchronous RL, async one-off pipelining, or fully async RL with in-flight weight updates.
- Optimizations: Sequence packing and Flash Attention 2.
Supported Environments in SkyRL-Gym
mainSkyRL-Gym provides several types of environments designed for Reinforcement Learning with LLMs, ranging from single-domain tasks to complex multi-tool scenarios.
Basic Environments
These are single-domain environments focused on specific capabilities:
- Search: Web search and information retrieval.
- SQL: Database querying and management.
- Math: Mathematical reasoning and problem-solving.
- Simple Code (LCB): Basic coding tasks using the LCB framework.
Multi-Tool Environment
In this environment, the model must decide which tool to invoke at each step. This is used to showcase complex workflows, such as combining Search and Python Code execution. Note that these environments use different tool parsing logic compared to basic environments.
Mix Dataset Environment
These environments handle multi-domain datasets. They allow for switching between different environment initializations. While you can reuse the environment object, the agent loop logic must be adjusted to handle the domain switching.
Explore the SkyRL API Reference
mainThe SkyRL API is organized into several core modules for Reinforcement Learning (RL) with LLMs.
SkyRL Core
- Backends: Implementations for computation, including
AbstractBackend,JaxBackend, andSkyRLTrainBackend. - Tinker Engine: The orchestration engine used for managing RL training processes.
- TX Models: Configuration and model interfaces like
ModelConfig,ModelForCausalLM, andCausalLMOutput. - Entrypoints: Execution interfaces such as
BasePPOExpandEvalOnlyEntrypoint. - Configuration: Dataclasses used to define training parameters.
- SFT (Supervised Fine-Tuning): Includes
SFTTrainer,SFTConfig, and various samplers likeStatefulSequentialSamplerandDataMixingSampler. - Callbacks: Mechanisms for monitoring and controlling training via
TrainingCallback,CallbackInput,TrainingControl, andCallbackHandler. - SkyRL-Train Backend: Low-level training components including data interfaces (
TensorBatch,TrainingInput), generators (GeneratorInterface), and trainers (RayPPOTrainer,Worker).
SkyRL-Gym
- Environment: Interfaces for RL environments, including
Env,BaseTextEnv, andEnvStepOutput. - Tools: Tool management via
ToolGroupand thetooldecorator.
- Backends: Implementations for computation, including
Overview of SkyRLGymGenerator
mainThe
SkyRLGymGeneratoris an implementation of theGeneratorInterfacedesigned for use with SkyRL Gym environments. It is intended for use withfsdpandmegatronbackends.It utilizes a
RemoteInferenceClient(acting as an LLM endpoint) to generate responses and returnsGeneratorOutput(containingresponse_idsandloss_masks) to the training loop.Key implementation details:
- Single-turn generation: Uses
generate_batched()whenconfig.generator.batchedisTrue. This is suitable for tasks like math problems where a single response is required. - Multi-turn generation: Uses
agent_loop()whenconfig.generator.batchedisFalse. This is used for interactive environments where the model performs multiple turns of interaction.
- Single-turn generation: Uses
What is On-Policy Distillation in SkyRL?
mainOn-Policy Distillation is a training technique that combines on-policy Reinforcement Learning (RL) with the dense reward signals found in distillation.
The Workflow:
- Collect on-policy samples from a student model.
- Use a teacher model to grade each token from those student samples.
- Update the student policy based on the teacher's grading.
- Repeat the process.
In SkyRL, this is implemented by modifying the standard training loop to replace the reference model with a teacher model and adjusting the reward/advantage computation logic to utilize reverse KL loss.
What is Fully Async Training in SkyRL?
mainFully async training (also known as in-flight weight update or multi-turn partial rollout) is a training paradigm designed to solve the 'straggler issue' in long-horizon tasks.
Unlike synchronous training (where training and generation are colocated) or one-step off-policy training (where generation and training are split but stalled by slow generation batches), fully async training allows the trainer to update weights without waiting for all generation workers to finish.
The Workflow:
- The trainer finishes a training step.
- It pauses ongoing trajectory generation.
- It updates the generator weights in-flight.
- It resumes generation using the new weights.
This allows a single trajectory to be generated by multiple model versions, maximizing throughput by preventing slow generation workers from stalling the entire pipeline.
What is step-wise training and when to use it
mainStep-wise training decomposes a multi-turn trajectory into $N$ separate training samples (one per LLM turn) instead of a single
(prompt, response)pair.Why use it?
- Avoids Re-tokenization Drift: Standard re-tokenization of a full conversation string can result in different token IDs than what the model actually generated (due to BPE non-uniqueness or chat template changes). Step-wise training uses the exact token IDs and logprobs from the inference engine, which is required for accurate rollout correction like TIS (Truncated Importance Sampling).
- Handles Complex Context Management: It allows for non-strictly-appending chat histories. If an agent harness summarizes context, strips thinking tokens, or resets windows between turns, step-wise training can represent these discontinuities because each turn is treated as an independent sample.
Trade-offs
- Complexity: Training time grows as $O(T^2)$ vs $O(T)$ because each turn has a growing prompt prefix. You can mitigate this using the
generator.merge_stepwise_outputconfig flag if the history is linearly appending. - Batching: A batch of $T$ trajectories with $M$ turns per trajectory produces $T \times M$ training samples. Each mini-batch contains sequences for exactly
policy_mini_batch_sizeprompts, regardless of turn count.
Overview of Clip-Cov and KL-Cov Policy Loss
mainClip-Cov and KL-Cov are policy loss methods designed to improve training stability in LLM reinforcement learning by using covariance-based token selection.
- Clip-Cov: Integrates standard PPO clipping with a covariance-based correction masking mechanism.
- KL-Cov: Applies KL regularization specifically to tokens that are selected based on their covariance values.
Understand the Tinker API in SkyRL
mainSkyRL implements the Tinker API, a minimal interface designed for post-training Large Language Models (LLMs). The core philosophy of the Tinker API is the separation of algorithm logic from infrastructure logic.
By writing training scripts against the Tinker API, you can implement complex algorithmic logic (like RL or SFT) while SkyRL handles the underlying infrastructure concerns such as worker management, batching, and weight transfers. This allows you to run scripts written for the Tinker API on SkyRL's high-performance backends (FSDP, Megatron, vLLM) with zero code changes.
How Trajectory lifecycle works
mainA
Trajectoryinstance manages the execution flow for a single instance from a batch. It follows a three-step lifecycle:initialize_trajectory: Sets up the necessary runtime environment for the agent.generate_trajectory: Runs the agent loop to produce the final conversation and task results.evaluate_trajectory: Parses the final result and evaluates it against the task requirements.
Results from both
generate_trajectoryandevaluate_trajectoryare stored in the.resultattribute of the trajectory object. Each trajectory initializes its ownAgentinstance to handle LLM interactions.