Calculates gradients.

Calculates logits_max.

Calculates predicted logits.

Call any DDP pre-forward hooks which are used to launch async data parallel param gather. Any other pre-forward hooks are not allowed.

Returns True if new timeouts can be computed. `.on_timeouts_updated()` resets this property back to False.

Verifies if this strategy is compatible with `loaded_backend`.

Verifies if this strategy is compatible with `loaded_version`.

Clear embedding activation buffer.

Clips gradient of an iterable of parameters in fp32 by total norm. Note that the gradients are modified in place. Args: parameters (

Compute and return grad norm, also clip grads.

Simple save and load sanity check, without any equality tests.

Test model loading with different vocab size (caused by TP padding).

Contrast distortion.

Apply conversion to val. Recursively apply conversion if `val` #is a nested tuple/list structure.

Converts a sequence of ids in BPE tokens using the vocab.

Build the split matrix from one or optionally two contributing split vectors. Ex. a standard conversion: [0.99, 0.01, 0.0] -> [(0, 0.99), (0

Converts a sequence of tokens into ids using the vocab.

Convert tokens to IDs.

Wrapper for autograd function: forward: copy, backward allreduce

Create a bwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.

Create a fwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.

Adds IOBytes write request on non-coordinator ranks.

This function solve the math problem below: There is an equation: index = sum(idx[i] * stride[i]) And

Sets the global memory buffer to None

Set the groups to none.

Helper function to shut down the rerun machine instance.

Convert embedding ids to text Args: ids (numpy.ndarray): The ids to convert Returns: str: The converted text

Perform top-k routing on a subset of expert parallel ranks. Selects N ranks for each token, then conducts top-k selection among experts on these

Disable forward pre-hooks needed for param all-gather overlap with forward compute. Skip synchronous param all-gather if `param_sync`

Applies the save distribution computed with `determine_main_replica_uniform_distribution`. Based on rank assignment, sets replica ids of the shar

Get the dtype from the code Args: value (int): The code Returns: Type[numpy.number]: The dtype

Embed a single text string. Primarily used for on-the-fly embeddings, particularly during analysis or debugging. For large scale, use

Embed a range of a text dataset. Args: embedder (Embedder): Embedder used for embedding a text dataset. text_dataset

Instantiates a ShardedObject from a unique key. Args: unique_key: a string of the form <key>/shard_<global_offset

Encode sub-dataset block, to be later added to index. Encode the data subset, generally in blocks of 1M vectors each. For each block,

Finalize the document, for use with IndexedDatasetBuilder.add_item

Exchange tensors loaded by different ranks using the specified exchange_algo. Args: loaded_tensors (Dict[_ShardId, torch.Tensor]): mappin

difflm arguments.

Extract a dict consisting of only ShardedTensor and ShardedTensorFactory objects from a given state dict with any objects. Args: shar

Finds the checkpoint for rank 0 without knowing if we are using pipeline parallelism/expert parallelism or not. Since the checkpoint naming s

Finish embedding wgrad compute.

Finishes grad sync (all-reduce or reduce-scatter) communication operations for all model gradients. When overlap_grad_reduce

Apply flattened range to a tensor.

Run a forward pass for each of the models and return the respective embeddings.

Replay the forward graph of the passed runner.

Forward step.

Apply rotary positional embedding to input tensor T in `thd` format with CP support.

Forward pass of the router gate. Args: input (torch.Tensor): Input tensor. Returns: torch.Tensor: Logits ten

model_kwargs: model: Mask predictor. prompt: A tensor of shape (1, L). steps: Sampling steps, less than o

Get number of microbatches.

ADLR autoresume object. It can be None so no need to check if it is initialized.

Generate a batch.

Generate a batch

Randomly get a crop bounding box.

Get current running global batch size, taking into account number of DP replicas might be incompatible with true global batch size if `decrease_ba

Get DB chunk as GPT token ids.

Get the document indices This method is slated for deprecation. Returns: numpy.ndarray: The document indices

Path of empty index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Path to the empty (traine

Get expert data parallel group.

Return a new fp8 context in cudagraph mode.

Generates matrix of frequencies based on positions in the sequence, used to create positional encodings

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

Mix the native spec with TENorm. This is essentially the native local spec except for the layernorm implementation is using TENorm from Trans

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

Calculate the norm of gradients in fp32. This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and added functionality to handle mode

Individual DB's directory. Args: project_dir (str): Path to Retro project dir. prefix (str): Unique relative path to dataset with

Get the trtllm model weights split per gpu Given the trtllm mapping information (tp, pp rank etc) we split the model weights in a list, with

Get main_grads that should be taken into account to compute the grad norm. Filter parameters based on: - grad should not be

Paths to merged datasets. Args: project_dir (str): Path to Retro project dir. Returns: A dict of chunk databases, one for ea

Get meta from key, which must be either 'src' or 'dst'.

Get current global batch size.

Helper method to get the microbatch_id within model chunk given the iteration number.

Build the model.