If distributed is initialized print only on rank 0.

Return arguments.

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

Return tokenizer.

Write timers to a tensorboard writer

Start the timer.

Stop the timer.

check if the filepath for a text tokenizer exists

Set random seed for reproducability.

Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. Parameters: key_states (`jt.Va

given a generator of metrics for each of the data points X_i, write the metrics, text, and labels to a csv file

Get the model parallel group the caller rank belongs to.

Return my rank for the model parallel group.

A unified checkpoint name.

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

Return world size for the model parallel group.

Return timers.

Track the rng state.

encode text using text tokenizer and shift Id values for command tokens

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

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

get command token corresponding to `name`

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

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

Simple linear layer with weight initialization.

Init method based on N(0, sigma).

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

Load a model checkpoint and return the iteration.

Customized load.

Log a group of timers.

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

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

Make sure the input variable is not None.

Parse entire arguments.

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

Given something that might be a URL (or might be a local path), determine which. If it's a URL, download the file and cache it, and retur

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

Instantiate a PreTrainedBertModel from a pre-trained model file. Download and cache the pre-trained model file if needed.

Get cuda rng tracker.

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

Build the model.

Initialize torch.distributed.

Load the state dicts of each of the models

Save a model checkpoint.

Use this function to override the state dict for saving checkpoints.

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.

Converts a sequence of ids in BPE tokens using the vocab.