If distributed is initialized, print only on rank 0.

Return arguments.

Start the timer.

Stop the timer.

Return timers.

Write timers to a tensorboard writer

Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

Return tokenizer.

Return my rank for the tensor model parallel group.

Set random seed for reproducability.

Return world size for the tensor model parallel group.

Init method based on N(0, sigma).

Fork the cuda rng state, perform operations, and exit with the original state.

Helper method to get the model chunk ID given the iteration number.

Get the tensor model parallel group the caller rank belongs to.

Log a group of timers.

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

Communicate tensors between stages. Used as helper method in other communication methods that are used in megatron/schedules.py. Takes the fo

Track the rng state.

Note that this call will sync across all ranks.

Remove new lines and multiple spaces and adjust end of sentence dot.

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

Simple linear layer with weight initialization.

Load a model checkpoint and return the iteration. strict (bool): whether to strictly enforce that the keys in :attr:`state_dict` of the ch

Make sure the input variable is not None.

Reduce a tensor of losses across all GPUs.

A unified checkpoint name.

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

Build the model.

Allocate a chunk of memory from the buffer to tensor and copy the values.

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

Get cpus rng tracker.

Get cuda rng tracker.

Return world size for the pipeline model parallel group.

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

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

Load the state dicts of each of the models

If distributed is initialized, print only on last rank.

Save a model checkpoint.

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

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

Backward step through passed-in output tensor. If last stage, output_tensor_grad is None, otherwise gradient of loss with respect to stage's

Initialize tokenizer.

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

Calculate the elapsed time.

Helper function to evaluate and dump results on screen.

Forward step for passed-in model. If first stage, input tensor is obtained from data_iterator, otherwise passed-in input_tensor is used.

Return my rank for the pipeline model parallel group.

Initialize torch.distributed.

Load the state dicts of each of the models

Customized load.

Main training program. This function will run the followings in the order provided: 1) initialize Megatron. 2) setup model, optim

Set to the initial state (no tracker).

See megatron.model.transformer.set_input_tensor()

Set lr for all parameters groups.

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

Concat with special tokens and pad sequence to self.max_seq_length

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

Embed a batch of tokens using the model

Fix up query/key/value matrix ordering if checkpoint version is smaller than 2.0

Build masks and position id for left to right model.