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

Return caller's rank in the expert-model-parallel group.

Get the expert parallel rng tracker name

Get project-relative directory of GPT bin/idx datasets.

Load indexed dataset meta-infos. Args: project_dir (str): Path to Retro project dir. Returns: List of metadata for each data

Path to indexed dataset meta-infos. Args: project_dir (str): Path to Retro project dir. Returns: Path to the `indexed_datase

Get paths of all database blocks of an individual dataset. Args: project_dir (str): Path to Retro project dir. prefix (str): Uniq

Simple linear layer with weight initialization.

Get merged dataset. Args: project_dir (str): Path to Retro project dir. chunk_length (int): GPT chunk length (e.g., 64).

Generate mapping from optimizer state ids to model sharded parameters. Args: model_sharded_state_dict: sharded state dict with all model

Return the global rank that follows the caller in the pipeline, for each pipeline-parallel group that the rank is part of. If it is just part

Return the global rank that precedes the caller in the pipeline, for each pipeline-parallel group that the rank is part of. If it is just par

Get root directory of all saved query data. Args: project_dir (str): Retro project dir. Returns: Path to query sub-directory

This method needs to be implemented based on whether the derived class is "self-attn" or "cross-attn".

T5 decoder block spec for Transformer Engine Args: config (TransformerConfig): config, containing number of layers for decoder

T5 encoder block spec for Transformer Engine Args: config (TransformerConfig): config, containing number of layers for encoder

Get directory for of saved embedding blocks. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: P

Get path to merged training embeddings. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Path t

Get root directory for embeddings (blocks and merged data). Args: config (RetroPreprocessingConfig): Retro preprocessing config. Ret

Return the virtual pipeline-parallel rank.

Method to check if there are requests pending This method returns False only when there are no active requests or waiting requests.

Initialize Megatron, tokenizers, and datasets.

Initialize the state from a pretrained BERT model on iteration zero of ICT pretraining

Check if a tensor is a Transformer Engine Float8Tensor

Distributes the load and calls underlying strategy only for parts of the state dict. Steps: 1. Load metadata is exchanged between the

Set required arguments from the checkpoint specified in the arguments. Will overwrite arguments that have a non-None default value, but w

Load parameter state (i.e., parameter & optimizer tensors) from DP 0 rank, using the new checkpoint format with coalesced state across buckets

Load metadata of Zarr arrays. Args: checkpoint_dir (str): checkpoint root directory get_shape_dtype_fn (str -> ((int, ...), np.dt

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

Merge individual DBs into single DB. Args: project_dir (str): Retro project dir. indexed_dataset_infos (List[Dict]): Preprocessin

Builds the model. If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model. Args:

Reformulated global shape of the flattened N-D ShardedTensor. N-D tensor global shape [X, Y, Z] and local shape [x, y, z] is reformulated int

Opens a Zarr file array with Tensorstore with basic setting. Args: arr_path (Path): path to a Zarr (Tensorstore) array

Turn optimizer state dict to sharded state dict based on model state dict *in-place*. Can be used to add sharding information to most common opti

Get the dtype to use for an index of a certain cardinality Args: cardinality (Optional[int]): The number of elements to be indexe

LM logits using word embedding weights.

Parse the dataset split ratios from a string Args: split (str): The train valid test split string e.g. "99,1,0" Returns: Lis

Parses the given S3 path returning correspsonding bucket and key. Args: path (str): The S3 path Returns: Tuple[str, str]: A

Automatically regroup the buckets of input buffers and return a list of bucket groups. In some scenarios, we need to put buckets from differ

Permute the tokens and probs based on the mask. Tokens with the same designated expert will be grouped together. The shape of mask is [tokens,

Reshapes tensor from [s, b] to [b, s] for upcoming loss masking.

Performs projection of student tensor to match teacher's size if necessary.

Randomly initialize model params.

Wrapper for p2p_communication.recv_forward used with non-interleaving schedule.

Wrapper for autograd function: forward: all reduce, backward copy

Reset timer.

Set to the initial state (no tracker).

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

Set PyTorch JIT layer fusion options.

Shard buffer into data_parallel_world_size chunks of equal size.

Generate a sharded state dictionary for the transformer layer. Args: prefix (str, optional): Prefix to be added to all k

Convert a conversation to tokens. Args: conversation (List[Dict]): Sequence of system/user/assistant messages. Mu

Apply capacity and padding to the top-k selection. Args: logits (torch.Tensor): Logits tensor. topk (int): The number of experts t

Single training step.

Log training information such as losses, timing, ....

Unset global vars. Useful for multiple runs. See `tests/unit_tests/ckpt_converter/test_ckpt_converter.py` for an example.

Update request pool status This method will full up the active request pool, if it has less than max batch size elements from the wai

This method verifies a result and possibly triggers a re-run. Args: result: result to verify. rejection_func: functio

Checks if either of the passed strategies can handle sharded objects. Args: sharded_strategy (Union[SaveShardedStrategy, LoadShardedStrat

Runs basic whitespace cleaning and splitting on a piece of text.

Returns a list of the text of each token Args: uncased: lower cases text

Enter the context introduced by the 'with' keyword Returns: _IndexWriter: The instance

Exit the context introduced by the 'with' keyword Args: exc_type (Optional[Type[BaseException]]): Exception type exc

Initialize the MoE Token Dispatcher.

Initialize the Router module. Args: config (TransformerConfig): Configuration object for the Transformer model.