Return a dictionary containing name to a shape and dtype.

Create a Warmup-Stable-Decay learning rate scheduler. The schedule follows three phases: 1. Warmup: Learning rate increases linearly fro

Compute per-group mean/std/count in pure PyTorch. mean_g = sum / count std_g = sqrt( max( (sum2 - sum^2/count) / max(count-1, 1), eps )

Groupid reset.

Identify which fields in an image dataset contain image, input, and output. Uses LLM to analyze field names and determine the mappin

Index a tensor dict with a tensor of indices.

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

Initialize class state shared across all instances.

Init colocated rollout server, rollout engine and hybrid engine colocated in same ray placement group but in separate processes. Args

Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process. Args: worker_group: RayWorke

Initialize resource pool manager.

Initialize search execution pool.

Initialize visual execution pool.

Initialize distributed training workers using Ray backend. Creates: 1. Ray resource pools from configuration 2. Worker groups

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

Initialize system prompt tokens for chat templates that support them. Args: tokenizer: The tokenizer with a chat template **

Check if the solution is correct according to Minerva criteria. Args: solution_str: The solution string to check gt: The ground t

Check if the prediction is correct using strict boxed answer criteria. Args: pred: The prediction string gt: The ground truth ans

Return True if the current router_replay_action is REPLAY_BACKWARD for the local router instances. This inspects the first local RouterReplay

Check the availability of NPU

check if the package version is greater than or equal to the minimum version

Compute KL divergence given logprob and ref_logprob. Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104

Launch an SGLang HTTP server process and wait for it to be ready. This function starts a new process running an SGLang HTTP server, then waits

Load all PDF documents from configured corpus paths. Returns: List of PDF file paths

DEPRECATED. Directly use `load_extern_object` instead.

Load optimizer parameters back to GPU. Handles ChainedOptimizer. Args: optimizers: Optimizer or ChainedOptimizer instance.

Load a module from a path. Args: module_path (str): The path to the module. Either - `pkg_path`, e.g.,

Load merged state_dict to sharded Megatron module in training.

Load training configuration from a YAML file. Args: yaml_path: Path to YAML configuration file Returns: Appropriate con

Log validation generation to clearml as table

Log validation generation to mlflow as artifacts

Log samples to swanlab as text

Log samples to tensorboard as text

Calculate and log metrics related to sequence length imbalance before and after partitioning. Args: seqlen_list (List[int]): A list

A memory efficient implementation of logprobs_from_logits

Main entry point for PPO training with Hydra configuration management. Args: config_dict: Hydra configuration dictionary containing train

Creates a pre-wrap hook that replace the output layer with a value head. Args: hidden_size (int): The hidden size of the model's transfor

Compute variance of tensor with masked values.

Merge recorded router top-k indices across sequence-parallel ranks for all router instances, then pack/unpack them to align with the original

Given a role worker mapping, do we need reward model.

Offload optimizer parameters to CPU. Supports both Megatron optimizers and `ChainedOptimizer`, which wraps a list of underlying optimizers.

On group commit backward.

Perform an optimization step using the optimizer.

Zero the gradients of the optimizer.

Zero gradients and enforce FSDP grad-clipping logic.

pad the tokens such that the total length is a multiple of sp world size Args: unpad_tokens: (total_nnz, ...). Tokens after removing padd