The response from a tool execution.

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

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

A distributed profiler class for collecting performance metrics across multiple ranks. This profiler is designed to work in distributed training

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

Sharding manager to support data resharding when using FSDP + Ulysses

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

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

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

Configuration for Mcore module wrapper.

The data model for async rollout.

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

Agent loop manager that manages a group of agent loop workers.

Agent loop output.

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

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

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

Configuration for model merger operations. Args: operation (str): Operation type - 'merge' or 'test'. backend (str): Backend type

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

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

Internal agent loop output with padded sequences.

Adaptive KL controller described in the paper: https://arxiv.org/pdf/1909.08593.pdf

Compared to synchronize, this uses more memory because of the buffer but achieves better performance due to the overlapping. D2h and h2d copying a

A class to manage multiple OpenAI compatible LLM servers. This class provides - Load balance: least requests load balancing - Sticky sess

A class to manage multiple OpenAI-compatible LLM servers with *least in-flight requests* load balancing while preserving sticky sessions.

Configuration for efficient entropy kernel operations. Args: _backward (BackwardEnum): Backward computation method. Defaults to BackwardE

Context-manager that offloads/recovers tensors through an offload hander. The hook just offloads/recovers the tensor object to the handler throug

Configuration for FSDP checkpointing. Args: FSDP_version (int): Version of FSDP being used. world_size (int): Number of processes

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

Model merger for FSDP (Fully Sharded Data Parallel) checkpoints. This class handles the conversion of FSDP distributed checkpoints into Hugg

Sharding manager for FSDP models with vLLM inference engine integration. Manages parameter synchronization between FSDP training models and vLLM

Fixed KL controller.

LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla

LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev

A local logger that logs messages to the console. Args: print_to_console (bool): Whether to print to the console.

Model merger for Megatron-LM distributed checkpoints. This class handles the conversion of Megatron-LM distributed checkpoints into HuggingF

A sharding manager for Megatron-style training & inference with SGLang. This class manages the sharding of model parameters between training and

A sharding manager that bridges Megatron-LM training with vLLM inference. This class handles the parameter sharding and communication between:

A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying memory. It must have a unique type

A nested version of SimpleNamespace that recursively converts dictionaries to namespaces. This class allows for dot notation access to nested dic

The parsed schema of a tool in OpenAI format.

The schema of parameters in OpenAI format.

The parsed schema of a tool in OpenAI format.

The tool call in OpenAI format.

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

Multi-headed attention from 'Attention Is All You Need' paper

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig