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

GPT2Model adapted for pipeline parallelism. The largest change is flattening the GPTModel class so we can express it as a sequence of layers

Loader for IndexedDataset

Duplicate output to both stdout/err and file

Group of timers.

Anneals the learning rate.

Wraps the data_iterator to count the number of characters in a batch

Character Level Tokenizer

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

An adapter to run NeoX models on LM Evaluation Harness (https://github.com/EleutherAI/lm-evaluation-harness) tasks. Args: model: A N

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

A class to measure the gradient noise scale of a model while training (cf. https://arxiv.org/abs/1812.06162). The core thesis of the paper i

Designed to Integrate the pretrained OpenAI GPT2 Tokenizers from HF

Designed to Integrate HF's Tokenizer library.

data class containing all configurations NeoXArgs inherits from a number of small configuration classes

Args for deepspeed runner (deepspeed.launcher.runner). Every argument included here will be passed as command line argument to deepspeed.laun

Checks if the past n iterations have been skipped due to overflow, and exits training if that happens.

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

T5 Relative Position Bias parallelized in the heads dimension Based on https://github.com/lucidrains/x-transformers/blob/6b93c21be0d0a679da6f7b96

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

Implements SM3 algorithm. It has been proposed in `Memory-Efficient Adaptive Optimization`_. Arguments: params (iterable): iterable of

Designed to Integrate SP's Tokenizer.

Used to convert a deepspeed PipelineModule to an nn.Sequential like model whilst retaining activation checkpointing.

Tokenizer from OpenAI's tiktoken implementation

Timer.

Original GPT2 BPE tokenizer.

MADGRAD_: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimization. .. _MADGRAD: https://arxiv.org/abs/2101.1

Abstract class for tokenizer.

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

Dataset registry class to automatically download / extract datasets

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

Extends Embedding to forward attention_mask through the pipeline.

Args for deepspeed config Every argument included here will be included in deepspeed config json #TODO this list is not complete as compa

LR Scheduler Arguments

Logging Arguments

Model Arguments

Optimizer Arguments

Misc. Arguments

Parallelism Arguments

Text Generation arguments

Tokenizer Arguments

Training Arguments

Just a helper class to pass presents through to the output when doing inference with a Pipe Parallel model

A Parallel Linear Layer transforming the transformer outputs from hidden_size -> vocab_size

Another helper class to pass presents through to the output when doing inference with a Pipe Parallel model

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

Extends ParallelTransformerLayer to forward attention_mask through the pipeline.

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

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

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

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: Note th

Pass the input to the model parallel region.

Gather the input from model parallel region and concatinate.