Vocab range from global vocab size.

Receive happens between the layers with size [s, b, h].

Receive happens between the layers with size [seq_len, batch_size, hidden_size].

Apply rotary positional embedding to input tensor T. check https://kexue.fm/archives/8265 for detailed formulas Args: t (Tensor): In

Broadcast data from TP rank 0 to other ranks.

Build each MidLevelDataset split from a single LowLevelDataset Args: dataset_path (Optional[str]): The path on disk which defines

Casts weights to the expected datatype. When appropriate scaling factor is found inside self.scales, the weight gets scaled before the cas

Download the object at the given S3 path to the given local file system path Args: client (S3Client): The S3 client s3_path (str

Report only min and max times across all ranks.

Return dataset iterator.

Return dataset iterator.

Build masks and position id for left to right model. Args: data (torch.Tensor): The data tenor that holds the tokens from the dataset

Return a dict containing the main param and optimizer states corresponding to the input model_param. The structure of the returned di

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

args: y_true : 3-d ndarray in [batch_size, img_rows, img_cols] y_pred : 3-d ndarray in [batch_size, img_rows, img_cols] retur

Check if a tensor is not none, is cuda, and is contiguous.

Get the text (token ids) and document ids for a given index Args: idx (int): The index into the dataset Returns:

All-reduce the input tensor across model parallel group.

Reduce-scatter tensors on the last dimension.

Reshards the input tensor to match the sharding configuration of the reference tensor if the reference tensor is a DTensor. Otherwise, return

Set the main param and optimizer states corresponding to the input model_param. The structure of the input `tensors`: tensors = {

Create a Faiss Flat index with inner product as the metric to search against

Split the tensor along its first dimension and keep the corresponding slice.

Split the tensor along its last dimension and keep the corresponding slice.

Unshards the input tensor if it is a DTensor and otherwise returns the tensor unmodified. Args: tensor (Union[torch.Tensor, DTen

Provide function that calculates accuracies.

Add a single item to the dataset Args: tensor (torch.Tensor): The item to add to the data file mode (int, optional):

Add an incoming request This method will add the request to either the active pool or the waiting pool depending on the batch size.

Replaces prefixes *only in keys matching* with one of prefixes in the map. Args: sharded_state_dict (ShardedStateDict): sharded state dic

This function applies rotary positional embedding to the target tensor t using precomputed cos and sin of size (seq_len, d_rot / 2)

Helper method to run backward step with model split into chunks (run set_virtual_pipeline_model_parallel_rank() before calling backwar

Run beam search and post-process outputs, i.e., detokenize, move to cpu and convert to list.

Broadcast a list of float values. Args: data_parallel (bool): Broadcast across a single data parallel model replica.

Broadcast a list of integer values. Args: data_parallel (bool): Broadcast across a single data parallel model replica.

Data loader. Note that batch-size is the local (per GPU) batch-size.

Build the DistributedDataset Return None if and only if the underlying dataset class is not built on the current rank and torch.distr

Build token types and paddings, trim if needed, and pad if needed.

Calculate l2 norm of parameters

Test model saving and loading with different TP/PP

This method is called to configure the Singleton instance It should be called once per instantiation per process. Note:

Converts `text` to Unicode (if it's not already), assuming utf-8 input.

Copy tensor values from last stage into the first stage. Note that the input tensor is updated in place.

Count number of zeros in model's gradients.

Should be called at the end of each schedule function, (e.g. forward_backward_pipelining_with_interleaving) in `megatron.core.pipeline_paralle

Computes the save distribution. Should be used in conjunction with `distribute_main_replicas_with_precomputed_distribution` which applies the

Detokenize the generated tokens.

Disable asynchronous grad reductions

Embed a batch of tokens using the model

Embed a text dataset.

Enable forward pre-hooks needed for param all-gather overlap with forward compute.

Helper function to evaluate and dump results on screen.

Extract a dict consisting of only LocalNonpersistentObjects from a given state dict. Args: sharded_state_dict: state dict possibly contai

Extract a dict consisting of only ShardedBase from a given state dict with any objects. Args: sharded_state_dict: state dict possibly con

The flash decoding kernel will do the following in a single execution: 1. Compute RoPE embedding with precomputed cos & sin tensors

Forward pass data, and gather parallel output tensors.

Helper method to run forward step with model split into chunks (run set_virtual_pipeline_model_parallel_rank() before calling forward_

Wrapper for autograd function: forward: AG, backward: split <last dim>

Utility to generate the all the output tokens and probabilities for the prompts . This utility generates the output tokens for a static batch

Slice batch input along sequence dimension into multiple chunks, which are parallelized across GPUs in a context parallel group.

Cosine and sine values for RoPE are precomputed for all positions up to the maximum sequence length

Gets the data parallel group of the given tensor if it is a DTensor.

Return world size for the data parallel group.

Sub-directory for DB data. Args: project_dir (str): Path to Retro project dir. Returns: Path of the DB sub-directory within

Use this spec to use lower-level Transformer Engine modules (required for fp8 training). Args: num_experts (int, optional): Number of ex

Get empty index (i.e., trained, but unpopulated). Args: config (RetroPreprocessingConfig): Retro preprocessing config. R