Determine if peers are sending, and where in data structure to put received tensors. Return a boolean if the pipeline stage expects to

Adds a given strategy to the registry of default strategies. Args: action (StrategyAction): specifies save/load and sharded/common

Reset internal offset to the start and return self.

Which row parallel linear module the backend uses

Run personalization evaluation for a scenario. Returns dict with scores, average, weight_version.

The forward pass of the model for inference Args: inference_input (Dict[str, Any]): A dict containing the inputs for the VLM mode

Save the trajectory of the agent to a file. Args: agent: The agent to save the trajectory of. path: The path to save the trajecto

Batched recv from next rank and send to previous rank in pipeline.

set the scale of the aux loss. Args: scale (torch.Tensor): The scale value to set. Please ensure that the scale passed in

Set the rng states. For efficiency purposes, we do not check the size of seed for compatibility.

Setup the global config file.

Provides implementation for sharded_state_dict method for non-MegatronModules. Tries to call `module.sharded_state_dict` when possible, other

Return True when a periodic action (eval/save/checkpoint) should run. Args: rollout_id: The current rollout index (0-based).

Stop communication for the module.

Remove ANSI escape sequences from text.

Converts list of tokens to it's ids.

Validate if the ShardedTensors and ShardedObjects from multiple processes define correct sharding. Local ShardedTensors and ShardedObject metadat

Wait for the bucket to be ready.

Initialize self.

Perform AlltoAll communication on a context parallel group, transform the input tensor from shape [global-sequence/context-parallel-size,

Build the final interaction result with all collected data. Args: res: InteractionResult instance to populate

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

Make sure the input variable is not None.

Helper function to lazy initialize the backward configuration.

Given a model param, get the index sub-range of the param that this data-parallel rank owns.

Initialize affine weight for model parallel on GPU.

Test fp8_param with gpt_model.

Send generation prompt to Tinker teacher model and return decoded text.

Truncate tokens and labels if they exceed packing sequence length.

Zero out the gradient for a group of parameters. Note: copied from torch.optim.optimizer.

Utility to add more attributes to sampling params Use this method to pass in a custom dictionary to add more sampling parameter attributes.

Fast path to add dummy requests without allocating real KV blocks.

Add request to inference context. Args: request_id (int): Unique ID of request. prompt (Union[str, Tensor]): Prompt a

Allocates layers according to the requested distribution of layer types.

Append given string to progress log.

Wrapper for controller.generate_output_tokens_dynamic_batch(), to match vLLM API. Uses `asyncio` for continuous generation which allo

Builds a computation schedule plan for the model. This function creates a schedule plan for a model chunk, including preprocessing, t

Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.

Clear the mtp losses.

Which module for sequential layernorm and linear

Test model saving and loading with different TP/PP

Take a response and a golden answer and return a score. Supports tagged or boxed answers. Uses the final answer in the response string to com

Performs a conv1d on one context parallel rank, using slices of the weight and bias from the convolution that would be run when cp_si

Convert model weights to trtllm model weights This method goes through each layer in the model state dict and converts to equivalent trtllm m

Delete all CUDA graphs.

Cross-process throttle + exponential backoff on 429 for DeleteInstance. Also treat NotFound as success (idempotent cleanup).

Detaches a tensor and stores it for backward computation.

Calculates the elapsed time and restarts timer. Args: reset (bool, optional): Resets timer before restarting. Defaults to True.

Return matching values in a single list Args: x (Union[dict, list]) : state dict to process. Top-level argument must be a dict or list

Run forward passes only and collect non-loss outputs (e.g., logprobs). The model is put into evaluation mode, a forward-only pipeline pass is

Constructor which turns a ShardedTensor into CheckpointableShardedTensor Args: sh_ten (ShardedTensor): a sharded tensor to wrap

Gather an unevenly sharded DTensor distributed across multiple ranks, reconstructing the full (unsharded) tensor on each rank. This func

Generate a state dict from given model, optimizer, scheduler, rng state and others.

Returns a slice of D relevant to the current context parallel rank

Playwright cannot get the url of active tab directly, so we need to use accessibility tree to get the active tab info. This

Returns the inference data given context window This function gets called iteratively in a loop . Given the start and end context pos

Get the path to the data file from the prefix Args: path_prefix (str): The prefix Returns: str: The path to the data file

Calculate the capacity of each expert. Args: num_tokens (int): num of the input tokens. num_experts (int): num of the expert

Get the current review stage of the MR.

Get the FSDP buffer for the given bucket ID.