Parse all arguments.

Read bytes into a numpy array. Args: dtype (Type[numpy.number]): Data-type of the returned array. count (int): Numbe

Get the blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list Args: blend (Optional[List[str]]): The

Return True if in the last pipeline-model-parallel stage, False otherwise.

Return `True` if context uses static batching.

Parse the Google-style docstring into its components. :returns: parsed docstring

Remove a string and all nodes with this text as prefix from the trie. Args: text: String to remove (will also remove all

Print warning for deprecated `inference_params`.

Recursive diff of dicts. Args: x1 (object): left dict x2 (object): right dict prefix (tuple): tracks recursive calls. Use

Get number of microbatches.

Return timers.

Load a function from a module. :param path: The path to the function, e.g. "module.submodule.function". :return: The function object.

Save the sharded part of the state dict.

Ensure that numerator is divisible by the denominator and return the division value.

Entrypoint method that adds the argument group to the parser. Args: parser: The parser to add arguments to title: Tit

Insert a string and its corresponding token IDs, log probabilities, and loss mask into the trie. Args: text: String to in

Check if the current rank is in the grid.

Return True if in the first pipeline model-parallel stage, False otherwise.

Returns iterator over (nested) values of a given dict or list.

Copied from torch.optim.optimizer

Creates a ProcessGroup.

Get cuda rng tracker.

Every metrics-dict value is a list of 1D tensor, i.e., [torch.Tensor] with shapes exactly the same as log_probs. All metrics will be aggreg

Retrieve a single item from the dataset with the option to only return a portion of the item. get(idx) is the same as [idx] but get()

Load checkpoint tensors without any sharding and plain structure. NOTE: common state dict is NOT included. Args: checkpoint_dir (str

Pop NVTX range from stack. If msg is not provided, use the calling function's path. Args: msg (str, optional): Message to associate with

Check if the feature flag is enabled.

Push NVTX range onto stack. If msg is not provided, use the calling function's path. Args: msg (str, optional): Message to associate with

Run step() until agent is finished. Return exit status & message

Build the model.

Import the package.

Functional access to `.decode_mode`, to match dynamic context.

Wraps tensors from transformer layers with ShardedTensor or ShardedObject. For a given `state_dict`, wraps: - all _extra_states with ShardedO

Pad token ID.

Set random seed for reproducability.

Set the NCCL process group options. Args: pg_name (str): process group name nccl_comm_cfgs (dict): nccl communicator configuratio

Set global variables, initialize distributed, and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for

Customized load.

The PAD token id Raises: NotImplementedError: Non-abstract, optional attribute

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

Run pipeline. Running a pipeline consists of: - Save checkpoint (includes initializing params & forward passing data). - Con

Generates the rewrited solution.

Checkpoint function.

Get an attribute from a wrapped model. If return_model_obj is true, return the object that has the 'attr' attribute; otherwise, return the att

Determine the directory name for this rank's checkpoint.

Retrieves the appropriate forward_backward function given the configuration of parallel_state. Returns a function that will perform all of th

Return fp8 context manager. Arguments: config (TransformerConfig): Configuration object. layer_no (int): *Global* lay

Helper method to get the model chunk ID given the iteration number.

Check if a tensor is a FP8 tensor.

Checks if given `replica_id` is considered as main. "Main" replica is: - integer 0 - or an iterable with all 0 elements It is the ap

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

Return the number of chunk handlers in the queue.

Internal function that safely checks and returns the rank of the caller.

Gets a file from the server.

Get the Gloo group for distributed communication.

Retrieve the Megatron optimizer for model chunks. We use separate optimizers for expert parameters and non-expert parameters. Args:

Reset the default value of a Megatron argument. :param parser: The argument parser. :param name: The name of the argument to reset. :

Reroute to the appropriate apply_rotary_pos_emb function depending on fused/unfused kernels, or bshd (conventional) / thd (packed seq) format

Clear the embedding data structures to save memory. The metadata ends up getting used, and is also much smaller in dimensiona

Close the S3 client

GPT block spec.

Get the index offset of current pipeline stage, given the level of pipelining.

Read metadata from file system, caching for subsequent calls. Returns: Metadata: Checkpoint metadata.

Convert a DTensor to a local tensor. Args: tensor (torch.Tensor): The tensor to convert. Returns: torch.Tensor: The local

Helper function to lazy initialize the forward configuration.

Add vectors to index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. text_dataset (GPTToTextDat

Reset all stored tensors to None.

Performs inference step for decoding.

Returns the config attribute, allowed to return None

Get rng states. Copy the dictionary so we have direct pointers to the states, not just a pointer to the dictionary.

Utility to build a ``TransformerBlock`` for tests.

Log data to tensorboard Args: data (dict): Dictionary containing metric names and values step (int): Current step/epo

Send forward activation tensor. Args: output_dict: A dictionary mapping module names to tensors.