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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101470
+Version: 0.16.0rc0.dev101524
 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, and distillation
+- **[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/).
 
 **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.dev101470 → megatron_core-0.16.0rc0.dev101524}/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, and distillation
+- **[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/).
 
 **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.dev101470 → megatron_core-0.16.0rc0.dev101524}/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2782,6 +2782,9 @@ 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.dev101470 → megatron_core-0.16.0rc0.dev101524}/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py

@@ -25,6 +25,8 @@ from torch.distributed.checkpoint.metadata import (
 from torch.distributed.checkpoint.planner import TensorWriteData, WriteItem, WriteItemType
 from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard
 
+from .utils import get_mesh_names
+
 
 def gather_and_compute_chunk_metadata(dtensor: DTensor) -> ChunkStorageMetadata:
     """
@@ -272,7 +274,14 @@ def gather_uneven_dtensor_to_full_tensor(
     if not device_mesh.mesh_dim_names:
         process_group = device_mesh.get_group()
     else:
-        process_group = device_mesh._flatten().get_group()
+        # Check if the fully-flattened mesh exists first.
+        full_flattened_mesh_dim_name = "_".join(device_mesh.mesh_dim_names)
+        if full_flattened_mesh_dim_name in get_mesh_names(device_mesh):
+            # Retrieve the existing flattened DeviceMesh ProcessGroup.
+            process_group = device_mesh[full_flattened_mesh_dim_name].get_group()
+        else:
+            # Create the _-separated flattened DeviceMesh ProcessGroup.
+            process_group = device_mesh._flatten().get_group()
 
     # Collect chunk metadata for uneven shards (update if missing)
     if not hasattr(dtensor._local_tensor, "__create_chunk_list__"):

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

@@ -167,13 +167,10 @@ def get_mesh_names(device_mesh: Optional[DeviceMesh] = None) -> list[str]:
         submesh_dim_name
         for child_mesh, root_mesh in _mesh_resources.child_to_root_mapping.items()
         for submesh_dim_name in (child_mesh.mesh_dim_names or [])
-        if root_mesh == device_mesh
+        # Add flattened or other unaccounted for children of the root mesh.
+        if root_mesh == device_mesh and submesh_dim_name not in mesh_dim_names
     ]
-    # Combine without duplicate dimensions.
-    for dim_name in submesh_dim_names:
-        if dim_name not in mesh_dim_names:
-            mesh_dim_names.append(dim_name)
-    return mesh_dim_names
+    return mesh_dim_names + submesh_dim_names
 
 
 def contains_submesh(

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

@@ -48,6 +48,28 @@ 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.dev101470 → megatron_core-0.16.0rc0.dev101524}/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 experimental_fn, null_decorator
+from megatron.core.utils import null_decorator
 
 try:
     import triton
@@ -279,7 +279,6 @@ 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.dev101470 → megatron_core-0.16.0rc0.dev101524}/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 experimental_fn, null_decorator
+from megatron.core.utils import null_decorator
 
 try:
     import triton
@@ -324,7 +324,6 @@ 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,
@@ -733,7 +732,6 @@ 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.dev101470 → megatron_core-0.16.0rc0.dev101524}/megatron/core/fusions/fused_pad_routing_map.py

@@ -1,10 +1,12 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 from unittest.mock import MagicMock
 
 import torch
 from packaging import version
 
-from megatron.core.utils import experimental_fn, null_decorator
+from megatron.core.jit import jit_fuser
+from megatron.core.utils import null_decorator
 
 try:
     import triton
@@ -68,7 +70,7 @@ def _pad_routing_map_kernel(
     tl.store(output_row_ptr + token_indices, output_row, mask=token_mask)
 
 
-@experimental_fn(introduced_with_version="0.13.0")
+@jit_fuser
 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.dev101470 → megatron_core-0.16.0rc0.dev101524}/megatron/core/inference/communication_utils.py

@@ -71,8 +71,7 @@ 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(tensor)
-        assert tensor.is_contiguous()
+        _is_cuda_contiguous(tensor)
     else:
         tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
 

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

@@ -1,8 +1,28 @@
 # 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."""
 
@@ -64,7 +84,7 @@ class MambaMetadata:
         """
         self.request_to_mamba_state_idx_cudagraph_only[0:num_active_requests] = active_mamba_indices
 
-    def allocate_slot(self) -> int:
+    def allocate_slot(self) -> Optional[int]:
         """
         Allocates a new slot for a request in the Mamba state buffers.
 

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

@@ -0,0 +1,119 @@
+# 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
+
+    Args:
+        context (DynamicInferenceContext): Dynamic inference context.
+        active_count (int): Total number of active blocks available in the buffer.
+            The full buffer size is 2*active_count, to accommodate an equal-size
+            space for paused requests that live on the CPU.
+    """
+
+    def __init__(self, context: "DynamicInferenceContext", total_count: int):
+
+        self.context = context
+
+        active_count = (total_count - 1) // 2  # -1 for dummy_block_idx (see below)
+        active_count = max(1, active_count)  # need at least one block
+        self.total_count = 2 * active_count + 1  # +1 for dummy_block_idx
+        self.total_avail = self.total_count - 1  # -1 for dummy_block_idx
+        self.active_count = active_count
+        self.paused_count = self.total_count - self.active_count - 1  # -1 for dummy_block_idx
+        self.dummy_block_idx = self.total_count - 1
+
+        # Initialize block pool as a "stack" data structure
+        self.block_bag = torch.arange(
+            self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+    def __str__(self):
+        return (
+            f"total avail {self.total_avail} / {self.total_count - 1}"
+            f"; active {self.active_count}"
+        )
+
+    def get_active_used(self):
+        """Compute number of active blocks used."""
+        return (
+            self.context.request_kv_block_counts[
+                self.context.paused_request_count : self.context.total_request_count
+            ]
+            .sum()
+            .item()
+        )
+
+    def get_paused_used(self):
+        """Compute number of paused blocks used."""
+        return (
+            self.context.request_kv_block_counts[: self.context.paused_request_count].sum().item()
+        )
+
+    def get_active_avail(self):
+        """Compute number of active blocks available."""
+        return self.active_count - self.get_active_used()
+
+    def get_paused_avail(self):
+        """Compute number of paused blocks available."""
+        return self.paused_count - self.get_paused_used()
+
+    def is_memory_available(self, num_blocks: int) -> bool:
+        """Check if memory blocks are available.
+
+        Args:
+            num_blocks (int): Number of blocks to check.
+
+        Return:
+            (bool) Is memory available?
+        """
+        return self.get_active_avail() >= num_blocks
+
+    def allocate_memory_blocks(self, num_blocks: int) -> Optional[Tensor]:
+        """Allocate memory blocks if available, else return None.
+
+        Args:
+            num_blocks (int): Number of blocks to allocate.
+
+        Return:
+            (Optional[Tensor]) Allocated block IDs.
+        """
+        if self.is_memory_available(num_blocks):
+            self.total_avail -= num_blocks
+            block_ids = self.block_bag[self.total_avail : (self.total_avail + num_blocks)]
+            assert num_blocks == block_ids.numel()
+            return block_ids
+        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.total_avail : (self.total_avail + num_blocks)] = blocks
+        self.total_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.total_avail = self.total_count - 1