Return arguments.

If distributed is initialized print only on rank 0.

Start the timer.

Stop the timer.

Retrieves a single item from the dataset with the option to only return a portion of the item. get(idx) is the same as [idx] but get

Get the data parallel group the caller rank belongs to.

Get the model parallel group the caller rank belongs to.

Return my rank for the model parallel group.

Return world size for the model parallel group.

Return tokenizer.

Tracker file rescords the latest chckpoint during training to restart from.

A unified checkpoint name.

Return timers.

Returns a 3-dimensional (3-D) attention mask :param source_block: 1-D array :param target_block: 1-D array

Save a model checkpoint.

Make sure the input variable is not None.

Embed a batch of tokens using the model

Returns a 2-dimensional (2-D) attention mask :param source_block: 1-D array :param target_block: 1-D array

Write timers to a tensorboard writer

Auxiliary function for checking the values in the checkpoint and setting them.

Sets the random number generator state of the current GPU. Argumentss: new_state (torch.ByteTensor): The desired state This function

Ensure that numerator is divisible by the denominator and return the division value.

Get cuda rng tracker.

Return T5 tokenizer.

Set global variables, initialize distributed, and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for

Make sure the input variable is not None.

Initialize tokenizer.

Clear the embedding data structures to save memory. The metadata ends up getting used, and is also much smaller in dimensionality so i

Converts a sequence of [tokens|ids] using the vocab.

Init method based on N(0, sigma).

Customized load.

Log a group of timers.

Init method based on N(0, sigma/sqrt(2*num_layers).

Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers.

For easy load.

Gather tensors and concatinate along the last dimension.

Initialize affine weight for model parallel. Build the master weight on all processes and scatter the relevant chunk.

Initialize affine weight for model parallel on GPU.

All-reduce the the input tensor across model parallel group.

Create a Faiss Flat index with inner product as the metric to search against

Create a Faiss Flat index with inner product as the metric to search against

Split the tensor along its last dimension and keep the corresponding slice.

Data loader. Note that batch-size is the local (per GPU) batch-size.

Data loader. Note that batch-size is the local (per GPU) batch-size.

Build token types and paddings, trim if needed, and pad if needed.

Converts `text` to Unicode (if it's not already), assuming utf-8 input.

Calculate the elapsed time.

Build filename's path if it does not already exists.

Build language model and return along with the key to save.

Simple linear layer with weight initialization.

Build the model.

Get rng states. Copy the dictionary so we have direct pointers to the states, not just a pointer to the dictionary.

Return tensorboard writer. It can be None so no need to check if it is initialized.

Load a model checkpoint and return the iteration.

selectively load retrieval models for indexing/retrieving from saved checkpoints

Populate members from instance saved to file

Load the state dicts of each of the models

Reduce a tensor of losses across all GPUs.

Set the rng states. For efficiency purposes, we do not check the size of seed for compatibility.

If no closure is supplied, :attr:`step` should be called after ``fp16_optimizer_obj.backward(loss)``. :attr:`step` updates th

Runs basic whitespace cleaning and splitting on a piece of text.

Returns a list of the text of each token Args: uncased: lower cases text

extracts samples only pertaining to this worker's batch

Initialize tokenizer.