Construct an index from an index type string. Args: index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-p

Get Tensor item data from the given bucket specified by the item ID.

Compute KV cache buffer utilization stats for the current step. Returns a dictionary with counts and percentages for both allocated block

All offsets start from the begining of the prompt.

Learning rate decay functions from: https://openreview.net/pdf?id=BJYwwY9ll pg. 4 Args: param_group (dict): parameter gro

Returns model_type attribute

Get the offset of the MTP layer.

Get the image name for a SWEBench instance.

Extract weight_version from response metadata.

GPT tokens to text.

Calculates the number of GPUs needed for a given sequence length and max sequence length per CP rank. This is used to determi

Initialize attention state so that every layer can use it. Args: construct_graph_dimensions (Optional[InferenceBatchDimensions]):

Initialize torch.distributed.

Check if the experimental attention variant is a linear attention variant.

Distributes the load and calls underlying strategy only for parts of the state dict. Steps: 1. Load metadata is exchanged between the

Load the state dicts of each of the models

Customized load.

Logs a message to the given rank.

Helper for instantiating a ShardedTensor where the `tp_axis` dimension is sharded across TP group. Optionally, can provide offsets which prep

Gets the number of transformer layers.

Handle key events.

Pre-processing gradients before the optimizer step, returns whether inf/nan is found.

Add a chunk handler to the backward queue.

Report current memory usage.

Set to the initial state (no tracker).

Set to the initial state (no tracker).

Runs the agent by sending a prompt to the LLM.

Batched recv from previous rank and send to next rank in pipeline.

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

Activate or deactivate flag that controls Megatron-FSDP model synchronization. When activated, the model parameters and gradients wil

Return the sharded state dictionary of the module.

Gets the emedding weight or output logit weights when share embedding and output weights set to True. Returns: Tensor: During pre

Initiates all necessary param all-gathers for this bucket. When ddp_config.overlap_param_gather is set to True, dispatches an asynch

Checkpointing with Transformer-Engine.

Copy the unpadded tensor to the tensor_buf, pad the tensor_buf with zero or the last value of the tensor, depending on whethe

Converts text to ids.

Train index on a representative set of vectors. Args: config (RetroPreprocessingConfig): Retro preprocessing config.

Type text character by character using pilot.press() to simulate real user input. This properly tests focus behavior and input handling, unlike s

Unset global vars. Useful for multiple runs. See `tests/unit_tests/ckpt_converter/test_ckpt_converter.py` for an example.

Unset microbatches calculator. Useful for multiple runs. See `tests/unit_tests/ckpt_converter/test_ckpt_converter.py` for an example.

Returns the original object.

Turn instance of base class into registered subclass.

Args: request_query_lengths: (>real_batch_size,) request_kv_length_offsets: (>real_batch_size,) request_t

Update number of microbatches. Args: consumed_samples (int): Number of samples consumed. consistency_check (bool, opt

Surround <image> with image tags such as <img> and </img>.

Apply rotary positional embedding to input tensor T. check https://kexue.fm/archives/8265 for detailed formulas Args: t (Tensor): In

Build each MidLevelDataset split from a single LowLevelDataset Args: dataset_path (Optional[str]): The path on disk which defines

Check if arg meta is a supported type for cudagraph input/outputs.

Helper to create MLA config.

Helper to create MLA config.

Helper to create indexer.

Helper to create sparse attention.

Helper method to distribute counts according to weights. This implementation ensures the most balanced distribution possible while ma

General function to download a specific folder from a GitHub repository. Args: repo_url: GitHub repository URL (.git) sparse

Perform a forward pass through the attention layer and the layernorms before and after the attention operations. Args:

Perform a forward pass through the feed-forward layer. Args: hidden_states (Tensor): Transformed hidden states before th

Slice ``cand_tensor[k*(t), t, ...]`` per token. Returns ``[R, *]``.

Gather tensors and concatenate along the first dimension. Args: input_tensor (torch.Tensor): A tensor to be gathered.

Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

Find next available port starting from start_port.

Helper to create a DynamicInferenceContext.

Return the coordinates of the slice of the item that is contained in this buffer shard. In other words, this returns the coordinates

Return dataset iterator.

Determine the contribution of the MegatronDataset splits to the BlendedDataset splits Args: normalized_weights (List[float]): e.g. [0.3,

Checks whether a path is in MSC path (msc://profile/path/to/file) Args: path (str): The path Returns: bool: True if the path

Move all possible GPU tensors to CPU and make a synchronization at the expected point.