Return arguments.

If distributed is initialized print only on rank 0.

Start 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

Stop the timer.

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.

A unified checkpoint name.

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

Return timers.

Make sure the input variable is not None.

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. Arguments: new_state (torch.ByteTensor): The desired state This function i

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

Get cuda rng tracker.

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

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

Initialize tokenizer.

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 T0 tokenizer.

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.

Save a model checkpoint.

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.