A wrapper class for Ray actors with initialization arguments. This class extends ClassWithInitArgs to provide additional functionality for co

A group of Ray workers that can be managed collectively. This class extends WorkerGroup to provide Ray-specific functionality for creating an

A metric aggregator for collecting and aggregating numeric values. This class accumulates numeric values (int, float, or scalar tensors) and

A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions. It contains a batch (TensorDict

The response from a tool execution.

HTTP-based adapter for SGLang engines. This adapter allows interaction with SGLang engines through HTTP requests instead of direct engine cal

A dispatcher that delegates to specific profilers based on config.tool. Supported tools: - nsys: NsightSystemsProfiler - npu: NPUProfiler

Worker profiler config. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. Args: d

Async token bucket for rate limiting with variable token consumption. The token bucket algorithm is a classic rate limiting technique that allows

Checkpoint engine manager to coordinate weight synchronization between trainer and rollout replicas. - ME: model engine, FSDP, MCore, VeOmni, exp

Configuration for critic model training. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.

FSDP optimizer configuration extending base OptimizerConfig. Args: optimizer (str): Optimizer class name (e.g., "AdamW", "AdamW8bit", "_A

Base optimizer configuration. Args: lr (float): learning rate. Must be specified. lr_warmup_steps_ratio (float): Warmup steps rat

Reward manager with rate limiting for API-based reward functions. This manager implements a sophisticated three-layer rate limiting system de

NPU profiler too; config.

RewardLoopManager run in single controller. This class will create reward loop workers and manage them.

RolloutSkip skips sequence generation during rollout by attempting to load previously dumped data. If no dumped data is found, it generates n

A unified tracking interface for logging experiment data to multiple backends. This class provides a centralized way to log experiment metrics, p

Manage FSDP checkpointing in SPMD training. - Saves/loads per-rank sharded model & optimizer states - Persists full lr_scheduler and RNG

Load and preprocess RLHF data from Parquet files. - Caches files locally. - Reads into a HuggingFace Dataset and tokenizes prompts.

Define a resource pool specification. Resource pool will be initialized first.

Minimal stand-in for HFModelConfig.

Configuration for FSDP (Fully Sharded Data Parallel). The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclas

Configuration for actor model training. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.

Asynchronous HTTP-based adapter for SGLang engines. This class inherits from HttpServerAdapter and adds async capabilities for non-blocking H

Uniform dispatch wrapper for batch data operations. All type-specific logic (isinstance checks) is centralized here so that callers (e.g. dec

Configuration for FSDP-based critic model training. The inheritance from CriticConfig provides all base critic configuration plus FSDP-specific s

Used to count mfu during training loop Example: flops_counter = FlopsCounter(config) flops_achieved, flops_promised = flops_

Configuration for Megatron parallelism. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.

Configuration for checkpoint engine to update weights from trainer to rollout

Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve.

Simulates a causal LM with embed_tokens, decoder layers, and lm_head.

Dataset for multi-turn conversations where each assistant response should be trained Args: data_files (str or list): Path(s) to Parq

Async broadcast operation with NCCL in separate thread. Args: rank (int): The rank of the current process. group_name (str): The

Async broadcast operation with HCCL in separate thread. Args: rank (int): The rank of the current process. group_name (str): The

Configuration for CPUMasterModel training. Supports any HuggingFace decoder-only model (Llama, Qwen, Mistral, Phi, etc.). Args: mode

FSDP DataParallel PPO Actor or Ref worker Args: config (ActorConfig): Actor config actor_module (nn.Module): Actor or ref module

Sharding manager to support data resharding when using FSDP + Ulysses

A custom linear layer implementation for the last layer of a model. This layer extends PyTorch's Linear module with functionality specifical

Configuration for Mcore module wrapper.

Mcore optimizer configuration extending base OptimizerConfig. Args: optimizer (str): Optimizer name; default is "adam". lr (float

The schema of a tool in OpenAI format.

Encapsulates a readable operation to remote agent. 1. send metadata to remote agent 2. wait until remote agent read complete. Args:

Torch profiler tool config.

A distributed worker that handles initialization and configuration for distributed training. This class manages worker initialization, configurat

Agent loop output.

A class to manage multiple OpenAI compatible LLM servers. This class provides - Load balance: least in-flight requests load balancing via glo

Receive model weights via bucketed IPC transfer over ZMQ. Receives weight tensors from BucketedWeightSender and passes each bucket to a

CPU master with explicit recompute - TRUE async pipeline. Supports any HuggingFace decoder-only model and VLM. Handles: - Uniform layers (Lla

A demo interaction for calculating the reward of gsm8k. - `start_interaction`: start a interaction instance for a trajectory. - `generate_res

Checkpoint manager for Megatron-LM distributed training. This class manages the saving and loading of model checkpoints in a Megatron-LM

Multi-head attention with Grouped Query Attention (GQA) support. Uses PyTorch's nn.Linear for projections to leverage optimized kernels.

Manages FP32 master weights, BF16 working copy, gradients, and optimizer state.

Precision debugger tool config (msprobe).

Root Mean Square Layer Normalization.

A custom collate_fn that handles batching of sequences. 1. for variable-length sequences, convert them into NestedTensors. 2. for fixed-l

TrainingWorker provides a Tinker-like API (https://thinkingmachines.ai/tinker/) as a RayWorkerGroup to a single controller. Currently, we onl

Single transformer decoder layer with pre-normalization. Architecture: x = x + Attention(RMSNorm(x)) x = x + MLP(RMSNorm(x))

This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on th

AdamW optimizer with proper bias correction and gradient clipping. Implements the algorithm from "Decoupled Weight Decay Regularization" (Lo

Teacher-specific async client used for distillation logprob computation.

Send model weights via bucketed IPC transfer over ZMQ. Packs weight tensors into a fixed-size communication buffer and sends them in buc

Universal SFT dataset with multi-turn, VLM, and thinking support. Features: - Multi-turn conversations (sharegpt format) - VLM image toke

Checkpoint handler handles the path, global_step of a checkpoint folder. Currently, it only works with a single model. We can expand it t

DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait for data so that asynchronous

Device check configuration: encapsulates device-specific validation rules

An env loop manages interactions between models and vectorized environments. It's designed for computationally intensive environments, such as rob

Configuration for FSDP actor models. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. Ar

FullyAsyncLLMServerManager supports resume generation on partial rollout, making rollout interruption invisible to the AgentLoop.

MLP with SwiGLU activation (used in LLaMA/Qwen models). Uses PyTorch's nn.Linear for better performance.

A configurable Multi-Layer Perceptron (MLP) module. It supports dynamic layer construction, multiple activation functions, and various we