Create sub-directory for this index. Args: config (RetroPreprocessingConfig): Retro preprocessing config. Returns: Path to i

Get the leader rank for a given grid and direction. We elect leader rank for each dp replica, the first tp-cp rank in the group in th

Get main_grads that should be taken into account to compute the grad norm. Filter parameters based on: - grad should not be

Helper function to get module spec for MLP/MoE

Create an NVTX range context manager.

Get the position embedding on the current context parallel rank. Args: pos_emb (Tensor): Positional embedding tensor seq_dim (int

Collect rng state across data parallel ranks.

Return num_samples samples

Retrieve an Megatron-FSDP-registered submesh by name(s).

Calculate correct sample mean for CP

T5 decoder block spec for Transformer Engine Args: config (TransformerConfig): config, containing number of layers for decoder

T5 encoder block spec for Transformer Engine Args: config (TransformerConfig): config, containing number of layers for encoder

Included for backwards compatibility.

Filter and flatten all tensors from args and kwargs.

Helper function to get timestamp in ms Returns: [int]: [timestamp in ms]

Get user's email, prioritizing public profile, then recent commits.

Weight decay incr functions Args: param_group (dict): parameter group from the optimizer.

Returns a tuple of int and slice objects representing a slice of the global tensor that this ShardedTensor corresponds to.

Init method based on N(0, sigma).

Return if this iteration we run decode only implementation.

Return True if on first or last pipeline stage, taking into account virtual pipeline parallelism.

First tensor and pipeline parallel rank.

Load common part of the checkpoint.

Return {full_key: tensor} for every tensor whose key matches *pattern*.

Args: xlsx_file (Workbook): concerned excel book sheet_name (str): sheet name options (Dict[str, List[str]]): dict like

Make index object verbose. Args: index (faiss.Index): Faiss object to set verbose. verbose (bool): Sets whether index

Create a thread cooperative group.

Create a thread cooperative group.

Make_viewless util func

Returns the Mamba state tensors for the given layer.

Simple map-reduce implementation following `more_itertools.map_reduce` interface.

Mark proxy as failed

Gives an operator's context, return a configuration key if necessary, or sentinel (None) denoting no matchers matched.

Finalizes all available calls. This method must be called on all ranks. Args: blocking (bool, optional): if True, will w

Returns checkpoint config if `checkpoint_dir` is a distributed checkpoint and None otherwise Args: checkpoint_dir: checkpoint directory

Check if there is MTP on the current rank. Behavior: - If a custom pipeline model parallel layout is provided in the config:

Do non-exponentiated normalization Args: weights (List[float]): The weights Returns: List[float]: The normalized weights

Called when the container gains focus.

Pack sequences into bins using a greedy first-fit algorithm.

Method to be applied after the convolution and SSM

Trigger post-process for int4/fp4 quantization on all rollout engines.

Prepares the inference input data. Args: prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len] e

Print the trie structure in a readable format.

Performs Quick-GELU-based GEGLU activation : quick_gelu(y1) * (y2 + offset). Args: y: Input tensor split into two halves on the last dime

Randomly initialize model params.

Receive backward activation tensor. Args: tensor_shape: Expected gradient tensor shape Returns: A dictionary

Receive tensor from previous rank in pipeline (forward receive).

Receive from previous pipeline stage and update the input buffer inplace. Args: recv_buffer: Buffer to receive data into pp_group

Release the lock, allowing others to acquire.

Release reference, this helps avoid memory leak.

Remove an instance from the predictions file.

Request input from user. Returns input text (empty string if confirmed without reason).

Reset timer.

Reset by removing all requests and reset all state.

Apply Hadamard rotation activation. Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L424-L428

Deprecated in favor of round_up_tokens and round_up_requests.

Runs the model with all-to-all overlap optimization and captures outputs and gradients. Args: model: The transformer model to run.

Runs the model in reference mode and captures outputs and gradients. Args: model: The transformer model to run. input_tensor

Send forward activation tensor and receive backward activation tensor. Args: output_dict: A dictionary mapping module names to te

Send output to the next pipeline stage. Args: tensor: Tensor to send pp_group: Custom process group (if None, uses global state)

Converts the instance into a serializable dictionary. Returns: (dict) A dictionary representation of the instance suitable for

Set the current microbatch.

Toggle MoE drop-and-pad for decode. Applies ``capacity_factor`` to the router and all token dispatchers so decode runs with fixed shapes

Update a Tensor item managed by the `DataParallelBuffer` instance, i.e. store (a shard of) the Tensor in this buffer's datastore.

Method to set the operating mode

Utility function to set the DynamicInferenceContext rounder.

Set symmetric all-reduce functionality across all eligible modules. This method traverses the model's module hierarchy to find all m

Set sequence_sr for the grad_fn of a tensor to control the backward order. For older PyTorch version, do nothing (backward order is not chang

Sets up embedding layer in first stage and output layer in last stage. This function initalizes word embeddings in the final stage when we ar

Gets the embedding weight or output logit weights when share input embedding and output weights set to True or when use Multi-Token Prediction

Block forward pre-hook for certain configurations.

Shutdown NVML.

Pause CUDA memory for all tracked tensors.

缩放后图片满足以下条件: 1. 长宽能被 factor 整除 2. pixels 总数被限制在 [min_pixels, max_pixels] 内 3. 最长边限制在 max_long_side 内 4. 保证其长宽比基本不变

Initializes ZMQ communication to connect the engine with an inference coordinator. This asynchronous method sets up the distributed communica

The state dict contains all non-DP-rank-dependent (i.e., non-parameter- related) optimizer variables. The returned state dict can be

For easy load.