Ensure uniqueness of saved objects.

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

Compilie JIT functions before the main training steps

No interleaving is supported.

Add to index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. text_dataset (GPTToTextDataset): T

Add data for set of blocks :param row_id: 1D array of unique int ids for the blocks :param block_embeds: 2D array of embeddin

Read codes from disk, and add them to the index. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Add the embedding of each block to the underlying FAISS index

Prepend a given prefix to all ShardedBase objects in a given state dict *in-place*. Args: sharded_state_dict (ShardedStateDict): sharded

Entry point for adding to the index. We select whether to add to a new index, or validate an existing index. Args: config (RetroPrep

Sum aux losses across PP.

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

Add noise to the input tensor. Refer to https://arxiv.org/abs/2101.03961. Args: input (Tensor): Input tensor. Re

Applies auxiliary loss to the MoE layer. Args: loss_func (callable): The loss function to be used. probs (torch.Tenso

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

A baseline implementation of applying RoPE for `thd` format. Args: t (Tensor): Input tensor T is of shape [t, h, d] cu_seqlens(Te

Encourages the router's logits to remain small to enhance stability. Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) f

Assert that a tensor is not a view (i.e., its '._base' field is not set).

Apply loss-based load balancing to the logits tensor. Args: logits (torch.Tensor): The logits tensor. Returns:

Reduce a tensor of losses across all GPUs.

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

Creates the extended attention mask Converts the attention mask of dimension [batch size, 1, seq len] to [batch size, 1, seq len, seq len] an

Build the model.

Broadcast a list of interger values.

Syncs parameters across all DP ranks.

Split each document within block into consecutive retro_gpt_chunk_length size chunks. Args: config (RetroPreprocessingConfig): For DB bui

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

Get train, valid, test GPT chunk datasets. Args: project_dir (str): Retro project dir. gpt_datasets (dict): Mapping of 'train', '

Process a single indexed dataset & extract chunks. Args: config (RetroPreprocessingConfig): Retro preprocessing config. dataset_i

Iterate each indexed dataset & process its chunks. Args: config (RetroPreprocessingConfig): Retro preprocessing config. indexed_d

Abstract method implementation Args: dataset_path (str): The real path prefix to the IndexedDataset .bin and .idx files

Merge individual dataset components into single database. This method merges databases for DB types: - 'sampled': used for training the vecto

Convert to numpy and return a sample consumed by the batch producer.

Initialize tokenizer.

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

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

Build pretraining data loaders.

Build pretraining datasets.

Returns list of utf-8 byte and a corresponding list of unicode strings. The reversible bpe codes work on unicode strings. This means you

Calculate l2 norm of parameters

Check for autoresume signal and exit if it is received.

Ensure fixed arguments for a model are the same for the input arguments and the one retrieved from checkpoint.

Clear the auxiliary losses.

Clips gradient norm of an iterable of parameters whose gradients are in fp32. This is adapted from torch.nn.utils.clip_grad.clip_grad_norm

Get the code from the dtype Args: value (Type[numpy.number]): The dtype Returns: int: The code

Compile C++ helper functions at runtime. Make sure this is invoked on a single process.

Converts a sequence of tokens (string) in a single string.

Builds core config object from namespace args from given dataclass Raises exception if argument missing in args Args: args(SimpleNam

Counts the number of zeros in gradients associated with the passed-in list of parameters. Args: parameters (Iterable[Tensor] or Tenso

Directly call C++ autograd engine. To make the 'deallocate_output_tensor' (above) optimization work, the C++ autograd engine must be called d

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

T5 decoder TE spec (uses Transformer Engine components).

Group tensors by chunk and then pick the tensor with the lowest rank. NOTE: with proper loading overlap, loading from randomized ranks

Cross attention for a standard encoder-decoder model.

Convert embedding ids to text Args: ids (numpy.ndarray): The ids to convert Returns: str: The converted text

Embed block of chunks. Args: config (RetroPreprocessingConfig): Retro preprocessing config. gpt_dataset (GPTChunkDataset): Chunk

Embed a batch of tokens using the block model

Embed a batch of tokens using the query model

Embed DB chunks. Store chunks in blocks on disk. These blocks will later be merged into a single dataset for training the index. Args: