Memory optimization arguments.

Text generate arguments.

blocking call avoid explicit memory copy in _allgather_params

Build an array with length = number-of-epochs * number-of-dcuments. 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 a looped dataloader with infinite size.

Build the size on rank 0 and broadcast.

Build lambada dataset.

Build the range [0, size) and shuffle.

Traing and validation dataloaders.

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

Forward method with activation checkpointing.

Performs invalid character removal and whitespace cleanup on text.

Divide params into with-weight-decay and without-weight-decay groups. Layernorms and baises will have no weight decay but the rest will.

Return a numpy array of byte offsets given a list of sizes. Multiplies values in the sizes array by dtype size (bytes),

Set autoresume start time.

Initialize torch.distributed and mpu.

Initialize manually allocated static memory.

Checks whether CP is the codepoint of a CJK character.

Checks whether `chars` is a control character.

Checks whether `chars` is a punctuation character.

Checks whether `chars` is a whitespace character.

Create the all-reduce hook for backprop.

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

Total number of tokens in the dataset.

Parse entire arguments.

Print arguments.

Splits punctuation on a piece of text.

Strips accents from a piece of text.

Initialize ADLR autoresume.

Set random seed for reproducability.

Set tensorboard writer.

Initialize timers.

Set wandb writer.

Adds whitespace around any CJK character.

Train the model.

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

Provide function that calculates accuracies.

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

Add the embedding of each block to the underlying FAISS index

Reduce gradients across data parallel ranks.

Goes through one epoch of the dataloader and adds all data to this instance's BlockData. The copy of BlockData is saved as a

Helper function to select and build dataset.