Recursively sets an use_cache to `value` on a list of pytorch modules, if they have a use_cache attribute. use_cache is used to decide whethe

Send signal to all workers to terminate if we've finished the process

Recursively clears the kv cache on all layers

generator for getting default values.

Launch processes and gracefully handle failures.

Exponential moving average

Ensure that numerator is divisible by the denominator.

instantiates NeoXArgs while reading values from input dict

retrieves the activation function specified in neox_args

A unified checkpoint name.

Get coord data for coord check. Train the models in `models` with data from `dataloader` and optimizer specified by `optimizer` and `lr` for `

Build masks and position id for left to right model.

Calculate the global rank corresponding to a local rank zero in the model parallel group.

Return set of symbol pairs in a word. Word is represented as tuple of symbols (symbols being variable-length strings).

Get the pipe parallel group the caller rank belongs to.

Get Weights and Biases API key from ENV or .netrc file. Otherwise return None

True if is the local main process

Load a model checkpoint and return the iteration.

Local rank of process

Fills the input Tensor with a (semi) orthogonal matrix, as described in Exact solutions to the nonlinear dynamics of learning in deep linear neura

Reduce a tensor of losses across all GPUs.

Initializes the model for evaluation or inference (doesn't load optimizer states, etc.) from command line args. use_cache: bool Whet

Set lr for all parameters groups.

iterator producing text completions neox_args: NeoXArgs. model: a Megatron model. context_tokens: the prompt to complete; unpadded l

tokenizes dataset

Sets up the model for training by turning off k/v caching.

test that there are no duplicate arguments

Root Mean Square Layer Normalization :param dim: model size :param p: partial RMSNorm, valid value [0, 1], default -1.0 (

Build an array with length = number-of-epochs * number-of-documents. Each index is mapped to a corresponding document.

Build doc-idx, sample-idx, and shuffle-idx. doc-idx: is an array (ordered) of documents to be used in training. sample-idx: is the start docum

Build the size on rank 0 and broadcast.

Build the range [0, size) and shuffle.

Check that all the keys have the same target data type.

Gather logits from all data parallel ranks

Scatters the inputs to all data parallel ranks.

num_splits: whether to parallelize over the seqlen_k dimension (num_splits > 1) or not (num_splits = 1). num_splits=0 means it will be set by

num_splits: how much to parallelize over the seqlen_q dimension. num_splits=0 means it will be set by an internal heuristic. We're exposing n

Get slopes for Alibi positional embedding n : int = number of heads. For best performance, restrict n to a power of 2.

Initialize torch.distributed and mpu.

Based on number of samples and sequence length, calculate how many epochs will be needed.

Total number of tokens in the dataset.

Rescale parameters to convert SP initialization to μP initialization. Warning: This method is NOT idempotent and should be called only once

Set random seed for reproducibility.

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

Write arguments to tensorboard.

Backward step.

Returns a shuffled, bounded cartesian product of the input sequence. Designed to cover as wide a range of permutations as possible with a lim

XXX

Build train, valid, and test datasets.