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

Applies a monkey patch to the forward method of a given model class to enable Ulysses sequence parallelism input slicing.

Release GPU occupation temporarily.

Replace LoRA parameter keys with base layer equivalents. Transforms LoRA parameter names to their corresponding base layer names for proper w

Resume rollout weights or kv cache in GPU memory. Args: tags: weights or kv_cache.

Save checkpoint using FSDPCheckpointManager with improved tracking

Shutdown the HTTP server and clean up resources. This method performs the following cleanup operations: 1. Unregisters the worker fro

Split a RayResourcePool into multiple SubRayResourcePool. resouce_pool can also be a SubRayResourcePool (have been splited) for multiple-time

Decorator to add a timeout to a function. Args: seconds: The timeout duration in seconds. use_signals: (Deprecated) This is

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

Converts a string in a sequence of tokens. Args: text (`str`): The sequence to be encoded. a

Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GP

broadcast tensor across mp_group

broadcast tensor across mp_group

Acquire a token from the bucket.

Convert arbitrary group labels to a contiguous 1-D torch.long tensor (0..G-1). Args: index: Any iterable of labels or tensor/ndarray

Build an optimizer based on the configuration. Dynamically imports and instantiates an optimizer class from the specified module. Args:

calculate rollout vs actor logprobs diff, for debugging purpose Args: data: DataProto the data batch to calculate

Check and disable incompatible configurations for older Megatron version. Args: original_config (dict): The original model configura

collect lora params or full params if base model is not ready in vllm work with if isinstance(self.module._fsdp_wrapped_module, PeftModel)

Compute off-policy diagnostic metrics (helper function). This helper function operates on raw tensors and is used internally by: - compute_ro

Converts a sequence of tokens in a single string.

Compute the score for a given solution based on the data source. Args: data_source (str): The source dataset identifier which determines

Generates a memory snapshot and saves it as a pickle file in a specified directory. The files are organized by timestamp in subdirect

Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation groups asynchronously. This mea

Fallback implementation for older transformers versions. Returns False to disable features that require this function.

Get the advantage estimator function with a given name. Args: name_or_enum: `(str)` or `(AdvantageEstimator)` The name or enu

Create a constant LR schedule with a linear warmup phase. Args: optimizer (Optimizer): Wrapped optimizer. num_warmup_steps (

Create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer

Get all registered dispatch and collect dp_ranks. Returns: dict[str, int]: A dictionary mapping mesh names to the

Build the model.

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

A filter function to return the PPO Ray runtime environment. To avoid repeat of some environment variables that are already set.

Get the resource pool of the worker_cls

Get the resource pool of the worker_cls

Get the rollout class by name. Args: rollout_name: The name of the rollout. mode: The mode of the rollout, sync: spmd mode, async

Init standalone rollout server, create new resource pool for this rollout.

Init resource pool and worker group

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 weights for mcore GPT model.

Given a config, do we need critic

Generate data using the provided data generation strategy. Note: This method is intended to change the dataset after each training ba

Get chunk data of this tp rank since we do all gather in preprocess.

Postprocess packed sequences

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

Set the dispatch and collect dp_ranks for all registered meshes. Args: mesh_name (str): Mesh name to set dispatch and collect dp_

Initialize distributed environment if not already initialized.

Validate an OmegaConf DictConfig. Args: config (DictConfig): The OmegaConf DictConfig to validate. use_reference_policy (bool): i

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group

Build multimodal embeddings and attention mask

Compute loss with optional sequence parallelism and remove padding features

Computes information about prompts and responses from a batch. This is an internal helper function that extracts masks and lengths for promp

Creates the train and validation dataloaders.

Dynamically import attention functions based on available hardware.