Clear the auxiliary losses.

Get all input files and an output path for a merged file.

Return the moe layer wise tracker.

Get TE version from __version__; if not available use pip's. Use caching.

Check if config logger is enabled.

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

r"""Generate rank enumeration before calling new_subgroups_by_enumeration This function returns ranks grouped by the specified dimensions, bu

Which column parallel linear module the backend uses

Compile C++ helper functions at runtime. Make sure this is invoked on a single process.

Decoding function.

Maps dicts and lists *in-place* with a given function.

Wrapper for autograd function: forward: AG, backward: RS <first dim>

Retrieves a default strategy for a given action, backend and version.

r"""Get a process group based on a list of dimension names Args: dims: Name of leading dimensions to create process group

Which linear module TE backend uses

Sets tp attributes to tensor

Which module to use for activation function

Track the rng state.

Function to run process tokens with input embeddings

Get current global batch size.

Get the fine-grained offload context

Return one logger. It can be None so no need to check if it is initialized.

Return caller's rank for the tensor-model-parallel group.

Checks if the internal RNG state has been set with set_states().

Encodes the input dict (dict_data) using the JSONEncoderWithMcoreTypes and dumps to disk, as specified via path

Entry-point for creating viewless tensors. This method should be used, rather than calling 'MakeViewlessTensor' or '_kernel_make_viewles

Receive forward activation tensor. Args: tensor_shape: Expected activation tensor shape Returns: A dictionar

Run a coroutine in the background event loop.

Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers.

Set random seed for reproducability.

Dummy add_request.

Which module to use for attention

Simple context manager for timing functions/code blocks. Args: name (str): Label describing the code being measured. logger (logg

Execute a command in the bubblewrap environment and return the result as a dict.

Return matching and nonmatching values. Keeps hierarchy. Args: x (Union[dict, list]) : state dict to process. Top-level argument must be

Wrapper for autograd function: forward: AG, backward: split <last dim>

Creates an asyncio loop if necessary and then returns the current asyncio loop.

Get the singleton instance of PipelineOffloadManager.

Build the model.

Determine the number of transformer layers to build for the current pipeline stage. Args: config (TransformerConfig): Configuration o

Helper function to return the singleton instance of the rerun machine.

Returns (potentially) a sub-tensor from the self.buffer for the given shape.

Import a named object from a module in the context of this function. TODO: make this importer module more robust, at least make sure there ar

Resets the model's gradients and optionally its parameters. Args: model: The model to reset. params: Optional dictionary of

Set the tensors in the store from keyword arguments.

Can be used either as a decorator or a context manager: @timer def func(): ... or with timer("block_name"): ..

Track the rng state.

Stop and remove the Docker container.

Convert checkpoint

Destroy number of microbatches calculator.

Execute a command in a Singularity container and return the result as a dict.

Execute a command in the Docker container and return the result as a dict.

Filter and slice a list of SWEBench instances.

Allows to construct a *flattened* ShardedTensor given offset specified in process ranks. Args: key (str): data (torch

Get the key for the bucket.

Tracker file rescords the latest chckpoint during training to restart from.

Get an initialized model object from any kind of user input or settings.

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

Check if the current rank is in the grid.

Check if the current process is the last stage of the pipeline

Utility to tokenize the input prompts. Args: prompt (str): The input prompt. Returns: List[int]: Returns the

Do one rollout training

Check if it is the first validation step.

Get the data parallel process group.

GPT Multi-Token Prediction (MTP) block spec.

Return tensorboard writer. It can be None so no need to check if it is initialized.