Serializable data structure for holding data for blocks -- embeddings and necessary metadata for ART

Similar to normal implementation of distributed sampler, except implementation is at the batch sampler level, instead of just the sampler leve

Class that manages dynamic loss scaling. It is recommended to use :class:`DynamicLossScaler` indirectly, by supplying ``dynamic_loss_scale

Loader for IndexedDataset

Parallel self-attention layer abstract class. Self-attention layer takes input with size [b, s, h] and returns output of the same size.

Transformer class.

A class to represent a list of tokenized text.

Original BERT wordpiece tokenizer.

Anneals the learning rate.

Runs basic tokenization (punctuation splitting, lower casing, etc.).

Masked LM head for Bert Arguments: mpu_vocab_size: model parallel size of vocabulary. hidden_size: hidden size init_metho

Tracker for the cuda RNG states. Using the `add` method, a cuda rng state is initialized based on the input `seed` and is assigned to `name`.

Language model embeddings. Arguments: hidden_size: hidden size vocab_size: vocabulary size max_sequence_length: maximum s

:class:`FP16_Optimizer` is designed to wrap an existing PyTorch optimizer, and manage static or dynamic loss scaling and master weights in a

Wrapper object for a BlockData which similarity search via FAISS under the hood

fused operation: scaling + mask + softmax Arguments: input_in_fp16: flag to indicate if input in fp16 data format.

Class that manages a static loss scale. This class is intended to interact with :class:`FP16_Optimizer`, and should not be directly manipu

Contiguous memory buffer. Allocate a contiguous memory of type `dtype` and size `numel`. It is used to reduce memory fragmentation. Usage

Open Retrieval Evidence dataset class.

MLP. MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the s

A single transformer layer. Transformer layer takes input with size [b, s, h] and returns an output of the same size.

Dividing by the length (raised to some power (default 0.6))

Pooler layer. Pool hidden states of a specific token (for example start of the sequence) and add a linear transformation followed by a tanh.

Open-Retrieval Question Answer pairs dataset.

Masked LM head for T5 Arguments: mpu_vocab_size: model parallel size of vocabulary. hidden_size: hidden size init_method:

Group of timers.

Transformer language model. Arguments: transformer_hparams: transformer hyperparameters attention_mask_func: a function that take

Open Retrieval Evidence dataset class.

Runs WordPiece tokenziation.

Timer.

Utility module that implements:: def forward(self, input): return input.half()

Abstract class for tokenizer.

Bert Language model.

This function is adapted from torch.utils.checkpoint with two main changes: 1) torch.cuda.set_rng_state is replaced with `_set_cuda_

Linear layer with column parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, ..

Convert model to half precision in a batchnorm-safe way.

Runs end-to-end tokenziation.

Megatron specific extentions of torch Module.

Open Retrieval base dataset class.

Based off of pytorch RandomSampler and DistributedSampler. Essentially a RandomSampler, but this class lets the user set an epoch like Distrib

A ring of memory buffers.

Linear layer with row parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its sec

Fused operation which performs following three operations in sequence 1. Scale the tensor. 2. Apply the mask. 3. Perform

Fused operation which performs following three operations in sequence 1. Scale the tensor. 2. Apply upper triangular mask (typi

T5 Language model.

Base tokenizer class. Tokenizers implement tokenize, which should return a Tokens class.

Embedding parallelized in the vocabulary dimension. This is mainly adapted from torch.nn.Embedding and all the default values are kept. A

Split the vocabulary into `world_size` chunks amd return the first and last index of the vocabulary belonging to the `rank` partition:

Pass the input to the model parallel region.

Gather the input from model parallel region and concatinate.

All-redcue the input from the model parallel region.

Split the input and keep only the corresponding chuck to the rank.