Tracker file rescords the latest chckpoint during training to restart from.

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

Calculate the pipeline sharding configuration for Megatron-LM. Args: layer_num: Total number of layers in the model. pp_size: Num

Get the enum for the reduction method of cross entropy.

Get the generation prompt ids for rollout engine. Because rollout engine(SGLang) requires the ids to be a list, we need to convert t

Get the weight converter for given model architecture.

Get the last learning rate from the optimizer parameter scheduler.

Get the optimizer parameter scheduler for Megatron.

Group tensors into buckets based on a specified size in megabytes. Args: iterable: An iterator of tuples containing tensor names and

Generate and return an NPU profiler object. Args: option (DictConfig): The options to control npu profiler. role (str

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

The answer will be considered correct if: (a) it normalizes to the same string as the ground truth answer OR (b) sympy can simplify t

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

Check if the given string is an IPv6 address Args: ip_str: The IP address string to check Returns: bool: Returns True i

Check if padding is enabled for the DataProto. Returns: bool: True if padding is enabled, False otherwise.

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

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

Load merged state_dict to sharded Megatron module in training.

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

Make minibatch iterator for updating the actor Args: data (DataProto): a DataProto containing keys ``input_ids``

End a mark range in the profiler. Args: range_id (str): The id of the mark range to end.

Start a mark range in the profiler. Args: message (str, optional): The message to be displayed in the profiler. Defaults to N

Exact match of math if and only if: 1. numerical equal: both can convert to float and are equal 2. symbolic equal: both can convert to sy

Convert a MCP Tool to an OpenAI ChatCompletionTool.

Generate an MD5 hash of a path string. This function is used to create unique identifiers for paths, typically for creating cache directories

Transform the model name in each model_chunk in each pp stage into the name in inference engine

Offload.

On group commit forward.

Zero out gradients of all parameters before starting a new backward pass.

pad a 2D list (e.g. responses, logprobs) to a 2D tensor.

Split a repo_id like "owner/name" into (owner, name).

_summary_ For fused kernels, the output is a dictionary with keys like 'log_probs', 'entropy', etc. This function post-processes each tensor i

Converts a video dict into a [n_frames, 3, H, W] tensor Add video sample FPS in a future MR

Recover left padding from result return result

Reload.

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

Remove left padding from input_ids, attention_mask and position_ids return new_input_ids, new_attention_mask, new_position_ids

Note that this operation is in-place

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

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

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

A context manager to add attributes to a trace for the configured backend.

Rotates half the hidden dims of the input.

Run agent loop to interact with LLM server and environment. Args: sampling_params (Dict[str, Any]): LLM sampling params.

Set ulysses sequence parallel process group.

Shard or partition data for distributed training or parallel execution. Args: data: Data structure to be sharded across

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

Offload model weights and discard kv cache.

Sleep all rollout server instances.

Split the elements in a tuple/interval, while handling well-formatted commas in large numbers

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

Reconstruct or gather sharded data back to a unified format. Args: data: Sharded data structure to reconstruct.

Update the policy with an iterator of DataProto Args: data (DataProto): an iterator over the DataProto that returns by

Validate critic configuration with runtime parameters. Args: n_gpus: Total number of GPUs available train_batch_size:

TODO(zhangchi.usc1992): We may change the implementation later