megatron-core 0.14.0rc2__tar.gz → 0.14.0rc4__tar.gz

{megatron_core-0.14.0rc2/megatron_core.egg-info → megatron_core-0.14.0rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.14.0rc2
+Version: 0.14.0rc4
 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>
@@ -31,6 +31,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch
 Requires-Dist: numpy<2.0.0
+Requires-Dist: packaging~=25.0
 Provides-Extra: mlm
 Requires-Dist: flask-restful; extra == "mlm"
 Requires-Dist: sentencepiece; extra == "mlm"
@@ -38,14 +39,16 @@ Requires-Dist: tiktoken; extra == "mlm"
 Requires-Dist: wandb; extra == "mlm"
 Provides-Extra: dev
 Requires-Dist: tqdm; extra == "dev"
-Requires-Dist: einops; extra == "dev"
-Requires-Dist: tensorstore!=0.1.46,!=0.1.72; extra == "dev"
-Requires-Dist: nvtx; extra == "dev"
-Requires-Dist: transformers; extra == "dev"
-Requires-Dist: multi-storage-client; extra == "dev"
+Requires-Dist: einops~=0.8; extra == "dev"
+Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "dev"
+Requires-Dist: nvtx~=0.2; extra == "dev"
+Requires-Dist: transformers~=4.53; extra == "dev"
+Requires-Dist: multi-storage-client~=0.20.3; extra == "dev"
+Requires-Dist: opentelemetry-api~=1.33.1; extra == "dev"
 Requires-Dist: setuptools<80.0.0; extra == "dev"
-Requires-Dist: nvidia-modelopt[torch]~=0.31.0; sys_platform != "darwin" and extra == "dev"
-Requires-Dist: megatron-energon[av_decode]<7; extra == "dev"
+Requires-Dist: nvidia-modelopt[torch]<0.32.0,>=0.31.0a0; sys_platform != "darwin" and extra == "dev"
+Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "dev"
+Requires-Dist: flashinfer-python; extra == "dev"
 Provides-Extra: lts
 Requires-Dist: tqdm; extra == "lts"
 Requires-Dist: einops; extra == "lts"

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/dist_checkpointing/strategies/torch.py

@@ -374,7 +374,8 @@ def _unwrap_pyt_sharded_tensor(sh_ten: TorchShardedTensor) -> List[torch.Tensor]
             ten = ten.view(-1)
         else:
             for _ in range(mcore_sh_ten.prepend_axis_num):
-                ten = ten.squeeze(0)
+                assert ten.size(0) == 1
+                ten = ten[0]  # NOTE: ten.squeeze(0) uses more memory for FP8 tensors
         ret_tensors.append(ten)
     return ret_tensors
 

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/dist_checkpointing/validation.py

@@ -375,28 +375,34 @@ def maybe_report_missing_and_unexpected_keys(
 def _validate_common_state_dict(common_state_dict: CommonStateDict) -> None:
     """Validate consistancy across ranks for the common state dict
 
-    We save the common state dict only on rank 0. We validate to make sure that the common dict is consistant across ranks before saving.
+    We save the common state dict only on rank 0. We validate to make sure that the common dict is consistent across ranks before saving.
 
     Args:
         common_state_dict: The common state dict present in all ransk
     """
+    if not torch.distributed.is_initialized():
+        return
 
-    # Gather the common state dict across ranks onto rank 0 for comparison
+    # Broadcast the common state dict from rank 0 to all other ranks
+    # Each rank will do a comparison with its local rank vs the broadcasted state dict from rank 0
     rank = torch.distributed.get_rank()
-    other_rank_state_dicts = [None] * torch.distributed.get_world_size() if rank == 0 else None
-    torch.distributed.gather_object(common_state_dict, other_rank_state_dicts)
-    common_state_dict_diff = {}
-    if rank == 0:
-        assert other_rank_state_dicts
-        main_rank_state_dict = common_state_dict
-        for rank, rank_state_dict in enumerate(other_rank_state_dicts[1:], 1):
-            only_left, only_right, mismatch = diff(main_rank_state_dict, rank_state_dict)
-            if only_left or only_right or mismatch:
-                common_state_dict_diff[rank] = (only_left, only_right, mismatch)
-
-        if len(common_state_dict_diff) != 0:
+
+    object_list = [common_state_dict] if rank == 0 else [None]
+    torch.distributed.broadcast_object_list(object_list, src=0)
+    rank0_state_dict = object_list[0]
+
+    # Skip comparing rank 0 with itself
+    if rank > 0:
+        current_rank_state_dict = common_state_dict
+        only_in_rank0, only_in_current_rank, mismatch = diff(
+            rank0_state_dict, current_rank_state_dict
+        )
+        if only_in_rank0 or only_in_current_rank or mismatch:
             logger.warning(
-                f"There is difference in the common state dict in different ranks. The differences are {common_state_dict_diff}"
+                f"Rank {rank} common state dict differs from rank 0 common state dict. "
+                f"Keys only on rank 0: {only_in_rank0}, "
+                f"Keys only on {rank}: {only_in_current_rank}, "
+                f"Mismatched keys: {mismatch}"
             )
 
 

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/distributed/custom_fsdp/fully_sharded_data_parallel.py

@@ -217,6 +217,16 @@ class FullyShardedDataParallel(_BaseDataParallel):
 
         self.module.apply(unmap_weight_tensor)
 
+        for param in self.module.parameters():
+            if not hasattr(param, 'grad_added_to_main_grad'):
+                # This is to ensure that the param.grad_added_to_main_grad is set to False
+                # when the parameter is created.
+                param.grad_added_to_main_grad = False
+            if not hasattr(param, '__fsdp_param__'):
+                # This is to ensure that the param.__fsdp_param__ is set to True
+                # when the parameter is created.
+                param.__fsdp_param__ = True
+
     def _init_fsdp_param_and_grad_buffer(self):
         if self.config.calculate_per_token_loss:
             # We don't need to scale the gradients in this case.

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/extensions/transformer_engine.py

@@ -889,25 +889,7 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
             if packed_seq_params is not None
             else {}
         )
-        # overwrite self.qkv_format depending on self.config.apply_rope_fusion, which can be set
-        # after init
-        if self.config.apply_rope_fusion and is_te_min_version("0.13.0", check_equality=False):
-            self.qkv_format = "bshd"
-
-        qkv_format = packed_seq_kwargs.get("qkv_format", self.qkv_format)
-
-        # WAR for peak memory usage.
-        # See https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/2388
-        if self.config.apply_rope_fusion and qkv_format == "bshd":
-            query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
-            # In PyTorch, the following two tensors are in fact the same:
-            #   Tensor with shape (1, S, H, D) and stride (S*H*D, H*D, D, 1)
-            #   Tensor with shape (1, S, H, D) and stride (H*D, H*D, D, 1)
-            # Stride for a dimension that is 1 has no meaning, so tensors created two different ways
-            # can have same shape but different strides.
-            # We unify them to the first one to pass the stride check in TE
-            if value.shape == key.shape and value.shape[0] == 1 and value.stride() != key.stride():
-                value = value.as_strided(value.shape, key.stride())
+        qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
 
         attention_bias_kwargs = {}
         if attention_bias is not None:
@@ -942,10 +924,7 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
                 query, key, value, attention_mask, **attention_bias_kwargs, **packed_seq_kwargs
             )
 
-        if self.config.apply_rope_fusion and qkv_format == "bshd":
-            return core_attn_out.transpose(0, 1)
-        else:
-            return core_attn_out
+        return core_attn_out
 
 
 if HAVE_TE and is_te_min_version("1.9.0.dev0"):
@@ -1633,10 +1612,8 @@ try:
         else:
             if interleaved:
                 raise ValueError("Only TE >= 2.3.0 supports interleaved fused RoPE.")
-            if is_te_min_version("1.4.0.dev0"):
-                return apply_rotary_pos_emb(t, freqs, tensor_format="sbhd", fused=True)
-            else:
-                raise ValueError("Only TE >= 1.4.0.dev0 supports fused RoPE.")
+
+            return apply_rotary_pos_emb(t, freqs, tensor_format="sbhd", fused=True)
 
     def fused_apply_rotary_pos_emb_thd(
         t: torch.Tensor,
@@ -1659,6 +1636,7 @@ try:
                 cp_rank=cp_rank,
             )
         else:
+            assert cp_size == 1, "Only TE >= 1.12 supports RoPE fusion for THD format with CP."
             return apply_rotary_pos_emb(
                 t, freqs, tensor_format="thd", fused=True, cu_seqlens=cu_seqlens
             )

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/extensions/transformer_engine_spec_provider.py

@@ -8,6 +8,7 @@ from megatron.core.extensions.transformer_engine import (
     TEColumnParallelLinear,
     TEDotProductAttention,
     TELayerNormColumnParallelLinear,
+    TELinear,
     TENorm,
     TERowParallelGroupedLinear,
     TERowParallelLinear,
@@ -23,6 +24,10 @@ from megatron.core.utils import get_te_version, is_te_min_version
 class TESpecProvider(BackendSpecProvider):
     """A protocol for providing the submodules used in Spec building."""
 
+    def linear(self) -> type:
+        """Which linear module TE backend uses"""
+        return TELinear
+
     def column_parallel_linear(self) -> type:
         """Which column parallel linear module TE backend uses"""
         return TEColumnParallelLinear

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/fp8_utils.py

@@ -2,7 +2,9 @@
 
 """Utility functions related to FP8 that are used throughout Megatron core"""
 
+import weakref
 from contextlib import nullcontext
+from functools import wraps
 from typing import List, Optional
 
 import torch
@@ -53,6 +55,29 @@ except (ImportError, ModuleNotFoundError):
     # MXFP8Tensor not found
     HAVE_TE_MXFP8TENSOR = False
 
+if HAVE_TE:
+    from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelLinear,
+        TELayerNormColumnParallelLinear,
+        TELinear,
+        TERowParallelLinear,
+    )
+
+    TE_LINEAR_TYPES = (
+        TELinear,
+        TEColumnParallelLinear,
+        TERowParallelLinear,
+        TELayerNormColumnParallelLinear,
+    )
+else:
+    TE_LINEAR_TYPES = ()
+
+try:
+    from megatron.core.extensions.transformer_engine import Fp8Padding, Fp8Unpadding
+except ImportError:
+    Fp8Padding = None
+    Fp8Unpadding = None
+
 
 def is_float8tensor(tensor: torch.Tensor) -> bool:
     """Check if a tensor is a Transformer Engine Float8Tensor.
@@ -511,3 +536,97 @@ else:
     def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
         """Returns dummy fp8 context manager since TE is not available."""
         return nullcontext()
+
+
+if HAVE_TE:
+    from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+
+    # Modules that have been wrapped for inference for fp8
+    _fp8_inference_wrapped_modules = weakref.WeakSet()
+
+    def _wrap_te_linear_for_padding(module: torch.nn.Module):
+        """Wrap a TE linear module to automatically pad sequences for FP8 inference.
+
+        Modifies the module's forward method to:
+        1. Pad input sequences to FP8 alignment requirements
+        2. Run the original forward pass
+        3. Unpad outputs to original sequence length
+
+        Args:
+            module: A Transformer Engine linear layer (TELinear, TEColumnParallelLinear, etc.)
+        """
+        if module in _fp8_inference_wrapped_modules:
+            return
+        _pad_func = Fp8Padding(1)
+        _unpad_func = Fp8Unpadding(1)
+
+        original_forward = module.forward
+
+        @wraps(original_forward)
+        def padded_forward(input_tensor, *args, **kwargs):
+            # Only do padding for fp8 if we are in fp8 context
+            if not FP8GlobalStateManager.is_fp8_enabled():
+                return original_forward(input_tensor, *args, **kwargs)
+
+            seq_len, batch_size, hidden_size = input_tensor.shape
+            # Reshape to (S, B*H) to pad sequence dimension
+            input_2d = input_tensor.reshape(seq_len, -1)
+            # Pad the sequence dimension
+            padded_input_2d, _ = _pad_func(input_2d, [seq_len])
+            padded_seq_len = padded_input_2d.shape[0]
+
+            # Reshape back to (padded_S, B, H)
+            padded_input_3d = padded_input_2d.view(padded_seq_len, batch_size, hidden_size)
+            output = original_forward(padded_input_3d, *args, **kwargs)
+
+            # Handle output
+            if isinstance(output, tuple):
+                output_tensor = output[0]
+                other_outputs = output[1:]
+            else:
+                output_tensor = output
+                other_outputs = ()
+
+            # Unpad output - reshape to 2D, unpad, reshape back
+            _, _, output_hidden_size = output_tensor.shape
+            output_2d = output_tensor.reshape(padded_seq_len, -1)
+            unpadded_output_2d = _unpad_func(output_2d, [seq_len])
+            unpadded_output = unpadded_output_2d.reshape(seq_len, batch_size, output_hidden_size)
+
+            if other_outputs:
+                return (unpadded_output,) + other_outputs
+            else:
+                return unpadded_output
+
+        module.forward = padded_forward
+        _fp8_inference_wrapped_modules.add(module)
+
+    def prepare_model_for_fp8_inference(model):
+        """Prepare a model for FP8 inference by wrapping TE linear layers with padding support.
+
+        FP8 TE Gemms have specific shape requirements. This function wraps all Transformer
+        Engine linear layers in the model to automatically pad/unpad sequences during inference.
+
+        Args:
+            model (model (GPTModel): Model containing TE linear layers.
+
+        Returns:
+            GPTModel: The same model with wrapped linear layers (modified in-place).
+
+        """
+        assert Fp8Padding and Fp8Unpadding, "TE version does not have FP8 padding functions"
+        # Find and wrap all TE linear layers
+        for module in model.modules():
+            if isinstance(module, TE_LINEAR_TYPES):
+                _wrap_te_linear_for_padding(module)
+
+        return model
+
+else:
+
+    def prepare_model_for_fp8_inference(model):
+        """If trys using prepare_model_for_fp8_inference without TE we error"""
+        raise RuntimeError(
+            "prepare_model_for_fp8_inference requires Transformer Engine to be installed. "
+            "Please install transformer-engine to use FP8 inference."
+        )

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -28,16 +28,25 @@ if not HAVE_TRITON:
 
 
 @triton.jit
-def _get_thd_token_idx(cu_seqlens, pid_m, seq_num):
+def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size):
     token_idx = -1
+    this_seq_len = 0
     seq_idx = 0
-    last_cum_seqlen = tl.load(cu_seqlens)
+    last_cum_seqlen = tl.load(cu_seqlens) // cp_size
     while seq_idx < seq_num:
-        cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1)
+        cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size
         if token_idx == -1 and cur_cum_seqlen > pid_m:
             token_idx = pid_m - last_cum_seqlen
+            this_seq_len = cur_cum_seqlen - last_cum_seqlen
         last_cum_seqlen = cur_cum_seqlen
         seq_idx += 1
+    if cp_size > 1:
+        if token_idx < this_seq_len // 2:
+            token_idx = token_idx + cp_rank * this_seq_len // 2
+        else:
+            token_idx = (token_idx - this_seq_len // 2) + (
+                2 * cp_size - cp_rank - 1
+            ) * this_seq_len // 2
     return token_idx
 
 
@@ -68,6 +77,8 @@ def rotary_fwd_q_kernel(
     cu_seqlens_q,
     stride_x_seq,
     stride_x_nheads,
+    cp_rank,
+    cp_size,
     BLOCK_H: tl.constexpr,
 ):
     """
@@ -89,7 +100,7 @@ def rotary_fwd_q_kernel(
     if cu_seqlens_q is None:
         token_idx = pid_m // batch_size
     else:
-        token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num)
+        token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size)
 
     cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
     sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
@@ -146,6 +157,8 @@ def rotary_bwd_q_kernel(
     cu_seqlens_q,
     stride_x_seq,
     stride_x_nheads,
+    cp_rank,
+    cp_size,
     BLOCK_H: tl.constexpr,
 ):
     """
@@ -165,7 +178,7 @@ def rotary_bwd_q_kernel(
     if cu_seqlens_q is None:
         token_idx = pid_m // batch_size
     else:
-        token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num)
+        token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size)
 
     cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
     sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
@@ -200,7 +213,18 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
     """
 
     @staticmethod
-    def forward(ctx, q, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, rotary_interleaved=False):
+    def forward(
+        ctx,
+        q,
+        cos,
+        sin,
+        qk_head_dim,
+        emb_dim,
+        cu_seqlens_q,
+        cp_rank,
+        cp_size,
+        rotary_interleaved=False,
+    ):
         """
         Forward function for ApplyMLARotaryEmbQ.
 
@@ -243,12 +267,16 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
             cu_seqlens_q,
             q.stride(0),
             q.stride(1),
+            cp_rank,
+            cp_size,
         )
         ctx.save_for_backward(cos, sin)
         ctx.qk_head_dim = qk_head_dim
         ctx.emb_dim = emb_dim
         ctx.cu_seqlens_q = cu_seqlens_q
         ctx.rotary_interleaved = rotary_interleaved
+        ctx.cp_rank = cp_rank
+        ctx.cp_size = cp_size
         if cu_seqlens_q is None:
             q = q.view(max_seqlen, batch_size, nheads, headdim)
         return q
@@ -268,7 +296,7 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
         seq_num = None
         if ctx.cu_seqlens_q is None:
             max_seqlen, batch_size, nheads, headdim = grad.shape
-            grad = grad.view(-1, nheads, headdim)
+            grad = grad.contiguous().view(-1, nheads, headdim)
             total_seqlen = grad.shape[0]
         else:
             seq_num = len(ctx.cu_seqlens_q) - 1
@@ -288,10 +316,12 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
             ctx.cu_seqlens_q,
             grad.stride(0),
             grad.stride(1),
+            ctx.cp_rank,
+            ctx.cp_size,
         )
         if ctx.cu_seqlens_q is None:
             grad = grad.view(max_seqlen, batch_size, nheads, headdim)
-        return grad, None, None, None, None, None, None
+        return grad, None, None, None, None, None, None, None, None
 
 
 @experimental_fn(introduced_with_version="0.13.0")
@@ -302,6 +332,8 @@ def fused_apply_mla_rope_for_q(
     qk_head_dim: int,
     emb_dim: int,
     cu_seqlens_q: Optional[torch.Tensor] = None,
+    cp_rank: int = 0,
+    cp_size: int = 1,
     rotary_interleaved: bool = False,
 ):
     """
@@ -327,7 +359,7 @@ def fused_apply_mla_rope_for_q(
         t: inplace modified input tensor
     """
     return ApplyMLARotaryEmbQ.apply(
-        t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, rotary_interleaved
+        t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved
     )
 
 
@@ -366,6 +398,8 @@ def rotary_fwd_kv_kernel(
     stride_k_nheads,
     stride_v_seq,
     stride_v_nheads,
+    cp_rank,
+    cp_size,
     BLOCK_H: tl.constexpr,
 ):
     """
@@ -394,7 +428,7 @@ def rotary_fwd_kv_kernel(
     if cu_seqlens_kv is None:
         token_idx = pid_m // batch_size
     else:
-        token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num)
+        token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size)
 
     cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
     sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2))
@@ -472,6 +506,8 @@ def rotary_bwd_kv_kernel(
     stride_dkv_seq,
     stride_dkv_nheads,
     stride_demb_seq,
+    cp_rank,
+    cp_size,
     BLOCK_H: tl.constexpr,
 ):
     """
@@ -496,7 +532,7 @@ def rotary_bwd_kv_kernel(
     if cu_seqlens_kv is None:
         token_idx = pid_m // batch_size
     else:
-        token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num)
+        token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size)
 
     dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads
     dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads
@@ -550,7 +586,18 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
 
     @staticmethod
     def forward(
-        ctx, kv, k_pos_emb, cos, sin, emb_dim, k_dim, v_dim, cu_seqlens_kv, rotary_interleaved=False
+        ctx,
+        kv,
+        k_pos_emb,
+        cos,
+        sin,
+        emb_dim,
+        k_dim,
+        v_dim,
+        cu_seqlens_kv,
+        cp_rank,
+        cp_size,
+        rotary_interleaved=False,
     ):
         """
         Forward function for ApplyMLARotaryEmbKV.
@@ -609,6 +656,8 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
             o_key.stride(1),
             o_value.stride(0),
             o_value.stride(1),
+            cp_rank,
+            cp_size,
         )
         ctx.save_for_backward(cos, sin)
         ctx.rotary_interleaved = rotary_interleaved
@@ -616,6 +665,8 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
         ctx.k_dim = k_dim
         ctx.v_dim = v_dim
         ctx.cu_seqlens_kv = cu_seqlens_kv
+        ctx.cp_rank = cp_rank
+        ctx.cp_size = cp_size
         if cu_seqlens_kv is None:
             o_key = o_key.view(max_seqlen, -1, nheads, emb_dim + k_dim)
             o_value = o_value.view(max_seqlen, -1, nheads, v_dim)
@@ -638,8 +689,8 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
         if ctx.cu_seqlens_kv is None:
             # sbhd
             max_seqlen, batch_size, nheads, _ = dk.shape
-            dk = dk.view(-1, nheads, ctx.emb_dim + ctx.k_dim)
-            dv = dv.view(-1, nheads, ctx.v_dim)
+            dk = dk.contiguous().view(-1, nheads, ctx.emb_dim + ctx.k_dim)
+            dv = dv.contiguous().view(-1, nheads, ctx.v_dim)
             total_seqlen = dk.shape[0]
         else:
             # thd
@@ -673,11 +724,13 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
             d_kv.stride(0),
             d_kv.stride(1),
             d_emb.stride(0),
+            ctx.cp_rank,
+            ctx.cp_size,
         )
         if ctx.cu_seqlens_kv is None:
             d_kv = d_kv.view(max_seqlen, batch_size, nheads, ctx.k_dim + ctx.v_dim)
             d_emb = d_emb.view(max_seqlen, batch_size, 1, ctx.emb_dim)
-        return d_kv, d_emb, None, None, None, None, None, None, None
+        return d_kv, d_emb, None, None, None, None, None, None, None, None, None
 
 
 @experimental_fn(introduced_with_version="0.13.0")
@@ -690,6 +743,8 @@ def fused_apply_mla_rope_for_kv(
     k_dim: int,
     v_dim: int,
     cu_seqlens_kv: Optional[torch.Tensor] = None,
+    cp_rank: int = 0,
+    cp_size: int = 1,
     rotary_interleaved: bool = False,
 ):
     """
@@ -715,5 +770,15 @@ def fused_apply_mla_rope_for_kv(
         value: [seq_len, batch_size, head_num, v_dim] or [total_seq_len, head_num, v_dim]
     """
     return ApplyMLARotaryEmbKV.apply(
-        kv, k_pos_emb, cos, sin, emb_dim, k_dim, v_dim, cu_seqlens_kv, rotary_interleaved
+        kv,
+        k_pos_emb,
+        cos,
+        sin,
+        emb_dim,
+        k_dim,
+        v_dim,
+        cu_seqlens_kv,
+        cp_rank,
+        cp_size,
+        rotary_interleaved,
     )

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc4}/megatron/core/inference/contexts/__init__.py

@@ -14,6 +14,7 @@ warnings.warn(
     DeprecationWarning,
 )
 from .dynamic_context import (
+    ActiveRequestCountOverflowError,
     ChunkOverflowError,
     ContextOverflowError,
     DynamicInferenceContext,