Update the module config with the override_config_kwargs. Args: module_config: The module config from Huggingface Transformers. ov

Compute the mean of tensor `x` over valid entries indicated by `valid` mask. Args: x: Tensor of shape (B, ...) containing values to avera

Wait for the broadcast operation to complete. Returns: dict[str, TensorMeta]: The bucket meta after broadcast.

Block until remote agent read complete.

Wait for the broadcast operation to complete. Returns: dict[str, TensorMeta]: The bucket meta after broadcast.

Reorder the data on single controller such that each dp rank gets similar total tokens.

broadcast tensor in tp shards across mp_group

Build multimodal embeddings and attention mask

Builds metadata used for sharded_state_dict versioning. The whole content metadata is passed to ``sharded_state_dict`` model and optimizer m

Compute the reward with colocate reward model.

Dump rollout/validation samples as JSONL.

Extract reward from already-computed rm_scores if available. This has been deprecated. Args: data: DataProto obj

Dynamically import attention functions based on available hardware.

Get attribute from skip config, supporting both dict and SkipConfig dataclass.

Get raw CUDA stream pointer.

Async prefetch layer to GPU.

Ensure Flash Attention implementation is preserved when layer is moved to GPU. HuggingFace may reset _attn_implementation during .to(device) or d

Inner function that handles the core timing logic. Args: name (str): The name/identifier for this timing measurement. timing_raw

Update the input_ids, attention_mask, position_ids, and loss_mask of the request in additive manner.

Assert that every group of force_group_size consecutive samples stays together. Args: batch_idx_list: Index lists returned by prepare_mic

Add items to the replay buffer. Args: partition_id (str): Partition of transfer queue, e.g. "train" or "val". items (

apply_chat_template to messages with special attention to template requiring at least one user message, e.g. Qwen3.5. Args: processor

Performs bootstrap resampling to estimate statistics of metrics. This function uses bootstrap resampling to estimate the mean and standard d

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 the attention mask for the response part of the sequence. This function extracts the portion of the attention mask that corresponds to th

Deprecated. Please use copy_to_local instead.

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

Apply one denoising step of the noise x_t at a given timestep. Args: state: (B, state_dim) robot state. prefix_pad_ma

Disable all adapters temporarily under the context in the model for LoRA

Embed visual and language inputs as the transformer prefix. Args: images: List of (B, C, H, W) tensors. img_masks: Li

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

Convert an iterable to an async iterator.

Extract reward tensor and extra info from batch data.

torch.nn.utils.clip_grad_norm_ cann't run on cpu parameter DTensor

Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the parameters from rank 0 to all ot

Generate tokens from prompt ids. Args: request_id (str): request id for sticky session. prompt_ids (List[int]): List

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

Return the registered subclass for ``model_config.architecture``.

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

Partition samples into k groups while keeping samples with the same uid together. Args: seqlen_list: List of sequence lengths for ea

Get the number of gpus in this cluster.

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

Compute per-group mean/std/count in pure PyTorch. mean_g = sum / count std_g = sqrt( max( (sum2 - sum^2/count) / max(count-1, 1), eps )

Return True if the current router_replay_action is REPLAY_FORWARD for the local router instances. This inspects the first local RouterReplay

Compute KL divergence given logprob and ref_logprob. Optionally using straight through to bind k2 on other kl penalty compute method for unbiased

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

Launch http server in each node.

Load model back to GPU.

Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_a

Offload an FSDP2-wrapped model to CPU (reshard, move to CPU, optional cache clear).

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

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

Reduces a dictionary of metric lists by computing the mean of each list. Args: metrics: A dictionary mapping metric names to lists o

Register a new rollout replica type.

Reload.

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