The forward pass of the model for inference Appropriate utility is called for the forward pass depending on the type of model parallelism use

Sets tp attributes to tensor

Sets up embedding layer in first stage and output layer in last stage. This function initalizes word embeddings in the final stage when we ar

Setup model and optimizer.

Constructor. Args: hidden_size_student: Size of the student's hidden dimension. hidden_size_teacher: Size of

Remove 'endoftext' string.

Saves the generated key and value tensors to the end of the buffers in inference_params. Returns the full size keys and values from t

Allocate memory to store kv cache during inference.

Surround <image> with image tags such as <img> and </img>.

Build the attention mask and postition ids for the input tokens.

Check if arg is a supported type for cudagraph input/outputs.

Forward method with selective activation checkpointing.

Make sure the input variable is not None.

Gather tensors and concatenate along the first dimension. Args: input_tensor (torch.Tensor): A tensor to be gathered.

Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

Load the base state_dict from the given directory If rank0 is true, just loads rank 0 checkpoint, ignoring arguments.

Reduce-scatter the input tensor across model parallel group. Args: input_ (torch.Tensor): The input tensor to be reduce-scattered.

act (position-wise halting) with cache and true halting. Every time the act halt function predicts the halt Args:

Wrapper for autograd function

Turn ShardedTensorFactories into ShardedTensors *in-place*. Args: sharded_state_dict (ShardedStateDict): state dict possibly

Apply merges defined by ShardedTensorFactories *in-place*. Args: x1 (StateDict): state dict loaded from the checkpoint x2 (Sharde

Build pretraining data iterators.

Check for autoresume signal and exit if it is received.

Computes hashes of all parameters in model, all-gathers hashes across DP replicas, and then checks for equality between the locally-computed hashe

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

Clear the embedding data structures to save memory. The metadata ends up getting used, and is also much smaller in dimensiona

Clear the auxiliary losses.

Concat with special tokens and pad sequence to self.max_seq_length

Converts a sequence of [tokens|ids] using the vocab.

Destroy number of microbatches calculator.

Embed a text dataset. Args: text_dataset (torch.utils.data.Dataset): Text dataset to embed. Each sample of the text dataset shoul

Encode text dataset, to be later added to index. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Convert text to a list of token IDs.

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

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

Fix up query/key/value matrix ordering if checkpoint version is smaller than 2.0

Forward.

This function is n abstraction over __getitem__ with support for slicing Args: idx (int): The index into the dataset

Get directory of saved encodings. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Path to the

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

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

Get current global batch size.

Get current running global batch size. If decrease_batch_size_if_needed is False, this just equals global batch size.

Get DB chunk as text.

Get file path to empty index (i.e., trained, but unpopulated). Args: config (RetroPreprocessingConfig): Retro preprocessing confi

Get the expert-tensor and expert-model group the caller rank belongs to.

Process group for reducing gradient statistics (num_zeros & norm). The two most common cases are: - Non-distributed optimizer (defaul

Create sub-directory for this index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Path to i

Retrieve the Megatron optimizer for model chunks. We use separate optimizers for expert parameters and non-expert parameters. Args:

Method to get the operating mode

Returns model_type attribute

Get the pipeline-model-parallel group the caller rank belongs to.

Filter and flatten all tensors from args and kwargs.

Helper function to get timestamp in ms Returns: [int]: [timestamp in ms]

Weight decay incr functions

Returns a tuple of int and slice objects representing a slice of the global tensor that this ShardedTensor corresponds to.

Init method based on N(0, sigma).

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

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

Load tensors metadata from the checkpoint. Returns a dictionary similar to a sharded state dict, but note that the dictionary keys are simply

Make index object verbose. Args: index (faiss.Index): Faiss object to set verbose. verbose (bool): Sets whether index

Simple map-reduce implementation following `more_itertools.map_reduce` interface.

Returns checkpoint config if `checkpoint_dir` is a distributed checkpoint and None otherwise Args: checkpoint_dir: checkpoint directory

Do non-exponentiated normalization Args: weights (List[float]): The weights Returns: List[float]: The normalized weights

A utility function for preparing model for inference This function is called before the forward pass. It puts the model in eval mode, builds

Reads a TensorBoard Events file from the input path, and returns the summary specified as input as a list. Args: path: str, path to t

Function to log the min/max metircs and the associated rank over a time period It finds the slowest and fastest rank among all ranks. It shou

Run pipeline. Running a pipeline consists of: - Save checkpoint (includes initializing params & forward passing data). - Con

Start a new async call and add it to a queue of active async calls. This method must be called on all ranks. Args: async

Set sequence_sr for the grad_fn of a tensor to control the backward order. For older PyTorch version, do nothing (backward order is not chang

Gets the emedding weight or output logit weights when share embedding and output weights set to True. Returns: Tensor: During pre

Split and sort the input tensor based on the split_sizes and sorted indices.

Initiates all necessary param all-gathers for this bucket. When ddp_config.overlap_param_gather is set to True, dispatches an asynch

The state dict contains all non-DP-rank-dependent (i.e., non-parameter- related) optimizer variables. The returned state dict can be

For easy load.

Restore the original order of tokens after permutation. If probs are provided, it will also apply them to the tokens before restoring the ord