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

If distributed is initialized, print only on rank 0.

Return arguments.

Start the timer.

Write timers to a tensorboard writer

Stop the timer.

Return tokenizer.

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

Return timers.

Return world size for the tensor model parallel group.

Log a group of timers.

broadcast_nz

train function for PanguAlpha

Return my rank for the tensor model parallel group.

r""" Set config according to the mode

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

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

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

Customized load.

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

Note that this call will sync across all ranks.

Build train, valid, and test datasets.

If distributed is initialized, print only on last rank.

Make sure the input variable is not None.

Track the rng state.

Clear the internal evaluation result.

Download the dataset from the obs. src_data_url (Str): should be the dataset path in the obs tgt_data_path (Str): the local d

Helper function to evaluate and dump results on screen.

Build masks and position id for left to right model.

Return world size for the pipeline model parallel group.

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

Customized load.

Customized load.

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

Buld dataloader given an input dataset.

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

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

A unified checkpoint name.

Get cuda rng tracker.

Customized load.

Save a model checkpoint.

Set lr for all parameters groups.

Parses each jsonl line and yields it as a dictionary

Calculate the elapsed time.

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

Text generation for incremental inference Inputs: model: the model for inferencing origin_inputs: the original inputs based

Get element from dictionary by path. If element is not present, recursively add empty dictionaries. Args: d (dict): the dictionary to

Return my rank for the pipeline model parallel group.

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

Init method based on N(0, sigma).

Replace fp16 linear with quantized linear

r"""reshape nd tensor to 2d, if n <= 2, keep original shape.

r"""reshape nd tensor to 2d, if n <= 2, keep original shape.

For easy load.

For easy load.

For easy load.

For easy load.

Tokenize a string.