[1/2] Enable DSA CP, THD with absorbed MLA

megatron/core/models/common/embeddings/rope_utils.py

@@ -220,49 +220,49 @@ def _apply_rotary_pos_emb_thd(
     cp_size = cp_group.size()
     cp_rank = cp_group.rank()
     seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()
+    sequence_splits = torch.split(t, seqlens)
+    total_seqlen = int(cu_seqlens[-1].item())
+    has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen
 
     # Handle two different frequency tensor formats:
     # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences
     #    -> Use offset-based mapping for exact positional correspondence
     # 2. Otherwise: freqs contains only max sequence length positions
     #    -> Use traditional mapping without offsets (map first :seqlen part)
-    if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]:
+    if has_packed_freqs:
         # CASE 1: Exact mapping with offsets
-        # Build packed freqs in one pass, then apply once to the whole packed tensor
-        sequence_splits = torch.split(t, seqlens)
-        freq_slices = []
+        local_freqs = []
         for i, x in enumerate(sequence_splits):
             # cu_seqlens[i] is the starting offset of this sequence in the original batch
             seq_start_offset = cu_seqlens[i].item()
-            freq_slices.append(
+            local_freqs.append(
                 _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset)
             )
-
-        freqs_packed = torch.cat(freq_slices, dim=0)
-
+        freqs = torch.cat(local_freqs, dim=0)
         return _apply_rotary_pos_emb_bshd(
             t.unsqueeze(1),
-            freqs_packed,
+            freqs,
             rotary_interleaved=rotary_interleaved,
             mla_rotary_interleaved=mla_rotary_interleaved,
             mscale=mscale,
         ).squeeze(1)
-    else:
-        # CASE 2: Traditional mapping without offsets
-        # Build packed freqs for all sequences using the standard mapping, then apply once
-        sequence_splits = torch.split(t, seqlens)
-        freqs_packed = torch.cat(
-            [_get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) for x in sequence_splits],
-            dim=0,
-        )
 
-        return _apply_rotary_pos_emb_bshd(
-            t.unsqueeze(1),
-            freqs_packed,
+    # CASE 2: Traditional mapping without offsets
+    output = torch.empty_like(t)
+    output_offset = 0
+    for x in sequence_splits:
+        freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs)
+        output_slice = _apply_rotary_pos_emb_bshd(
+            x.unsqueeze(1),
+            freq_slice,
             rotary_interleaved=rotary_interleaved,
             mla_rotary_interleaved=mla_rotary_interleaved,
             mscale=mscale,
         ).squeeze(1)
+        output.narrow(0, output_offset, x.size(0)).copy_(output_slice)
+        output_offset += x.size(0)
+
+    return output
 
 
 def apply_rotary_pos_emb(
@@ -283,6 +283,8 @@ def apply_rotary_pos_emb(
     # Keep for backward compatibility. Will deprecate in the future.
     if cp_group is None:
         cp_group = parallel_state.get_context_parallel_group()
+    if mla_rotary_interleaved is None:
+        mla_rotary_interleaved = config.multi_latent_attention
 
     if config.apply_rope_fusion:
         if cu_seqlens is None:

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -6,17 +6,17 @@
 from megatron.core.models.backends import BackendSpecProvider
 from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
 from megatron.core.transformer.enums import AttnMaskType, LayerType
+from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+    AbsorbedMLASelfAttention,
+    AbsorbedMLASelfAttentionSubmodules,
+)
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerSubmodules,
     DSAttention,
     DSAttentionSubmodules,
 )
 from megatron.core.transformer.identity_op import IdentityOp
-from megatron.core.transformer.multi_latent_attention import (
-    MLASelfAttention,
-    MLASelfAttentionSubmodules,
-)
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import (
     TransformerBlockSubmodules,
@@ -109,9 +109,9 @@ def get_dsa_module_spec_for_backend(
     )
 
     attention = ModuleSpec(
-        module=MLASelfAttention,
+        module=AbsorbedMLASelfAttention,
         params={"attn_mask_type": AttnMaskType.causal},
-        submodules=MLASelfAttentionSubmodules(
+        submodules=AbsorbedMLASelfAttentionSubmodules(
             linear_q_proj=backend.column_parallel_linear(),
             linear_q_down_proj=backend.linear(),
             linear_q_up_proj=backend.column_parallel_linear(),

megatron/core/models/hybrid/hybrid_layer_specs.py

@@ -26,6 +26,10 @@
 )
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+    AbsorbedMLASelfAttention,
+    AbsorbedMLASelfAttentionSubmodules,
+)
 from megatron.core.transformer.experimental_attention_variant.dsa import (
     DSAIndexer,
     DSAIndexerSubmodules,
@@ -135,9 +139,9 @@
             submodules=TransformerLayerSubmodules(
                 input_layernorm=TENorm,
                 self_attention=ModuleSpec(
-                    module=MLASelfAttention,
+                    module=AbsorbedMLASelfAttention,
                     params={"attn_mask_type": AttnMaskType.causal},
-                    submodules=MLASelfAttentionSubmodules(
+                    submodules=AbsorbedMLASelfAttentionSubmodules(
                         linear_q_proj=TEColumnParallelLinear,
                         linear_q_down_proj=TELinear,
                         linear_q_up_proj=TEColumnParallelLinear,

megatron/core/pipeline_parallel/hybrid_cp_schedule.py

@@ -545,9 +545,15 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             )
             sample["local_cp_size"] = torch.tensor(partner_cp_size, dtype=torch.int32)
             new_data_iterator = RerunDataIterator(iter([sample]))
-            return new_data_iterator
         else:
-            return None
+            partner_cp_size = 0
+            new_data_iterator = None
+
+        partner_cp_size_tensor = torch.tensor(
+            [partner_cp_size], dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        _broadcast(partner_cp_size_tensor)
+        return new_data_iterator, int(partner_cp_size_tensor.item())
 
     # We get data once per global batch and schedule the sub-samples.
     # TODO(pmannan): Should we wrap the data_iterator here instead of the training.py file?
@@ -579,7 +585,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             sample_ids_this_group = sample_id_groups[j][hdp_rank] if is_first_tp_rank else None
             for i in range(num_samples_this_group[j]):
                 # Call forward step for each sub-sample
-                new_data_iterator = _get_new_data_iterator(i, j)
+                new_data_iterator, cp_group_size = _get_new_data_iterator(i, j)
                 # TODO: Find the usage of current_microbatch and is_first_microbatch and
                 # how that may affect my usage.
                 output_tensor, num_tokens = forward_step(
@@ -590,7 +596,8 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                     input_tensor,
                     forward_data_store,
                     config,
-                    collect_non_loss_data,
+                    cp_group_size=cp_group_size,
+                    collect_non_loss_data=collect_non_loss_data,
                     is_first_microbatch=check_first_val_step(
                         first_val_step, forward_only, current_microbatch == 0
                     ),
@@ -599,9 +606,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                 current_microbatch += 1
                 total_num_tokens += num_tokens.item()
                 if not forward_only:
-                    backward_step(
-                        input_tensor, output_tensor, output_tensor_grad, model_type, config
-                    )
+                    backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
             # Create a barrier at end of each group.
             # This barrier ensures that all ranks are prepared to change assigned CP group sizes and
@@ -614,7 +619,7 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
     with no_sync_func():
         sample_ids_this_group = sample_id_groups[-1][hdp_rank] if is_first_tp_rank else None
         for i in range(num_samples_this_group[-1] - 1):
-            new_data_iterator = _get_new_data_iterator(i, -1)
+            new_data_iterator, cp_group_size = _get_new_data_iterator(i, -1)
             # Call forward step for each sub-sample
             output_tensor, num_tokens = forward_step(
                 forward_step_func,
@@ -624,7 +629,8 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
                 input_tensor,
                 forward_data_store,
                 config,
-                collect_non_loss_data,
+                cp_group_size=cp_group_size,
+                collect_non_loss_data=collect_non_loss_data,
                 is_first_microbatch=check_first_val_step(
                     first_val_step, forward_only, current_microbatch == 0
                 ),
@@ -633,11 +639,11 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
             current_microbatch += 1
             total_num_tokens += num_tokens.item()
             if not forward_only:
-                backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+                backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
     # The last sub-sample of the last group of the last microbatch is
     # run out of the context handler.
-    new_data_iterator = _get_new_data_iterator(-1, -1)
+    new_data_iterator, cp_group_size = _get_new_data_iterator(-1, -1)
     # Call forward step for each sub-sample
     output_tensor, num_tokens = forward_step(
         forward_step_func,
@@ -647,14 +653,15 @@ def _get_new_data_iterator(sample_id_in_group, group_id):
         input_tensor,
         forward_data_store,
         config,
-        collect_non_loss_data,
+        cp_group_size=cp_group_size,
+        collect_non_loss_data=collect_non_loss_data,
         is_first_microbatch=check_first_val_step(
             first_val_step, forward_only, current_microbatch == 0
         ),
         current_microbatch=current_microbatch,
     )
     total_num_tokens += num_tokens.item()
     if not forward_only:
-        backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+        backward_step(input_tensor, output_tensor, output_tensor_grad, config)
 
     return forward_data_store, total_num_tokens

megatron/core/pipeline_parallel/schedules.py

@@ -226,28 +226,58 @@ def get_tensor_device(tensor: Union[torch.Tensor, Dict[str, torch.Tensor]]):
     return tensor.device
 
 
-def _get_mtp_loss_scale(config, device: torch.device) -> torch.Tensor:
-    """Get the MTP loss scale on the output tensor device."""
+def _normalize_loss_scale(loss_scale, device: torch.device, scale_func_name: str) -> torch.Tensor:
+    """Normalize loss scale outputs to a size-1 tensor on the output tensor device."""
+    loss_scale = torch.as_tensor(loss_scale, device=device)
+    if loss_scale.numel() != 1:
+        raise ValueError(
+            f"{scale_func_name} must return a scalar or size-1 tensor for loss scaling, "
+            f"but returned a tensor with {loss_scale.numel()} elements."
+        )
+    return loss_scale
 
-    def _normalize_loss_scale(loss_scale, scale_func_name: str) -> torch.Tensor:
-        loss_scale = torch.as_tensor(loss_scale, device=device)
-        if loss_scale.numel() != 1:
-            raise ValueError(
-                f"{scale_func_name} must return a scalar or size-1 tensor for MTP loss scaling, "
-                f"but returned a tensor with {loss_scale.numel()} elements."
-            )
-        return loss_scale
 
-    mtp_grad_scale_func = getattr(config, 'mtp_grad_scale_func', None)
-    if mtp_grad_scale_func is not None:
-        return _normalize_loss_scale(mtp_grad_scale_func(), "mtp_grad_scale_func")
+def _compute_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Calculate the loss scale from grad_scale_func or default to 1."""
     if config.grad_scale_func is not None:
         return _normalize_loss_scale(
-            config.grad_scale_func(torch.ones(1, device=device)), "grad_scale_func"
+            config.grad_scale_func(torch.ones(1, device=device)), device, "grad_scale_func"
         )
     return torch.ones(1, device=device)
 
 
+def _get_moe_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Get the MoE loss scale on the output tensor device."""
+    moe_grad_scale_func = getattr(config, 'moe_grad_scale_func', None)
+    if moe_grad_scale_func is not None:
+        return _normalize_loss_scale(moe_grad_scale_func(), device, "moe_grad_scale_func")
+    return _compute_loss_scale(config, device)
+
+
+def _get_mtp_loss_scale(config, device: torch.device) -> torch.Tensor:
+    """Get the MTP loss scale on the output tensor device."""
+    mtp_grad_scale_func = getattr(config, 'mtp_grad_scale_func', None)
+    if mtp_grad_scale_func is not None:
+        return _normalize_loss_scale(mtp_grad_scale_func(), device, "mtp_grad_scale_func")
+    return _compute_loss_scale(config, device)
+
+
+def _get_experimental_attention_variant_loss_scale_func(config):
+    """Get the loss scale hook for experimental attention variants."""
+    loss_scale_func = getattr(config, 'experimental_attention_variant_loss_scale_func', None)
+    if loss_scale_func is not None:
+        return loss_scale_func
+
+    if getattr(config, 'experimental_attention_variant', None) == 'dsa':
+        from megatron.core.transformer.experimental_attention_variant.dsa import (
+            DSAIndexerLossAutoScaler,
+        )
+
+        return DSAIndexerLossAutoScaler.set_loss_scale
+
+    return None
+
+
 def forward_step_calc_loss(
     model,
     output_tensor,
@@ -262,9 +292,6 @@ def forward_step_calc_loss(
 ):
     """Calculate the loss and number of tokens for forward_step()"""
 
-    from megatron.core.transformer.experimental_attention_variant.dsa import (
-        DSAIndexerLossAutoScaler,
-    )
     from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler
 
     model_vp_stage = getattr(model, "vp_stage", None)
@@ -315,16 +342,8 @@ def forward_step_calc_loss(
     # Since we use a trick to do backward on the auxiliary loss, we need to set the scale
     # explicitly.
     if hasattr(config, 'num_moe_experts') and config.num_moe_experts is not None:
-        # Calculate the loss scale based on moe_grad_scale_func (preferred),
-        # grad_scale_func (fallback), or default to 1.
         device = get_tensor_device(output_tensor)
-        moe_grad_scale_func = getattr(config, 'moe_grad_scale_func', None)
-        if moe_grad_scale_func is not None:
-            loss_scale = moe_grad_scale_func()
-        elif config.grad_scale_func is not None:
-            loss_scale = config.grad_scale_func(torch.ones(1, device=device))
-        else:
-            loss_scale = torch.ones(1, device=device)
+        loss_scale = _get_moe_loss_scale(config, device)
         # Set the loss scale
         if config.calculate_per_token_loss:
             MoEAuxLossAutoScaler.set_loss_scale(loss_scale)
@@ -344,17 +363,20 @@ def forward_step_calc_loss(
         else:
             MTPLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
 
-    # Set the loss scale for DSA (Dynamic Sparse Attention) indexer loss.
-    if getattr(config, 'experimental_attention_variant', None) == 'dsa':
-        loss_scale = (
-            config.grad_scale_func(torch.ones(1, device=output_tensor.device))
-            if config.grad_scale_func is not None
-            else torch.ones(1, device=output_tensor.device)
-        )
+    # Set the loss scale for any experimental attention-variant auxiliary loss.
+    experimental_attention_variant_loss_scale_func = (
+        _get_experimental_attention_variant_loss_scale_func(config)
+    )
+    if experimental_attention_variant_loss_scale_func is not None:
+        device = get_tensor_device(output_tensor)
+        loss_scale = _compute_loss_scale(config, device)
         if config.calculate_per_token_loss:
-            DSAIndexerLossAutoScaler.set_loss_scale(loss_scale)
+            experimental_attention_variant_loss_scale_func(loss_scale)
         else:
-            DSAIndexerLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
+            cp_size_for_scaling = cp_group_size if cp_group_size is not None else 1
+            experimental_attention_variant_loss_scale_func(
+                loss_scale * cp_size_for_scaling / num_microbatches
+            )
 
     return output_tensor, num_tokens