Check correctness of code generation with a global timeout. The global timeout is to catch some extreme/rare cases not handled by the timeouts

Unified interface for computing IS weights and rejection masks. This function combines IS weight calculation (truncated) and rejection sampling (

Create a DataProto from a dict of tensors and non_tensors

Spawn to a dictionary of worker groups, each with a subset of method with prefix. Args: prefix_set: Set of prefixes to create wor

Start profiling for the current rank in the current training step.

Stop profiling for the current rank in the current training step.

[Deprecated] convert config Args: hf_config (PretrainedConfig): _description_ megatron_config (_type_): _description_ Return

Create a dataset. Arguments: data_paths: List of paths to data files. data_config: The data config. tokenizer (Tokenizer)

Gather a tensor across a process group and optionally unpad its padded elements. Args: x (Tensor): Input tensor to gather. g

Get ulysses sequence parallel world size.

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

Enable or disable expandable segments for cuda. Args: enable (bool): Whether to enable expandable segments. Used to avoid OOM.

This is used to do heavy initialization work that should shared across all instances. It's only called once. Args: config (DictCo

Compute the mean of `values` over elements selected by `mask`. Args: values (Tensor): Input tensor. mask (Tensor): Boolean o

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 a HTTP request with retry logic and consistent error handling. Args: endpoint (str): The API endpoint to call (without leadi

Create a sampler for the dataset. Arguments: data_config: The data config. dataset (Dataset): The dataset. Returns:

Execute a method on all workers synchronously. Args: method_name: Name of the method to execute *args: Positional arg

Helper to create DataProto for testing.

Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. Args: text (`str`, `List[str]` or

Return nccl backend type based on the device type. Returns: nccl backend type string.

Load a external data type based on the file path and type name

_summary_ Thread-safe directory creation function that ensures the directory is created even if multiple processes attempt to create it simult

Compute per-token log-probabilities for the given labels. Uses a Flash-Attention–based cross-entropy (if available) for efficient backward,

Determine if checkpoint should be saved based on capacity esi expiration. Args: max_steps_duration: Max estimated time (seconds) req

Convert this DataProto to TensorDict. Note that this requires tensordict version at least 0.10 Returns:

Split input batch and dispatch to agent loop workers. Args: prompts (BatchMeta): Input batch. Returns: Batch

Batch generate sequences in sync mode. Args: prompts: The input prompts. Returns: The output sequences.

Get ulysses sequence parallel process group.

Given a role worker mapping, do we need ref policy.

Context manager for temporarily setting an environment variable. This context manager ensures that environment variables are properly set and res

Copied from ManiSkill! Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu is not enabled then we use CPU, otherwise G

Computes throughput metrics for PPO training. This function calculates performance metrics related to token processing speed, including

Create a random mask given input_ids. Support left padding and right padding. Process: - Sample valid token length - Sample left_padding l

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

Whether the current rank is the first rank in model parallel group that contains model outputs

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

Apply monkey patch to the models for ulysses sequence parallel and fused kernel. In the end of this function forward function of the model i

The interface for calling tools. Each tool needs to implement this function, which is the workflow of the tool. Args: pa

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

Computes timing metrics for different processing stages in PPO training. This function calculates both raw timing metrics (in seconds) and p

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

broadcast tensor from rank0 across mp_group

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

Calculate global rank with support for CP/EP parallelism

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

Return the most recent checkpoint directory based on a tracker file. Args: path (str): Base directory containing the checkpoint trac

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

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

Override init_model to use FlowRLActor instead of DataParallelPPOActor.

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