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

T5 encoder local spec (uses Megatron-Core components).

T5 encoder TE spec (uses Transformer Engine components).

Ensure that numerator is divisible by the denominator.

Evaluation.

Extract data config from dataset. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: The config o

Extract a dict consisting of only ShardedTensor and ShardedTensorFactory objects from a given state dict with any objects. Args: sharded

Finds the checkpoint for rank 0 without knowing if we are using pipeline parallelism/expert parallelism or not. Since the checkpoint naming s

Finish parameter synchronization for all optimizers.

Flattens state dict into a single level dict. It's a copy of torch.distributed.checkpoint._nested_dict.flatten_state_dict which also accepts

Convert fp16/bf16 `val` to fp32

Convert fp32 `val` to fp16/bf16

Get list of block files, and create map. Args: dir (str): Path to directory containing saved block files. block_size

Gather tensors and concatinate along the first dimension.

Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor model parallel ranks. Returns a new Tensor with the gathere

Forward pass of the router gate. Args: input (torch.Tensor): Input tensor. Returns: torch.Tensor: Logits ten

OpenAI's gelu implementation.

OpenAI's gelu implementation.

Given prompts and input parameters, run inference and return: tokens: prompts plus the generated tokens. lengths: length of the prompt +

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

Main token generation function. Args: model: no interleaving is supported. tokens: prompt tokens extended to be of size [b, max-s

Get paths to all saved encodings. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Paths of all

Returns the output string with logged timer values according to configured options. Args: names (List[str]): Names of the timers

Return the auxiliary losses.

Generate a batch.

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

Compute number of blocks to load. This is computed by multiplying the total number of available blocks with the fraction of blocks to load.

Get samples mapping for a dataset over fixed size blocks. This function also requires a dataset of the titles for the source documents since their

Divide range [0, num_samples) to sequence of block ranges. This is a core method within the concept of block processing. The idea is to divid

Get all global ranks of the context parallel group that the caller rank belongs to.

Return my rank for the data parallel group.

Return world size for the data parallel group.

Return the causal upper triangular mask for softmax input.

Return my rank for the expert parallel group

Mix the native spec with TENorm. This is essentially the native local spec except for the layernorm implementation is using TENorm from Trans

Read index from disk. Args: config (RetroPreprocessingConfig): Retro preprocessing config. ondisk (bool): If `ondisk = True`, mem

Get an index class, given a type string. Args: index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-par-ad

Load individual dataset's chunk DB. Args: project_dir (str): Path to Retro project dir. ds_id (int): Index of dataset within blen

Load individual dataset's document offsets. Args: project_dir (str): Path to Retro project dir. ds_id (int): Index of dataset wit

Get current loss scale factor. NOTE: The output should be a CUDA tensor of size 1.

Learning rate decay functions from: https://openreview.net/pdf?id=BJYwwY9ll pg. 4

Build masks and position id for left to right model.

Get main_grads that should be taken into account to compute the grad norm. Filter parameters based on: - grad should not be

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

Get sampled dataset (for training the vector index). Args: project_dir (str): Path to Retro project dir. chunk_length (int): GPT

Get training dataset (for adding to the vector index). Args: project_dir (str): Path to Retro project dir. chunk_length (int): GP

Get directory containing neighbor IDs for a dataset (i.e., train, valid, or test). Args: project_dir (str): Retro project dir. ke

Specifically one epoch to be used in an indexing job.

Build the learning rate scheduler.

Generate mapping from optimizer state ids to model sharded parameters. Args: model_sharded_state_dict: sharded state dict with all model

Get internal representation of parameter state without any copies and modifications. This is referred to as "fully sharded bucket space" beca

This method needs to be implemented based on whether the derived class is "self-attn" or "cross-attn".

Retro encoder block spec. The retro encoder block consists of one customized Retro encoder layer (layer 1), and all of the following layers a

Get function that saves the data to storage along with its arguments. Allows the external caller to apply the save function synchrono

Gets start time of earliest job with same world size. Also returns the number of floating-point operations completed in last saved checkpoint

Train/valid/test num samples.

Get dataset splits from comma or '/' separated string list.

Get paths to saved embedding blocks. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Paths of

Weight decay incr functions

Gather meta-info about each indexed dataset. The returned info array allows for easy access to the configuration, and helps remove ambiguity.

Init method based on N(0, sigma).

Initialize the state from a pretrained BERT model on iteration zero of ICT pretraining

Set global variables, initialize distributed, and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for