Definition of a SubPolicy. A SubPolicy consists of two augmentation operations, each of those parametrized as operation, probability, magnitu

GPT-2 Language model.

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

Transformer class.

BERT-based encoder for queries or contexts used for learned information retrieval.

Dataset containing sentences and their blocks for an inverse cloze task.

BERT-based encoder for queries or blocks used for learned information retrieval.

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.

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

Vision Transformer Model.

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

Bert Language model.

Bert-based module for Biencoder model.

A struct for fully describing a fixed-size block of data as used in REALM :param start_idx: for first sentence of the block :param end_idx: f

A struct for fully describing a fixed-size block of data as used in REALM :param start_idx: for first sentence of the block :param end_idx: f

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

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

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

Float16 optimizer for fp16 and bf16 data types. Arguments: optimizer: base optimizer such as Adam or SGD clip_grad: clip gradeint

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

GPT-2 BPE tokenizer. Peculiarities: - Byte-level BPE

GPT-2 Language model.

Bert-based module for Inverse Cloze task.

Definition of an ImageNetPolicy. Implements a fixed AutoAugment data augmentation policy targeted at ImageNet training by randomly applying a

Object for taking one pass over a dataset and creating a BlockData of its embeddings

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

Open Retrieval Question Answering evaluation using Google NQ dataset.

Open Retrieval Evidence dataset class.

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

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

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

T5 Language model.

Group of timers.

Transformer language model. Arguments: transformer_hparams: transformer hyperparameters vocab_size: vocabulary size max_s

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

Runs WordPiece tokenziation.

Original GPT2 BPE tokenizer.

Timer.

Abstract class for tokenizer.

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, ..

DDP with contiguous buffers options to storre and accumulate gradients. This class: - has the potential to reduce memory fragmentation.

Abstract class for DDP.

Runs end-to-end tokenziation.

GLUE base dataset class.