Get the actual results by calling get() method on gen_batch_output Returns: tuple: (epoch, batch, gen_batch_result, futu

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 the corresponding torch attribute based on the device type string. Returns: module: The corresponding torch device namespace, or to

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

move the batch to device Args: device (torch.device, str): torch device Returns: DataProto: the current Data

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

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

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

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

Get the value associated with the given key. If the key does not exist, return the default value. Args: key (str): The attribute

If distributed is initialized, print only on rank 0.

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

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 tool instance. Args: instance_id: The instance id of the tool. Returns: The instance id of the tool

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

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

Union with another DataProto. Union batch and meta_info separately. Throw an error if - there are conflict keys in batch and they are

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

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

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

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

Sleep each rollout server.

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

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

Context manager for basic timing without NVTX markers. This utility function measures the execution time of code within its context and accum

_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

Execute the main PPO training workflow. This method sets up the distributed training environment, initializes workers, datasets, and

Split input batch and dispatch to env loop workers. Args: prompts (DataProto): Input batch. Returns: DataPro

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

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

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

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

Acquire tokens from the bucket, waiting if necessary. This method will block (using asyncio.sleep) until sufficient tokens are availa

r"""Works like shutil.copy() for file, and shutil.copytree for dir, and supports hdfs. Copy data and mode bits ("cp src dst"). Return the file's

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

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

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

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

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

Reduces a dictionary of metric lists by computing the mean, max, or min of each list. The reduce operation is determined by the key name:

r"""Works like os.makedirs() but supports hdfs. Super-mkdir; create a leaf directory and all intermediate ones. Works like mkdir, except tha

pad a 2D tensors (e.g. responses, logprobs) in the last dim to max_seq_length. input shape: [bs, seq_length] output shape: [bs, max_seq_l

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 the resource pool of the worker_cls

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

Verify the parameters of the function call

Load and initialize a reward manager based on the configuration. Args: config: PPO trainer configuration object containing reward_mo

Note that this operation is in-place

The interface of calling LLM for the agent. We prepend the system_message of this agent to the messages, and call LLM. Args: