Get the size of the dtype/code in bytes Args: key (Union[int, Type[numpy.number]]): The dtype or code Raises:

Return arguments.

If distributed is initialized, print only on rank 0.

If torch distributed is initialized, log only on rank Args: logger (logging.Logger): The logger to write the logs args (Tuple[An

Log on rank 0. Args: message (str): Message to log.

Start the timer. Args: barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.

Stop the timer. Args: barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.

Updates internal state in grad scaler based on whether NaNs are seen in grads or not.

Retrieve 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()

logs the timers passed in names to stdout. Example usage is to log average per step value for timer 'foo', this function can be called with

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

Return whether the IndexedDataset exists on disk at the prefix Args: path_prefix (str): The prefix to the index (.idx) and data (

Return world size for the tensor model parallel group.

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

Add vectors to index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. text_dataset (GPTToTextDat

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

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

Set random seed for reproducability.

Set the NCCL process group options. Args: pg_name (str): process group name nccl_comm_cfgs (dict): nccl communicator configuratio

Determine the directory name for this rank's checkpoint.

Return my rank for the tensor model parallel group.

The PAD token id Raises: NotImplementedError: Non-abstract, optional attribute

Get cuda rng tracker.

Maps dicts and lists *in-place* with a given function.

Return timers.

Write the index (.idx) file Args: sequence_lengths (List[int]): The length of each sequence sequence_modes (Optional

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

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

Make sure the input variable is not None.

Return matching and nonmatching values. Keeps hierarchy. Args: x (Union[dict, list]) : state dict to process. Top-level argument must be

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

Simple linear layer with weight initialization.

Return my rank for the pipeline model parallel group.

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

Return the global rank that follows the caller in the pipeline

Return the global rank that preceeds the caller in the pipeline

Return tokenizer.

Checks if given `replica_id` is considered as main. "Main" replica is: - integer 0 - or an iterable with all 0 elements It is the a

Note that this call will sync across all ranks.

Allows to construct the ShardedTensor given offset specified in process ranks. Args: key: unique key data: local tens

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

Wraps tensors from transformer layers with ShardedTensor or ShardedObject. For a given `state_dict`, wraps: - all _extra_states with ShardedO

Returns iterator over (nested) values of a given dict or list.

Function to share the input embeddings and output logit weights.

Write timers to a tensorboard writer. Note that we only report maximum time across ranks to tensorboard. Args: names (List[str]):

Gather tensors and concatinate along the last dimension.

Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field. This method should be called right after the output tensor has been

Enable asynchronous grad reductions

Get file path to index that has been populated with vectors. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Get an attribute from a wrapped model. If return_model_obj is true, return the object that has the 'attr' attribute; otherwise, return the att

Divide existing and missing blocks evenly across all ranks. See 'get_blocks()' above for description. The returned lists of existing and miss

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

Default returned here, but the distributed optimizer overrides this.

Return world size for the pipeline model parallel group.

Load the state dicts of each of the models

Helper for instantiating a ShardedTensor where the `tp_axis` dimension is sharded across TP group. Optionally, can provide offsets which prepend

If distributed is initialized, print only on last rank.

Make a directory, conditional on not being in validation mode. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Build each MidLevelDataset split from a single LowLevelDataset Args: dataset_path (Optional[str]): The path on disk which defines

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.

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

Zero out the gradient for a group of parameters. Note: copied from torch.optim.optimizer.

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

Broadcast a tensor from last pipeline stage to all ranks.

Broadcast tensor values from last stage into the first stage.