Register some load_state_dict prehooks to handle some known state_dict key mismatch. (legacy <-> modelopt) and (default te <-> modelopt)

Adds scaling factors to the distributed and single device converters. Args: converter (ModelWeightConverter): Converter, holding

All-reduce conditional embedding grads. Reduce grads across all the pp stages to ensure that parameters of the conditional embedders (e.

Reduce loss for reporting purposes.

Build the megatron datasets for a list of prefixes in parallel Args: prefixes (List[str]): The list of prefix strings

Build docstring element. :param text: docstring element text :param title: title of section containing element :return:

Build the range [0, size) and shuffle Args: num_samples (int): The size of the first shuffle range [0, num_samples) total_size (

Build the range [0, size) and shuffle Args: num_samples (int): The size of the first shuffle range [0, num_samples) total_size (

Initialize tokenizer.

Configure number of microbatches calculator. Can be used for initialization and reconfiguration. Args: rank (int): Rank o

Convert a dense model's state_dict to a MoE model's state_dict. This function takes the state dictionary of a dense model and modifies it to

Creates the predictions for the masked LM objective Args: token_ids (List[int]): The token ids target_sequence_length

Helper function to evaluate and dump results on screen.

Extract `size` bytes starting at `offset` bytes into the cache

Single forward step. Update the allocate memory flag so only the first time the memory is allocated.

Helper function to get module spec for MLP/MoE

Returns elapsed times of timers in names. Assumptions: - All the ranks call this function. - `names` are identical on

Determines the empty tensor to use for exchange. If shard_id is needed by this rank, it will be in the `unloaded_shards`. Otherwise, the meta

Get Megatron optimizer based on parameter groups. Args: config (OptimizerConfig): optimizer configuration object. model_chunks (l

Helper function to get module spec for MLP/MoE

Calculate the number of epochs Args: num_tokens_per_epoch (int): The number of tokens in a single epoch Returns:

Calculate the number of tokens in a single epoch Returns: int: The number of tokens in a single epoch

Returns parameter groups and buffer for optimizer. Args: model_chunks (List[MegatronModule]): model chunks to create parameter

Determine the contribution of the MegatronDataset splits to the BlendedDataset splits Args: normalized_weights (List[float]): e.g. [0.3,

Get TRTLLM Config Returns appropriate TRTLLM PretrainedConfig used by TRTLLM for building engine Args: export_config (Ex

This basically handles all the cases that we expect to see. Either the list None, or it's a singleton (the usual cases, since most ranks only

Check if a tensor is not none and is cuda.

Make a viewless tensor. View tensors have the undesirable side-affect of retaining a reference to the originally-viewed tensor, even after ma

Load the base state_dict from the given directory containing the global distributed checkpoint

Apply mask to the unreduced loss tensor.

Helper function that creates a new bucket. Also updates param->bucket mapping.

Padding tensor shape to multiples of 16.

Preprocess input data before input to language model. This function is adopted from https://github.com/huggingface/transformers/blob/

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

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

Processes the input data for model parallelism support. When using sequence parallelism (SP) or context parallelism (CP), the sequence is sha

Get memory used by current process. Returns (int): memory used by current process

Get the text (token ids) and document ids for a given index Args: idx (int): The index into the dataset Returns:

Group ShardedBase objects by keys and return mappings required for recreating the original dict.

Internal method that restores the state that was saved in _save_state().

Round `batch_size` down to nearest batch size divisible by `divisor`.

Unique id of the sharded object data. Should yield the same value for same data replicated on different ranks. Args: sharded_object

prefix should be module_name to make keys identical to sequetial ones.

Builds a factory that splits a given ShardedTensor into several independent chunks.

Add a new hypothesis to the list.

Adds a new finalize function to the request. Args: fn (Callable): function to add to the async request. This function

Add an entire IndexedDataset to the dataset Args: path_prefix (str): The index (.idx) and data (.bin) prefix

Entry point for adding to the index. We select whether to add to a new index, or validate an existing index. Args: config (RetroPrep

Convert a conversation to token IDs.

Calculate auxiliary loss, attach gradient function to activation and add to logging.

Applies N-D reformulation to a given sharded state dict. After applying the method and loading the reformulated state dict, the `restore_nd_f

Distributes the save across ranks by exchanging metadata. Exchanges metadata from the state dict and computes the uniform (as close a

Perform preparation and return an AsyncRequest to the external caller. Args: sharded_state_dict (ShardedStateDict): sharded state

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

Build the model.

Broadcast a tensor from last pipeline stage to all ranks.

Broadcast a list of values with a given type. Args: data_parallel (bool): Broadcast across a single data parallel model replica.

Syncs parameters across all DP ranks.

Helper function to select and build dataset.

Get train, valid, test GPT chunk datasets. Args: project_dir (str): Retro project dir. gpt_datasets (dict): Mapping of 'train', '

Convert to numpy and return a sample consumed by the batch producer.

Build token types and paddings, trim if needed, and pad if needed.

Build pretraining data iterators.

Build pretraining datasets.

Calculates cross entropy loss.