This is a Module Specification dataclass. Specification defines the location of the module (to import dynamically) or the imported module its

Configuration object for megatron-core transformers. The initialization function has an argument for each parameter, including those in Model

TemporaryDirectory with a fully named directory. Empties the dir if not empty.

Base checkpointing related exception

Definition of a SubPolicy. A SubPolicy consists of two augmentation operations, each of those parametrized as operation, probability, magnitu

Configuration class for specifying the submodules of a transformer layer. This class defines the structure and default implementations for v

This is an internal class, not for use outside this module Attributes: _rank (int): rank for the value _value (float) : the value

Builder class for the BlendedDataset and MegatronDataset classes Args: cls (Type[MegatronDataset]): The class to instantiate, must inheri

Configuration class for specifying the submodules of a self-attention.

GPT-2 Language model.

The low-level interface dataset class Args: path_prefix (str): The index (.idx) and data (.bin) prefix multimodal (bool): Whethe

Configuration object for Megatron Core datasets

Mixture of experts Layer **currently only supports no token dropping**. Args: BaseMoELayer (MegatronModule): Base class for MoE layers

Inference parameters sent along with the prompts. This class contains request-level attributes that control the sampling techniques used when

Transformer class.

Dataclass for specifying the submodules of a transformer block. This class defines the structure for configuring the layers and normalizatio

A docstring section.

Anneals learning rate and weight decay Args: optimizer (MegatronOptimizer): the optimizer to be used init_lr (float): initial lea

Represents a mapping between a local object and a global object. Global object is assumed to consist of many local objects distributed betwee

GPT-specific pipeline customizations. Args: src (Union[ModelMeta, Tuple]): Model meta for loading. dst (Union[ModelMeta, Tuple]):

Rotary Embedding for language model. Args: kv_channels (int): Projection weights dimension in multi-head attention. Obtained

Wraps arbitrary load strategy and distributes the load during `load`. See `load` method docs for details. Args: strategy (LoadSharde

A wrapper class for data iterators that adds replay capability. Args: iterable: data iterator that needs the replay capability. m

Wrapper for the Transformer-Engine's `DelayedScaling` layer.

A simple counter class This class is responsible for assigning request ids to incoming requests

Wraps arbitrary strategy and distributes the save during `save`. The save distribution happens without any *data* communication. Only the *me

Configuration class for specifying the submodules of a cross-attention.

Base configuration for Megatron Core Export These parameters control the export setting for trtllm

fused operation: scaling + mask + softmax Args: input_in_fp16: flag to indicate if input in fp16 data format. input_in_bf16:

Config for the model inference wrapper NOTE : All the arguments here are obtained from arguments.py file

Language model embeddings. Args: config (TransformerConfig): config object with all necessary configs for TransformerBlock vocab_

Object that should not be stored in a checkpoint, but restored locally. Wrapping any object inside the state dict with LocalNonpersistentObject

Allows to apply transformations to tensors before/after serialization. The essence of those transformations is that they can be applied to op

Bert Language model.

Configuration object for Megatron Core GPT datasets

LLaVA multi-modal model. Args: language_transformer_config (TransformerConfig): Transformer config for the language model. langua

Submodules for the MLA self-attention layer.

Configuration for optimizer.

Config options for different prompt formats.

An implementation of the Experts layer using a sequence of MLP layers. This class executes each expert sequentially.

Vision Transformer Model.

Manages a queue of async calls. Allows adding a new async call with `schedule_async_request` and finalizing active calls with `maybe_finalize

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

Inference parameters that are passed to the main model in order to efficienly calculate and store the context during inference.

Linear Embedding

MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state

Linear Embedding

Base configuration for Megatron Core The initialization function has an argument for each parameter.

Image to Patch Embedding

Transformer class.

A range represents a start and end points for indexing a shard from a full tensor. Args: start (int): Start index. end (

Singleton Class implementing per rank Straggler Detector It use cuda events to time operation of choice using the start and stop methods whic

Configuration for DistributedDataParallel.

Inference wrapper for GPT model

Dataset to convert GPT tokens to text. Args: gpt_dataset (MultiSplitGPTDataset): GPT dataset, which outputs GPT token samples. gp

Apply Gaussian Blur to the PIL image.

A generic data loader where the images are arranged in this way: :: root/dog/xxx.png root/dog/xxy.png root/dog/[...]/xxz.png

Class for one inference request Containing relevant data for an inference request

Args: config: The config of the model. submodules: Contains the module specs for the input and output linear layers. d_mo

MultimodalProjector will take the encoded input with input_size hidden state and project it into the hidden size of the language model for mu

Multimodal Tokenizer.

Serializable data structure for holding data for blocks -- embeddings and necessary metadata for Retriever

parameters to TEDotProductAttention and fused rope kernels for the `thd` (packed) sequence format

Base class for all parsing related errors.

BERT-based encoder for queries or contexts used for learned information retrieval.

A class for generating rank groups for different modes of parallelism.

Config when the data (.bin) file and the index (.idx) file are in S3 TODO: These parameters are few and can be consolidated with parameters speci

Represents a mapping between a local tensor and a global tensor. Global tensor is assumed to consist of many local tensors distributed betwee

T5 Language model.

Wrapper for the Transformer-Engine's `DotProductAttention` layer that also has "flash attention" enabled. Note that if Megatron's parall

TRTLLM Helper class to convert export and build TRTLLM model.

Async save strategy for the PyT Distributed format. The idea is to translate MCore ShardedTensors into PyT ShardedTensors and use the async-a

Original BERT wordpiece tokenizer.

Put multiple buckets into a group so that their communications can be aggregated together. Provides functionality to register when params in

CLIP ViT vision model. Args: transformer_config (TransformerConfig): Transformer config. transformer_layer_spec (ModuleSpec): Spe

Class capturing the code and rank calling a function.

ChainedOptimizer is designed for a collection of optimizers. These optimizers are responsible for different parts of multiple models for a tr

Documents backends used in the checkpoint. Checkpoint config keeps track of formats used for storing the sharded tensors (sharded_backend) an

`Cityscapes <http://www.cityscapes-dataset.com/>`_ Dataset. Args: root (string): Root directory of dataset where directory ``leftImg8bit``

Linear layer with column parallelism. The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, ..

Tracker for the cuda RNG states. Using the `add` method, a cuda rng state is initialized based on the input `seed` and is assigned to `name`.

Custom handler to filter out logging from code outside of Megatron Core, and dump to stdout.

The TRTLLM Converter class used for GPU (on device) conversion This class is used to convert models sharded and on gpus. (It assumes that the mod

Docstring meta information. Symbolizes lines in form of :param arg: description :raises ValueError: if something happens