Returns whether the CUDA Graphs have been created.

Return `True` if context uses dynamic batching.

Query if currently capturing.

Checks if the internal RNG state has been set wirth set_states().

Return True if in the first virtual pipeline model-parallel stage, False otherwise.

Parse the action from the message. Returns the action.

Parse config dictionary and return a schema instance. Expected config format: {"kitchen_config_type": "QLinearParams", "recipe_idx"

Make a directory, conditional on not being in validation mode. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Shard buffer into data_parallel_world_size chunks of equal size.

Function to share the input embeddings and output logit weights.

A simple generate function without using KV-cache.

Symbolic function for tracing.

Convert a conversation to tokens. Args: conversation (List[Dict]): Sequence of system/user/assistant messages. Mu

Single training step.

If distributed is initialized or rank is specified, warn only on rank 0.

Zeros out all grad buffers. Needs to be called at the beginning of each training iteration.

Communicate tensor shapes between sender and receiver ranks in the bridge. This is used to communicate tensor shapes before actual tensor com

Get value from src_dict according to the key, Copy and expand value to dist_dict with new key. The new key is generated by fo

Gather tensors and concatinate along the last dimension.

Return the random number generator state of the specified GPU. Arguments: device (int): The gpu to retrieve the rng state clone (

Ascertain whether a path is in S3 Args: path (str): The path Returns: bool: True if the path is in S3, False otherwise

Args: value (V): value to match rule (Dict[str, Union[str, V]]): rule dict like { "method": str

r"""Reorder dims based on the order of self.dim_names

We do some checks and return attention mask dimensions for self attention Transformer engine library underwent a lot of change. So we need to

Sets the random number generator state of the current GPU. Arguments: new_state (torch.ByteTensor): The desired state device (int

Pad vocab size so it is divisible by model parallel size and still having GPU friendly size.

Add event.

Applies tokenizer's chat template.

Execute weight gradient computation

Broadcast a tensor from last pipeline stage to all ranks. Args: size: Expected tensor size dtype: Expected tensor dtype t

Given size and type of a tensor on all ranks and the tensor value only on a specific rank, broadcast from that rank to all other ranks. Args:

Creates a random tensor for testing purposes. Returns: torch.Tensor: A random tensor of shape (1024, 1, 1024) with bfloat16 dtype

QK Clipping is a technique to clip the query and key attention logits to prevent the attention logits from exploding. This function i

QK Clipping is a technique to clip the query and key attention logits to prevent the attention logits from exploding. Per MuonClip us

Convert tokens to IDs.

Enable asynchronous grad reductions

Encode data batch into a list of tensors. Args: data_batch (Dict): Dictionary containing input data Retu

Load an image from path, ensure RGB, encode as PNG base64 string.

Build filename's path if it does not already exists.

Fetch a communication buffer for data-parallel operations. The size of the bucket is defined by the `DataParallelBuffer` instance.

Flush all staged chunks to the backward queue in reverse order.

Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.

Release the storage of a temporary communication bucket. If the bucket is temporary, this method frees its storage.

Freeze model modules. Make specific modules non-trainable by setting requires_grad to False. Args: freeze_language_model

Get file path to index that has been populated with vectors. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Get the context-parallel group the caller rank belongs to.

Get the default process groups for MoE. Returns: ProcessGroupCollection: The default process groups for MoE.

Get the expert parallel rng tracker name

Returns Mamba inference state config from the model if it is a hybrid model.

Provides human-readable labels for the tracked metadata fields.

Helper function to get module spec for MLP/MoE

Get the outer-FSDP process group.

Get parameter state (i.e., parameter & optimizer tensors). This method performs two steps: - For each DP rank, copy param & optimizer

Return caller's rank for the pipeline-model-parallel group.

Return world size for the pipeline-model-parallel group.

Return tensorboard writer. It can be None so no need to check if it is initialized.

Config: path (str): absolute path on the VM to fetch dest (str): file name of the downloaded file multi (bool) : optional

Test if context contains unfinished requests.

Initialize the tensor data of this ShardedTensor. Only called if `data` attribute is None. Args: device (Union[

Check if memory blocks are available. Args: num_blocks (int): Number of blocks to check. Return: (bool) Is m

Checks if ShardedTensor is flattened and more than 1-dimensional Args: sh_ten (Any): any object Returns: bool: whether the g

Check if the current process is the first stage of the pipeline

Load the state dicts of each of the models

Log on first rank in each pipeline stage Args: logger (logging.Logger): The logger to write the logs args (Tuple[Any]): All logg

Creates a distributed tensor (DTensor) from a local tensor with support for Megatron-FSDP and Tensor Parallel scenarios. This function i

Move a tensor to CPU if it is on GPU. Args: tensor (torch.Tensor or None): The tensor to move to CPU. as_numpy (bool): Whether to

A utility function for preparing model for inference The function gets called once before the auto regressive inference loop. It puts