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

Return the most recent checkpoint directory based on a tracker file. Args: path (str): Base directory containing the checkpoint trac

Create a DataProto from a dict of tensors and non_tensors

Create a DataProto from a TensorDict. This assumes that 1. All the tensor in tensor_dict have the same dim0 2. Only dim0 is the batch

Generate synthetic data samples using batch generation. Args: demo_examples: Optional list of demo examples (dict with '

Generate responses from LLM model. Args: prompts: Single prompt string or list of prompts n: Number of respo

Tracker file rescords the latest chckpoint during training to restart from.

Get the appropriate config class for a training method.

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

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 the forward function for given model architecture.

Get the last learning rate from the optimizer parameter scheduler.

Get the optimizer parameter scheduler for Megatron.

Get the policy loss with a given name. Args: name: `(str)` The name of the policy loss. Returns: `(callable)`: T

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

Gets the position ids for Qwen2-VL, it should be generated before sharding the sequence. The batch dim has been removed and the input_ids sho

Calculates partitions of indices from seqlen_list such that the sum of sequence lengths in each partition is balanced. Uses the Karmarkar-Kar

Choose the dimension that can divide fsdp_size to avoid padding

Function that gets visible devices keyword name. Returns: 'CUDA_VISIBLE_DEVICES' or `ASCEND_RT_VISIBLE_DEVICES`

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

Return True if the current router_replay_action is RECORD (R2) for the local router instances. This inspects the first local RouterReplay ins

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

Launch http server in each node.

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

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

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 mean of tensor with a masked values.

Normalize a final answer to a quantitative reasoning question. Args: final_answer: The answer string to normalize Returns: N

Converts a video dict into a [n_frames, 3, H, W] tensor Add video sample FPS in a future MR

Release weights and kv cache in GPU memory.

Release GPU memory occupation temporarily. Args: tags (Optional[List[str]], optional): List of tags to specify which memory to re

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

Run job task in background thread.

Select specific indices from the DataProto. Args: idxs (torch.Tensor or numpy.ndarray or list): Indices to select

Sets the router replay action for all router instances.

Sets the router replay action for this layer.

Sleep all rollout replica instances.

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

Perform a training step on a batch of data. Args: data: The input data for training, typically containing tensors and me

Context switch hybridengine to trainer mode.

Context switch hybridengine to trainer mode.

Update the Ray actor creation options. Args: options: Dictionary of options to update

Validate if a sample has the correct format. For structured generation samples (from Outlines), validation is minimal since Outlines

TODO(zhangchi.usc1992): We may change the implementation later

Initialize the FSDPEngine. Sets up distributed device meshes, LoRA, and offload policies based on config. Args:

Validate actor configuration parameters.

Set the value of an attribute. Check if the attr is mutable before setting the value.

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group