Allocate memory blocks if available, else return None. Args: num_blocks (int): Number of blocks to allocate. Return:

Backward step through passed-in output tensor. If last stage, output_tensor_grad is None, otherwise gradient of loss with respect to stage's

Initialize tokenizer.

Compute VQA accuracy.

Context manager that sets up a deterministic environment for reproducible testing. This function: 1. Sets environment variables to ensur

Release the output tensor storages and register the recompute function as a grad hook of the hook_tensor. Note: the caller s

Completes the stream by adding a sentinel value

Forward pass of RoPE embedding. Args: max_seq_len (int): Maximum size of sequence offset (int, optional): RoPE offset

Forward step for passed-in model. If it is the first stage, the input tensor is obtained from the data_iterator. Otherwise, the passed-in inp

Forward pass for DSA Indexer that returns both index scores and top-k indices. This is used when KL loss is enabled to compare index

Slice batch input along sequence dimension into multiple chunks, which are parallelized across GPUs in a context parallel group. Args:

Return the global GlobalMemoryBuffer object

Build masks and position id for left to right model.

This function is used to get the muon optimizer for the model chunks. It is used to get the muon optimizer for the model chunks. Args:

Select the best model class. If a model_class is provided (as shortcut name, or as full import path, e.g., "anthropic" or "minisweagent.model

This method needs to be implemented based on whether the derived class is "self-attn" or "cross-attn".

Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample. After squeezing batch dimension and applying temperature scaling, this

Allocate a buffer Tensor to persistently store the data for this (shard of) the buffer.

Initialize number of microbatches calculator. Supporting backward compatibility. Args: rank (int): Rank of the GPU, only rank

Safely parse JSON body from a FastAPI request. Returns an empty dict if parsing fails or the payload is not a JSON object.

Offset of a local chunk in a global array of chunks. Returns: Tuple[int, ...]: the offset of the whole local chunk in a global ar

Merge requests into a single suspend-agnostic request object. Args: tokenizer (MegatronTokenizer): The tokenizer. Return

Create a new router with updated configuration. Args: **kwargs: Configuration parameters to update in the default config.

Note that this call will sync across all ranks.

Helper function to rename model layer names to TRTLLM Layer names We go through each layer (keys) in the model state dict, and map it

Replaces the given prefix in *all* sharded keys in a given state dict. Errors out if some key does not begin with a given prefix. Args:

Roll the tensor input along the sequence dimension with Context Parallelism (CP) support. This function extends the original roll_tensor to suppo

Saves HF tokenizer files. Args: path (str): path where to save tokenizer files.

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

Check if two images are approximately the same by SSIM

Update context state after calling engine.step(). This method is responsible for: - Update prefill requests to decode requests.

Returns True if cuda graphs are being used for this step.

Codifies the constraints on metadata attributes.

Verifies if checkpoint metadata exists and matches given strategies. If no strategies are passed, they are determined based on the checkpoint met

Build an attention-mask having shape (bs, 1, q_len, kv_len) from source_block and target_block Args: source_block (torch.

Internal method to complete the input process.

Calculate expected sequence length given text tokens length and number of tiles.

Load the base state_dict from the given directory If rank0 is true, just loads rank 0 checkpoint, ignoring arguments.

Group messages into "pages" as shown by the UI.

Sends a generic control signal to the inference coordinator. Args: signal: The signal to send, typically a value from th

Unique id of the sharded tensor data. Should yield the same value for same data replicated on different ranks. Args: sharded_tensor

Add a model call with its cost, checking limits.

Submits a new inference request to the coordinator. This method sends the prompt and sampling parameters to the inference co

Align embeddings from different modalities based on special token positions in input_ids. Args: modality_embeddings: Dictionary m

Wrapper for autograd function

Call the OpenClaw-RL training server (port 30000).

Check that a field exists on queue message if necessary.

Compares two capture dictionaries to check if they contain the same values. Args: capture_ref: Reference capture dictionary.

Compute the approximate KL divergence between two distributions. Schulman blog: http://joschu.net/blog/kl-approx.html Args: log_

Config attention-mask for encoder_mask, decoder_mask, encoder_decoder_mask conditioned on transformer-implementation (e.g. TE vs local), TE ve

Convert tokens to IDs.

Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field. This method should be called right after the output tensor has been

Exchanges local metadata with `all_gather_object` to determine global metadata. Args: sharded_state_dict (ShardedStateDict): local sharde

Clean up and write the index (.idx) file Args: idx_path (str): The path to the index file

Fully shard the model and the optimizer for Megatron-FSDP. Wraps the model as an Megatron-FSDP module, and modifies the optimizer to be

Does the backend choose a single module for layernorm and linear

Divide existing and missing blocks evenly across all ranks. See 'get_blocks()' above for description. The returned lists of existing and miss

Create iterators and a micro-batch schedule for a rollout step. - If `use_dynamic_batch_size` is False, splits into fixed-size contiguous

Get default load sharded strategy.

Get the FSDP process group.

Get the path to the index file from the prefix Args: path_prefix (str): The prefix Returns: str: The path to the index file

Build language model and return along with the key to save.

Get a model name from any kind of user input or settings.

Get the number of image embeddings per image tile.

Get list of parameters wrapped in optimizer.

Get most recent request from a request record. Args: request_id (int): Request id. Returns: (DynamicInferenc

Get the shard or virtual shard of the bucket stored in this buffer.

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

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

Return world size for the tensor-model-parallel group.