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

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

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

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_

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

From Multi-Head Attention module https://github.com/jadore801120/attention-is-all-you-need-pytorch

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

A wrapper of textEnv for models. Returns valid actions at each step of the game.

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

Memory manager: responsible for storing & fetching per‑environment history records.

Linear interpolation between initial_p and final_p over schedule_timesteps. After this many timesteps pass final_p is returned. :para

Note that this trainer runs on the driver process on a single CPU/GPU node.

EnvironmentManager for SearchEnv.

AsyncLLMServerManager manage a group of vllm instances, i.e AsyncvLLMServer.

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

A basic LSTM cell.

A checkpoint manager that saves and loads - model - optimizer - lr_scheduler - extra_states in a SPMD way. We save -

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

Dataset for multi-turn conversations where each assistant response should be trained

The schema of parameters in OpenAI format.

The schema of a parameter in OpenAI format.

The schema of a function in OpenAI format.

The schema of a tool in OpenAI format.

This is an in-memory SFTDataset Arguments: config (OmegaConf): the data config

Track a series of values and provide access to smoothed values over a window or the global series average.

inputs: p: batch x time x inp_p mask_p: batch x time q: batch x time x inp_q mask_q:

A test class to be imported by load_extern_type

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

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

Interface for Textworld Env

Interface for Embodied (THOR) environment

A Ray-based distributed wrapper for AppWorld. - Creates multiple Ray actors, each holding a separate AppWorld instance. - Implements Gym-

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

The data model for async rollout.

Blackjack is a card game where the goal is to beat the dealer by obtaining cards that sum to closer to 21 (without going over 21) than the de