Apply NVFP4 patches to support dynamic weight updates. Call before model loading.

Apply ModelOpt fake quantization to a list of Megatron module chunks.

Apply rotary embeddings to query and key tensors. Args: q: Query tensor [B, T, num_heads, head_dim] k: Key t

Reroute to the appropriate apply_rotary_pos_emb function depending on bshd (conventional) / thd (packed seq) format In Qwen2-VL, the sha

Applies the monkey patch for MoE Router Replay functionality. This patch dynamically adds the 'enable_routing_replay' attribute to Transforme

Apply TiledMLP monkey patch based on model_type. This function MUST be called BEFORE model instantiation to take effect. It patches the MLP c

Assemble gen_batch_output from RolloutSample objects Assembles batches from RolloutSample objects, similar to the _post_generate_batch logic

Async copy parameters from CPU to GPU buffer

Send weights to the receiver. Accepts a sync generator or async iterator. Args: weights: Generator or async iterator yie

Backup base model weights to CPU with LoRA temporarily disabled. This function temporarily disables LoRA adapters, backs up the clean base model

Start the read operation.

Broadcast all tensors in a dictionary from source rank to all ranks. Iterates over all tensors in the dictionary and broadcasts each one from

Build a peft_config dict compatible with vLLM's PEFTHelper from megatron lora config. Args: lora_config: Megatron lora configuration dict

Build and configure the scheduler for the diffusion model. The returned scheduler has timesteps and sigmas already set. Args: model_c

Build arguments for SGLang profiler. Args: profiler_config (ProfilerConfig): The unified profiler configuration. tool_config

Build arguments and environment variables for vLLM profiler. Acts as an adapter to bridge verl's unified profiler config and vLLM's specific

Bulk offload group.

Calculate required CPU memory in GB. MegaTrain stores on CPU: - FP32 master weights: 4 bytes/param - Adam optimizer states (m + v): 8

Calculate fixed GPU memory usage in GB. Fixed components: - CUDA context: ~2 GB - Double-buffered layer templates: 2 * layer_size_bf1

Calculate a value based on majority voting. This function identifies the most common value for a specified vote key in the data, then re

Calculates a score for the interaction, potentially considering aspects like partial exposure & in-context task switching. sh

Calls the remote sandbox API to execute code with retry logic for Gateway Timeout, using increasing delay between retries. Logs internal call

Calls the remote search API to perform retrieval with retry logic for various errors, using increasing delay between retries. Logs internal

Perform ceiling division of two integers.

Compute ceiling division of a by b. Returns the smallest integer greater than or equal to a/b. Uses the identity: ceil(a/b) = floor((a + b -

A helper method to check if the passed module's key name matches any of the exclude modules in the adapter_config. Adapted from https://githu

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

A helper method to check if the passed module's key name matches any of the target modules in the adapter_config. Adapted from https://github

Stop worker threads/processes and cleanup resources.

Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py Args: token_level_rewards: `(torch.Tensor)`

Compute advantage for GRPO, operating only on Outcome reward (with only one scalar reward for each response). Args: token_level_

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 rollout correction metrics from log probabilities during training. This function is used in the actor to compute metrics using the CURREN

Compute teacher log probabilities for a single unpadded sequence.

Concatenate multiple TensorDicts along dimension zero. Combines a list of TensorDicts into a single TensorDict by concatenating all tensors a

Convert megatron lora target modules to HF-style module names. Args: megatron_modules: List of megatron-style module names. Returns:

Async copy from src to dst on given stream.

This function returns a FusedWorker class. `FusedWorker.{class_name}` -> FusedClass Use `class_name` as a param to directly access t

Load a Qwen3-1.7B model with only 2 layers from pretrained weights.

Create mapping from roles to worker classes Args: config: Configuration object Returns: dict: Mapping from roles to wor

Create a dataset.

Create a dataset.

name: name of the parameter train_params: training parameters infer_params (Iterable[torch.Tensor]): a iterator towards list of parameter

Embed state, action and time tokens as the transformer suffix. Args: state: (B, state_dim) robot state; ignored when pi05 is enab

Enable QAT fuse mode: sets up fusion siblings for weight scale fusion.

This function is used to interact with the environment.

The keys of the environment variables that are used to configure the Worker.

Context manager entry for switching the engine and model into evaluation mode. Usage: with engine.eval_mode():

Execute a method on rank zero worker asynchronously. Args: method_name: Name of the method to execute *args: Position

Process exported weights through QATWeightExporter for quantized weight sync.

Extract the equation from the solution string.

The training loop of PPO. The driver process only need to call the compute functions of the worker group through RPC to const

Full forward pass for one diffusion denoising step. Args: images: List of image tensors, each shaped (B, C, H, W) after batching.

Forward + backward with an externally provided loss function. Used by VERL integration where the loss is computed externally (PPO, DPO, etc.)

We assume: - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input - The communication sha

r""" Copy paste LLaMa's forward https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/model/llama.py This func

Generate text using the SGLang server asynchronously.

Generate sequences

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

Get data-parallel group size from device mesh.

Get the enum for the reduction method of cross entropy.