Function to create a new training engine instance based on the provided config. Args: key: A configuration object contain

Performs a single batch search for multiple queries (original search tool behavior). Args: url: The URL of API. query: query

Release the tool instance. Args: instance_id: The instance id of the tool.

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

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

Reorder the data on single controller such that each dp rank gets similar total tokens

model: AutoModelForCausalLM

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

Validate mutually exclusive micro batch size configuration options. Ensures that users don't set both deprecated micro_batch_size and

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

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

Pad a DataProto to size divisible by size_divisor Args: size_divisor (int): size divisor Returns: data: (DataProto): the pad

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

Get the rank of this worker in the distributed setup.

Sleep all rollout replica instances.

Load a simple tokenizer for testing.

Unpad the data proto with pad_size. i.e. `data[:-pad_size]`

Validate actor configuration with runtime parameters.

Initialize SearchTool with configuration and schema. Args: config: Configuration dictionary containing tool settings

given TP,DP,PP rank to get the global rank.

given TP,DP,PP rank to get the global rank.

Apply KL penalty to the token-level rewards. This function computes the KL divergence between the reference policy and current policy, then a

Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto

Compute the attention mask for the response part of the sequence. This function extracts the portion of the attention mask that corresponds to th

Compute reward for a batch of data. Args: data: DataProto object containing the input data. reward_fn: Reward function to com

Deprecated. Please use copy_to_local instead.

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

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

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

Extract the last LaTeX boxed expression from a string. Args: string: Input string containing LaTeX code Returns: The last bo

Normalize a final answer to a quantitative reasoning question. Args: final_answer: The answer string to normalize Returns: N

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

Process validation metrics into a structured format with statistical analysis. This function organizes validation metrics by data source and

Remove the LaTeX boxed command from a string. Args: s: String with format "\\boxed{content}" Returns: The content inside the

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

Resume GPU occupation.

Update the module config with the override_config_kwargs. Args: module_config: The module config from Huggingface Transformers. ov

broadcast tensor from rank0 across mp_group

broadcast tensor across mp_group

broadcast tensor from rank0 across mp_group

Dump rollout/validation samples as JSONL.

fetch tensor

fetch tensor

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

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

Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the parameters from rank 0 to all ot

Generate text using the SGLang server. Args: prompt (Optional[str], optional): Text prompt for generation. Defaults to None.

Load and return a custom reward function from external file. Dynamically imports a reward function from a specified file path and wraps it wi

Get the number of gpus in this cluster.

Convert huggingface PretrainedConfig to mcore TransformerConfig. Args: hf_config: The huggingface PretrainedConfig. dtype: The dt

Check if a path is a non-local (HDFS) path. Args: path (str): The path to check. Returns: bool: True if the path is an HDFS