megatron-core 0.16.0rc0.dev101399__tar.gz → 0.16.0rc0.dev101470__tar.gz

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101399
+Version: 0.16.0rc0.dev101470
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>
@@ -175,7 +175,7 @@ pip install --no-build-isolation .[mlm,dev]
 
 ```
 Megatron-LM/
-├── megatron/
+├── megatron/                    
 │   ├── core/                    # Megatron Core (kernels, parallelism, building blocks)
 │   │   ├── models/              # Transformer models
 │   │   ├── transformer/         # Transformer building blocks
@@ -204,7 +204,7 @@ Megatron-LM/
 
 - **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
 - **Research teams** exploring new architectures and training techniques
-- **Learning distributed training** concepts and best practices
+- **Learning distributed training** concepts and best practices  
 - **Quick experimentation** with proven model configurations
 
 **What you get:**
@@ -213,7 +213,7 @@ Megatron-LM/
 - End-to-end examples from data prep to evaluation
 - Research-focused tools and utilities
 
-### Megatron Core: Composable Library
+### Megatron Core: Composable Library  
 
 **Composable library** with GPU-optimized building blocks for custom training frameworks.
 
@@ -246,7 +246,7 @@ Megatron-LM/
 - **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
 - **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
 - **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
-- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
+- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, and distillation
 
 **Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
 
@@ -333,7 +333,7 @@ Our codebase efficiently trains models from 2B to 462B parameters across thousan
 **Benchmark Configuration:**
 
 - **Vocabulary size**: 131,072 tokens
-- **Sequence length**: 4096 tokens
+- **Sequence length**: 4096 tokens  
 - **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts
 - **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/README.md

@@ -99,7 +99,7 @@ pip install --no-build-isolation .[mlm,dev]
 
 ```
 Megatron-LM/
-├── megatron/
+├── megatron/                    
 │   ├── core/                    # Megatron Core (kernels, parallelism, building blocks)
 │   │   ├── models/              # Transformer models
 │   │   ├── transformer/         # Transformer building blocks
@@ -128,7 +128,7 @@ Megatron-LM/
 
 - **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
 - **Research teams** exploring new architectures and training techniques
-- **Learning distributed training** concepts and best practices
+- **Learning distributed training** concepts and best practices  
 - **Quick experimentation** with proven model configurations
 
 **What you get:**
@@ -137,7 +137,7 @@ Megatron-LM/
 - End-to-end examples from data prep to evaluation
 - Research-focused tools and utilities
 
-### Megatron Core: Composable Library
+### Megatron Core: Composable Library  
 
 **Composable library** with GPU-optimized building blocks for custom training frameworks.
 
@@ -170,7 +170,7 @@ Megatron-LM/
 - **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
 - **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
 - **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
-- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
+- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, and distillation
 
 **Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
 
@@ -257,7 +257,7 @@ Our codebase efficiently trains models from 2B to 462B parameters across thousan
 **Benchmark Configuration:**
 
 - **Vocabulary size**: 131,072 tokens
-- **Sequence length**: 4096 tokens
+- **Sequence length**: 4096 tokens  
 - **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts
 - **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2782,9 +2782,6 @@ class GradReducePipeline:
             outer_fsdp_group_grad_reduce (bool, optional): Whether to reduce gradients
                 across outer-DP groups. Defaults to False.
         """
-        # Sort parameters by their bucket IDs to ensure a deterministic processing order.
-        # Performing reduce-scatter operations out of order can lead to hangs.
-        params = sorted(list(params), key=lambda x: self.buffer.param_to_param_group[x])
         for param in params:
             bucket_id = self.buffer.param_to_param_group[param]
             param_group = self.buffer.parameter_groups[bucket_id]

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/fp4_utils.py

@@ -48,28 +48,6 @@ def is_nvfp4tensor(tensor: torch.Tensor) -> bool:
     return HAVE_TE_FP4_TENSOR_CLASS and isinstance(tensor, FP4_TENSOR_CLASS)
 
 
-def get_fp4_align_size(fp4_recipe: Fp4Recipe) -> int:
-    """
-    Get the alignment size required for FP4 GEMM.
-    FP4 GEMM requires Blackwell and later architectures.
-
-    The value 32 is a hardware requirement: TMA (Tensor Memory Accelerator) requires
-    a 16-byte aligned address for efficient memory access. Since FP4 uses 4 bits per value,
-    16 bytes (128 bits) corresponds to 32 FP4 values. Therefore, the alignment size for FP4
-    is 32. With this alignment, NVFP4 GEMM can be performed efficiently.
-
-    Note that since we are also random hadamard transform for NVFP4 training, we want
-    fused group nvfp4 quantize plus hadamard transform. Hadamard transform will leverage
-    tensor core instructions for better performance, while group quantize kernels also
-    prefer a more aligned size in token dimension M. Therefore, we apply align size 64
-    here for better performance in MOE.
-
-    Paper link: https://arxiv.org/pdf/2509.25149
-    """
-    # pylint: disable=unused-argument
-    return 64
-
-
 def dequantize_fp4_tensor(fp4_tensor: torch.Tensor) -> torch.Tensor:
     """Dequantize a fp4 tensor to a higher precision tensor."""
     if is_te_min_version("2.7.0.dev0"):

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/fusions/fused_indices_converter.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock
 import torch
 from packaging import version
 
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -279,6 +279,7 @@ class IndicesToMultihot(torch.autograd.Function):
         return None, grad_probs_indices, None, None
 
 
+@experimental_fn(introduced_with_version='0.11.0rc0')
 def fused_indices_to_multihot(indices, probs_indices, num_of_local_experts):
     """Convert moe topk indices to multihot representation.
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock
 import torch
 from packaging import version
 
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -324,6 +324,7 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
         return grad, None, None, None, None, None, None, None, None
 
 
+@experimental_fn(introduced_with_version="0.13.0")
 def fused_apply_mla_rope_for_q(
     t: torch.Tensor,
     cos: torch.Tensor,
@@ -732,6 +733,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
         return d_kv, d_emb, None, None, None, None, None, None, None, None, None
 
 
+@experimental_fn(introduced_with_version="0.13.0")
 def fused_apply_mla_rope_for_kv(
     kv: torch.Tensor,
     k_pos_emb: torch.Tensor,

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/fusions/fused_pad_routing_map.py

@@ -1,12 +1,10 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 from unittest.mock import MagicMock
 
 import torch
 from packaging import version
 
-from megatron.core.jit import jit_fuser
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -70,7 +68,7 @@ def _pad_routing_map_kernel(
     tl.store(output_row_ptr + token_indices, output_row, mask=token_mask)
 
 
-@jit_fuser
+@experimental_fn(introduced_with_version="0.13.0")
 def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
     """Fused version of pad_routing_map.
     Args:

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/inference/communication_utils.py

@@ -71,7 +71,8 @@ def broadcast_from_last_pipeline_stage(
             tensor.shape
         ), f"Expected tensor of shape {size} but got {list(tensor.shape)}"
         assert dtype == tensor.dtype, f"Expected tensor of type {dtype} but got {tensor.dtype}"
-        _is_cuda_contiguous(tensor)
+        _is_cuda(tensor)
+        assert tensor.is_contiguous()
     else:
         tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101470}/megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -1,28 +1,8 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
-
 import torch
 
 
-@dataclass
-class MambaInferenceStateConfig:
-    """Config for initializing Mamba model inference state tensors."""
-
-    layer_type_list: List[str]
-    """
-    A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
-    See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols.
-    """
-
-    mamba_conv_states_shape: Tuple[int]
-    """Mamba conv states shape per request."""
-
-    mamba_ssm_states_shape: Tuple[int]
-    """Mamba ssm states shape per request."""
-
-
 class MambaMetadata:
     """Manages the metadata tensors required for Mamba layers during inference."""
 
@@ -84,7 +64,7 @@ class MambaMetadata:
         """
         self.request_to_mamba_state_idx_cudagraph_only[0:num_active_requests] = active_mamba_indices
 
-    def allocate_slot(self) -> Optional[int]:
+    def allocate_slot(self) -> int:
         """
         Allocates a new slot for a request in the Mamba state buffers.
 

megatron_core-0.16.0rc0.dev101470/megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+class BlockAllocator:
+    """Allocator that manages blocks of memory for the KV cache.
+
+    This allocator is responsible for:
+    - Initializing a pool of block IDs
+    - Allocating blocks from the pool
+    - Releasing blocks back to the pool
+    - Managing the guaranteed block count for active requests
+
+    Args:
+        block_count_total (int): Total number of blocks available in the buffer.
+        gtd_block_count (int): Number of blocks reserved for guaranteed requests.
+    """
+
+    def __init__(self, block_count_total: int, gtd_block_count: int):
+        self.block_count_total = block_count_total
+        self.gtd_block_count = gtd_block_count
+
+        # Reserve last block ID as dummy block for decode-only inference steps
+        self.block_count_avail = self.block_count_total - 1
+        self.dummy_block_idx = self.block_count_total - 1
+
+        # Initialize block pool as a "stack" data structure
+        self.block_bag = torch.arange(
+            self.block_count_total, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+    def is_memory_available(self, num_blocks: int, safe: bool = False) -> bool:
+        """Check if memory blocks are available.
+
+        Use 'safe' to avoid all requests being deadlocked. A fraction of the KV cache
+        memory buffer is reserved to guarantee that a minimum number of active
+        requests can run on any given step.
+
+        Args:
+            num_blocks (int): Number of blocks to check.
+            safe (bool): Include extra space for guaranteeing ability to run
+                requests to completion.
+
+        Return:
+            (bool) Is memory available?
+        """
+        if safe:
+            return self.block_count_avail >= num_blocks + self.gtd_block_count
+        else:
+            return self.block_count_avail >= num_blocks
+
+    def allocate_memory_blocks(self, num_blocks: int = 1, safe: bool = False) -> Optional[Tensor]:
+        """Allocate memory blocks if available, else return None.
+
+        Args:
+            num_blocks (int): Number of blocks to allocate.
+            safe (bool): Include extra space for guaranteeing ability to run
+                requests to completion.
+
+        Return:
+            (Optional[Tensor]) Allocated block IDs.
+        """
+        if self.is_memory_available(num_blocks, safe):
+            self.block_count_avail -= num_blocks
+            return self.block_bag[self.block_count_avail : (self.block_count_avail + num_blocks)]
+        else:
+            return None
+
+    def release_memory_blocks(self, blocks: Tensor) -> None:
+        """Release memory blocks.
+
+        Args:
+            blocks (Tensor): Block IDs to release.
+
+        Return:
+            None
+        """
+        num_blocks = blocks.size(dim=0)
+        self.block_bag[self.block_count_avail : (self.block_count_avail + num_blocks)] = blocks
+        self.block_count_avail += num_blocks
+
+    def reset(self) -> None:
+        """Reset the allocator to initial state.
+
+        This resets the available block count to the entire memory pool
+        (except for the dummy block).
+        """
+        self.block_count_avail = self.block_count_total - 1