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

If distributed is initialized, print only on rank 0.

Return arguments.

Return generator over underlying items.

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

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

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

Allows to construct the ShardedTensor given offset specified in process ranks. Args: key (str): unique key data (torc

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

Initialize model parallel cuda seed. This function should be called after the model parallel is initialized. Also, no torch.cuda.manual_seed

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

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

Checks if the internal RNG state has been set wirth set_states().

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

Return a list of information on the exceptions that the function may raise.

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

Saving entrypoint. Extracts ShardedTensors from the given state dict. Rank 0 saves the "regular" part of the checkpoint to common torch file.

Check if minimum version of `transformer-engine` is installed.

Load the sharded part of the checkpoint.

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

Loading entrypoint. In the steps below, the following verbs refer to corresponding objects: - load = load from checkpoint - extract = ext

Customized load.

Use this spec to use lower-level Transformer Engine modules (required for fp8 training). Args: num_experts (int, optional): Number of ex

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

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

Return tokenizer.

Return timers.

Recursive diff of dicts. Args: x1 (object): left dict x2 (object): right dict prefix (tuple): tracks recursive calls. Use

Get number of microbatches.

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

Get rank group by input token. Args: token (str): Specify the ranks type that want to get. If we want

Read bytes into a numpy array. Args: dtype (Type[numpy.number]): Data-type of the returned array. count (int): Numbe

Save the sharded part of the state dict.

Save a model, optimizer and optionally dataloader checkpoint. Checkpointing context is used to persist some checkpointing state throughout a

Use this spec for an implementation using only modules in Megatron-Core. Args: num_experts (int, optional): Number of experts. Defaults

Builds a ShardedStateDict from the original tensor

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

Return world size for the tensor-model-parallel group.

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

No-op

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

Entry-point for creating viewless tensors. This method should be used, rather than calling 'MakeViewlessTensor' or '_kernel_make_viewles

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

Load checkpoint tensors without any sharding and plain structure. NOTE: common state dict is NOT included. Args: checkpoint_dir (str

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

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

Return one logger. It can be None so no need to check if it is initialized.

Return caller's rank for the tensor-model-parallel group.

Parse all arguments.

Check if minimum version of `torch` is installed.

args: y_data : [float] 4-d tensor in [batch_size, channels, img_rows, img_cols] threshold : [float] [0.0, 1.0] return 4-d bin

Get cuda rng tracker.

Build the model.

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

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

Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list Args: blend (Option

Encodes the input dict (dict_data) using the JSONEncoderWithMcoreTypes and dumps to disk, as specified via path

If distributed is initialized, print only on last rank.

The `RankGenerator` class produces a hyper-rectangle for a given set of tensor, pipeline, data, expert, and context parallelism. If we have an

Determine the directory name for this rank's checkpoint.

Get current global batch size.

Retrieves the appropriate forward_backward function given the configuration of parallel_state. Returns a function that will perform all of th

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

Get TE version from __version__; if not available use pip's. Use caching.

Set random seed for reproducability.

Performs inference step for decoding

Internal function that safely checks and returns the rank of the caller.

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

Retrieves a default strategy for a given action, backend and version.