Convert Hydra configs and other special types to regular Python types.

Create agent loop manager.

A func to sync attention result with alltoall in sequence parallel gather head dimension and scatter seq dim: e.g. seq_dim: 1, head_dim:

Get FSDP wrap policy for the module. Args: module: The module to get wrap policy for config: Configuration for wrap policy

Return the list of RouterReplay instances corresponding to the current micro-batch and local (pp_rank, vp_stage) layer range.

Converts peft model parameter name to base parameter name For example, base_model.model.model.embed_tokens.weight -> model.embed_toke

Preprocess packed sequences CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets s

Simple buffer registration (for standalone use).

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

Restore a batch from dynamic batching. Args: data (torch.Tensor): The input data. batch_idx_list (List[List[int]]): The list

A context manager to add attributes to a trace for the configured backend. Args: sample_index: Sample index for the trace. step:

Move model parameters, optimizer states, or both to the specified device. Note that this function executes irrespective of offload co

Move FSDP model and/or optimizer to CPU or GPU with offload support. Note that this function executes irrespective of offload config.

Get the GPU layer template for a given layer index and buffer slot.

Load CPU layer params to GPU buffer asynchronously. Uses pre-flattened pinned buffers for efficient DMA transfer.

Unflatten GPU buffer to the appropriate layer template parameters.

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

Compute advantage estimates for policy optimization. This function computes advantage estimates using various estimators like GAE, GRPO, REINFORC

Compute logits given a batch of data. Args: data (DataProto): a batch of data represented by DataProto. It must contain key ```in

Create Ray resource pools for distributed training. Initializes resource pools based on the resource pool specification, with each po

Remove old checkpoints to make room for a new one, keeping a safety buffer. With max_ckpt_to_keep=1, this does nothing - we keep the

Extract and process multi-modal inputs from a batch. Args: batch_data (list[dict[str, torch.Tensor]]): The batch containing potentia

Flush the cache of the server. This method repeatedly attempts to flush the server cache until successful. The flush operation will n

Get a generator that yields per-tensor parameters and optional peft config. Returns: Generator[tuple[str, torch.Tensor]]

Initialize interactions from configuration file. Args: interaction_config_file: Path to the interaction configuration file. Returns:

r"""Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch dataset. See https://pytorch.org/t

Given a config, do we need critic.

Unified entry point for `pad_input` across CUDA and NPU backends. Dynamically dispatches to the appropriate device-specific implementation:

Split a batch into micro-batches by total token count, with optional DP sync and padding. Args: batch (TensorDict): must include "at

Register a successfully saved checkpoint and enforce retention limit. Adds the new checkpoint path to tracking and removes excess ol

Run the environment interaction loop. This method orchestrates a pipelined process: 1. Resets environments to specified initi

Save parameter metadata for rebuild.

Configure CUDA memory allocator expandable segments setting. Expandable segments can help avoid out-of-memory (OOM) errors by allowing the me

Pad and slice input_ids to be divisible by sp_size Pad position_ids to be divisible by sp_size. Note both input_ids_rmpad and position_i

Make an async HTTP request with retry logic and consistent error handling. Args: endpoint (str): The API endpoint to call (withou

Sync CPU master params to GPU modules (call after optimizer step).

Computes various metrics from a batch of data for PPO training. This function calculates metrics related to scores, rewards, advantages, ret

This function should return a class instance that delegates the calls to every cls in cls_dict

Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken. Args: batch_seqlens (Lis

Finalizes the interaction session and releases any associated state or resources. Simulates: release state

Split input batch and dispatch to env loop workers. Args: prompts (DataProto): Input batch. Returns: DataPro

Get the checkpoint engine class. Args: backend: The backend of the checkpoint engine. Returns: The checkpoin

Build the model.

Get the software version and CANN toolkit version for NPU devices. Returns: tuple[str, str]: A tuple of (software_version, cann_version)

Build the inverse of an index mapping. Args: idx_map (Sequence[int]): Sequence where idx_map[i] = j. Returns: List[int]

Get the reward manager class with a given name. Args: name: `(str)` The name of the reward manager. Returns: `(t

Calculates partitions of indices from seqlen_list such that the sum of sequence lengths in each partition is balanced. Uses the Karmarkar-Kar

Perform inference on a batch of data. Args: data: The input data for inference, typically containing tensors and metadat

Check if distillation is enabled based on the provided configuration.

Creates a batch generator suitable for Megatron pipeline parallelism, handling virtual pipeline parallelism (VPP). If VPP is used (vpp_s

Puts a list of data into the Ray object store in parallel using a thread pool. Args: data_list (List[Any]): A list of Python objects

Prepare a batch for dynamic batching. Args: data (DataProto): The input data. max_token_len (int): The maximum token length

Rebuild tensor from shared memory.

Context manager for basic timing without NVTX markers. This utility function measures the execution time of code within its context and accum

Perform a single optimization step with bias correction.

Update the weights of the rollout model. Args: weights: A generator that yields the name of the weight tensor and the tensor itse

Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GP

Validate actor configuration parameters.

broadcast tensor across mp_group

Create a base TransformerConfig with common parameters across different model architectures. TODO: (ycl) use dataclass or converter config?

Build a minimal TensorDict accepted by prepare_micro_batches. Args: seq_lens: Effective sequence length for each sample. force_gr

Query the dispatch info for a given mesh name. Args: mesh_name (str): Name of the mesh to query dispatch info for

Load layer weights from shared-memory flat to GPU buffer. Direct shared-memory → GPU copy (non_blocking). Not as fast as pinned → GPU DMA, bu

Unflatten GPU buffer to the appropriate layer template parameters.

Apply monkey patch to the models for ulysses sequence parallel, fused kernel, tiled MLP and prefix grouper. In the end of this function forw

View buffer as a pinned tensor (no copy).

Automatically configure device name for different accelerators. For example, on Ascend NPU, this function defaults the trainer device to "npu"

Calculate approximate computational workload for transformer attention. Estimates FLOPs for dense transformer blocks based on sequence length usi