Convert TensorDict from prompt embeds with padding to no-padding format. For diffusion model training only. Currently we expect the prom

Calculate Shannon entropy from unnormalized logits. Computes H(p) = -sum(p * log(p)) using the numerically stable formula: entropy = logsumex

Copy back to CPU and release GPU slot.

Execute a method on all workers asynchronously. Args: method_name: Name of the method to execute *args: Positional ar

Create a worker group from existing detached workers. Args: name_prefix: Prefix for worker names worker_names: Names

Merge or unmerge LoRA adapters in FSDP module. For FSDP (v1), it gathers all model parameters to each device, which may cause OOM. For FSDP2,

Get RLHF dataset class. Args: data_config: The data config. Returns: dataset_cls: The dataset class.

Dynamically import and return the interaction class.

Get the last learning rate from the optimizer parameter scheduler.

Get the optimizer parameter scheduler for Megatron.

Count the number of MoE layers assigned to the current rank. When ``moe_layer_freq`` is 1 or unset, every transformer layer is an MoE layer, s

Returns the number of input images for the vision backbone. Returns: Number of images expected in the input

Returns the number of vision patches output by the vision backbone. Returns: Number of patches per image

Get placement groups and bundle indices for the replica.

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

Create a TensorDict from tensors and non-tensor data. Automatically handles nested structures in lists by converting them to NonTensorStack.

Get trajectory info. Args: step (int): global steps in the trainer. index (list): form datastore extra_info.index column.

Initialize a Mcore model. Args: tfconfig: The transformer config. hf_config: The HuggingFace config. pre_process: Op

Load a class from its fully qualified name. Args: fqn: Fully qualified class name (e.g., 'mypackage.module.ClassName'). descripti

Load weights for mcore GPT model.

End a mark range in the profiler. Args: range_id (str): The id of the mark range to end.

Start a mark range in the profiler. Args: message (str, optional): The message to be displayed in the profiler. Defaults to N

Compute sum of tensor values where mask is True. NaN values outside the mask are replaced with zeros to prevent contaminating the sum. A

Given the config, do we need reward model.

Perform an optimization step using the optimizer.

Extract images and videos from messages. Args: messages (list[dict]): Input messages. Returns: dict: Multi-m

Get the rank of this worker in the distributed setup.

Receive weights from sender and process each bucket via callback. Args: on_bucket_received: Callback function(weights: l

Record an event on this stream.

Reduce timing information across all processes. This function uses distributed communication to gather and sum the timing information from al

Decorator to register a advantage estimator function with a given name. Args: name_or_enum: `(str)` or `(AdvantageEstimator)`

Replace the current new_batch's content with that from the dumped_new_batch. In case of [Answer] mismatch.

Reset environments to specified state IDs. Args: state_ids: State IDs to reset environments to

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

Compute state features needed for SAC actor and critic. Args: s: Dictionary of tensors representing the states, with keys

Get named actor parameters for optimization/EMA updates. Returns: A list of (name, parameter) tuples representing actor-side trai

Sample a batch of data from the replay buffer. Args: partition_id (str): Partition of transfer queue, e.g. "train" or "val".

Sample a batch from all task-specific pools.

This is a convenience method to surface the language model's word embeddings, which is necessary for `finalize_model_grads._allreduce_word_emb

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

Calculate topk items take-up ratio, and min/max timing across all ranks.

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

Context switch hybridengine to trainer mode.

Split along the second dim into `n_split`, unfold it to the first dim (batch dim) Useful in passing grouped tensors that doesn't want to be sh

Union two tensordicts.

Update the policy with an iterator of DataProto Args: data (DataProto): an iterator over the DataProto that returns by

Validate critic configuration with runtime parameters. Args: n_gpus: Total number of GPUs available train_batch_size:

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