SkyRL Documentation

repository·main·Indexed 24 days ago

https://github.com/novasky-ai/skyrl

A 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.

Tokens
97.3K
Snippets
244
Records
431
Agent score
85%

What's inside SkyRL

  1. Overview of SkyRL components

    main

    SkyRL 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 from skyrl-train (modular training framework) and skyrl-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.
  2. Overview of SkyRL-Train features

    main

    SkyRL-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.
  3. Supported Environments in SkyRL-Gym

    main

    SkyRL-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.

  4. Explore the SkyRL API Reference

    main

    The SkyRL API is organized into several core modules for Reinforcement Learning (RL) with LLMs.

    SkyRL Core

    • Backends: Implementations for computation, including AbstractBackend, JaxBackend, and SkyRLTrainBackend.
    • Tinker Engine: The orchestration engine used for managing RL training processes.
    • TX Models: Configuration and model interfaces like ModelConfig, ModelForCausalLM, and CausalLMOutput.
    • Entrypoints: Execution interfaces such as BasePPOExp and EvalOnlyEntrypoint.
    • Configuration: Dataclasses used to define training parameters.
    • SFT (Supervised Fine-Tuning): Includes SFTTrainer, SFTConfig, and various samplers like StatefulSequentialSampler and DataMixingSampler.
    • Callbacks: Mechanisms for monitoring and controlling training via TrainingCallback, CallbackInput, TrainingControl, and CallbackHandler.
    • 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, and EnvStepOutput.
    • Tools: Tool management via ToolGroup and the tool decorator.
  5. Overview of SkyRLGymGenerator

    main

    The SkyRLGymGenerator is an implementation of the GeneratorInterface designed for use with SkyRL Gym environments. It is intended for use with fsdp and megatron backends.

    It utilizes a RemoteInferenceClient (acting as an LLM endpoint) to generate responses and returns GeneratorOutput (containing response_ids and loss_masks) to the training loop.

    Key implementation details:

    • Single-turn generation: Uses generate_batched() when config.generator.batched is True. This is suitable for tasks like math problems where a single response is required.
    • Multi-turn generation: Uses agent_loop() when config.generator.batched is False. This is used for interactive environments where the model performs multiple turns of interaction.
  6. What is On-Policy Distillation in SkyRL?

    main

    On-Policy Distillation is a training technique that combines on-policy Reinforcement Learning (RL) with the dense reward signals found in distillation.

    The Workflow:

    1. Collect on-policy samples from a student model.
    2. Use a teacher model to grade each token from those student samples.
    3. Update the student policy based on the teacher's grading.
    4. 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.

  7. What is Fully Async Training in SkyRL?

    main

    Fully 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:

    1. The trainer finishes a training step.
    2. It pauses ongoing trajectory generation.
    3. It updates the generator weights in-flight.
    4. 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.

  8. What is step-wise training and when to use it

    main

    Step-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?

    1. 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).
    2. 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_output config 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_size prompts, regardless of turn count.
  9. Overview of Clip-Cov and KL-Cov Policy Loss

    main

    Clip-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.
  10. Understand the Tinker API in SkyRL

    main

    SkyRL 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.

  11. How Trajectory lifecycle works

    main

    A Trajectory instance manages the execution flow for a single instance from a batch. It follows a three-step lifecycle:

    1. initialize_trajectory: Sets up the necessary runtime environment for the agent.
    2. generate_trajectory: Runs the agent loop to produce the final conversation and task results.
    3. evaluate_trajectory: Parses the final result and evaluates it against the task requirements.

    Results from both generate_trajectory and evaluate_trajectory are stored in the .result attribute of the trajectory object. Each trajectory initializes its own Agent instance to handle LLM interactions.