See megatron.legacy.model.transformer.set_input_tensor()

Sets the is_first_microbatch flag if it exists. When this flag is set, TE modules will update their fp8 parameter cache.

Set PyTorch JIT layer fusion options.

Set the layer number for the router.

set the scale of the aux loss. Args: scale (torch.Tensor): The scale value to set. Please ensure that the scale passed in

Add a list of additional tokens to the encoder. The additional tokens are indexed starting from the last index of the current

Setup model and optimizer.

Naive implementation which reuses gather/scatter from the legacy ckpt format. During saving, gathers the parameters state on DP rank 0 and sa

Sharded state dict where each noncontiguous buffer is a separate ShardedTensor. Results in fully parallel save and load without any inter-pro

Convert MCore ShardedTensor to PyT ShardedTensor. PyT requires information about all chunks. NOTE: this function assumes regular (grid) sharding

Sinkhorn based MoE routing function

Apply sinkhorn routing to the logits tensor. Args: logits (torch.Tensor): The logits tensor. Returns: torch.

Split a tensor along its last dimension. Args: tensor: input tensor. num_partitions: number of partitions to split t

Break a tensor into equal 1D chunks across tensor parallel ranks. Returns a Tensor or View with this rank's portion of the data. Ar

swap between batches

Calculate the auxiliary loss for better load balacing. Please refer to the Switch Transformer paper (https://arxiv.org/abs/2101.03961) for detail

Ties the embedding and output weights in a given sharded state dict. Args: sharded_state_dict (ShardedStateDict): state dict with

Dispatch tokens to experts. Args: tokens (torch.Tensor): Input tokens. indices (torch.Tensor): indices tensor.

Restores the expert output to its original ordering. Args: expert_output (torch.Tensor): The output tensor from the expert models

Tokenize a string.

Utility function for tracking progress

Train the model function.

Train index. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Entry point for training the index. We select whether to train a new index, or validate an existing index. Args: config (RetroPrepro

Train index on embedded DB chunks. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Single training step.

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

Update center used for teacher output.

Set n_chunks_train & n_chunks sampled for each individual DB. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Validation method for validating saved neighbor IDs. Args: f (h5py.File): File containing save neighbor IDs.

Performs cross entropy loss when logits are split across tensor parallel ranks Args: vocab_parallel_logits: logits split across tens

Args: windows: (num_windows*B, window_size, window_size, C) window_size (int): Window size H (int): Height of image

Args: windows: (num_windows*B, window_size, window_size, C) window_size (int): Window size H (int): Height of image

Write arguments to tensorboard.

Encourages the router's logits to remain small to enhance stability. Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for d

Zeros out all grad buffers. Needs to be called at the beginning of each training iteration.

Zero out all parameters in embedding.

Define call method for ImageNetPolicy class.

Define call method for SubPolicy class.

Call timer with name and log level.

Invocation of the forward methods. Note that self.inference_params is being modified by the forward step.

Clean up the object

Clean up the object

Get an ICT example of a pseudo-query and the block of text from which it was extracted

Args: index (int): Index Returns: tuple: (sample, target) where target is class_index of the target class.

Get the data associated with an indexed sample.

Get the data associated with an indexed sample.

Return a sample that contains a dummy image, text sequence and the associated labels and cost and attention masks. Args: idx (int

Abstract method implementation Args: idx (int): The index into the dataset Returns: Dict[str, Union[int, nu

Abstract method implementation Args: idx (int): The index into the dataset Returns: Dict[str, Union[int, nu