Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

Get a rollout replica class by name.

Move model parameters, optimizer states, or both to the specified device. Args: device: Target device identifier.

Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split. Args: split_size (int): the size

Pop a subset of the DataProto via `batch_keys` and `meta_info_keys` Args: batch_keys (list, optional): a list of strings indicati

Return current device id based on the device type. Returns: device index

Return the corresponding torch attribute based on the device type string. Returns: module: The corresponding torch device namespace, or to

If distributed is initialized, print only on rank 0.

Update method for fixed KL controller (no-op). Args: current_kl (float): Current KL divergence value (unused). n_step

Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

Log GPU memory usage information. Args: head (str): A descriptive header for the memory usage log message. logger (logging.Logger

Extend self from other Dataset

Function that gets the torch.device based on the current machine. This currently only supports CPU, CUDA, NPU. Returns: device

_summary_ Log a message with rank information using a logger. This function logs the message only if `log_only_rank_0` is False or if the rank

Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split. Args: chunks (int): the number of

Convert an OmegaConf DictConfig to a dataclass. Args: config: The OmegaConf DictConfig or dict to convert. dataclass_type: T

Copy files/directories from HDFS to local cache with validation. Args: src (str): Source path - HDFS path (hdfs://...), local filesystem

Mark current step as complete.

Start a new step within current phase.

Create a tool instance. Args: instance_id: The instance id of the tool. Returns: The instance id of the tool

A func to sync embedding input with alltoall in sequence parallel gather sequence dimension and scatter head dim: e.g. seq_dim: 1, head_d

Concat a list of DataProto. The batch is concatenated among dim=0. The meta_info is merged, with special handling for metrics from different w

Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special tokens and clean up tokeniz

Create a huggingface pretrained tokenizer which correctness handles eos and pad tokens. Args: name (str): The name of the tokenizer.

Select a subset of the DataProto via batch_keys and meta_info_keys Args: batch_keys (list, optional): a list of strings indicatin

Create a DataProto from a dict of tensors. This assumes that 1. All the tensor in tensors have the same dim0 2. Only dim0 is the batch

Load configuration from a dictionary.

In megatron, the model and optimizer storage are: - bf16 parameter data chunked in model parallel group - fp32 grad chunked in model para

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

Aggregate the loss across global batch to ensure the loss is invariant to fsdp/megatron parallelism. NOTE: The returned loss has different b

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

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

Update current step progress.

More aggressive GPU memory cleanup function, tries to release PyTorch reserved but unallocated memory. Args: force_sync: Whether

Context manager for timing with platform markers. This utility function measures the execution time of code within its context, accumulates t

Register the dp_rank for a given mesh name. This function is meant to be called by the worker Args: mesh_name (str):

Get ulysses sequence parallel world size.

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

Sleep each rollout server.

Append values from new_data to lists in data. For each key in new_data, this function appends the corresponding value to a list stored under

Get ulysses sequence parallel process group.

Unified entry point for `index_first_axis` across CUDA and NPU backends. Dynamically dispatches to the appropriate device-specific implement

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

estimate the token and time usage of n iteration units Args: n: the number of iteration units

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

Repeat the batch data a specified number of times. Args: repeat_times (int): Number of times to repeat the data.

Combine base instruction with answer extraction instruction. Args: base_instruction: Task-specific output instruction

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

broadcast tensor from rank0 across mp_group

Calculate global rank with support for CP/EP parallelism

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:

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

Create a huggingface processor to process multimodal data. Args: name_or_path (str): The name of the processor. Returns: tra

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

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

Save the dataset to a JSONL file. Args: save_path: Path where the JSONL file should be saved.

Emit event to queue.

add token and time usage to this counter. Args: n_token: token cost time: time cost if add_

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

Create a Dataset instance from a list of samples. Args: sample_list: List of dictionaries containing the samples. Return

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

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