[split 2/4] Scale DSA indexer loss in pipeline schedules

megatron/core/pipeline_parallel/hybrid_cp_schedule.py

@@ -545,9 +545,17 @@ 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
+
+        # Keep this int32 to match the hybrid-CP batch metadata dtype
+        # (`local_cp_size`) used by get_batch_on_this_cp_rank.
+        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 +587,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 +598,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 +608,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 +621,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 +631,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 +641,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 +655,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,24 @@ 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)
+            # TODO: This path assumes static CP across outstanding pipeline microbatches.
+            # Hybrid/dynamic CP currently requires per-token loss and no PP; if that
+            # changes, carry the scale per autograd context instead of via a
+            # process-wide scaler hook.
+            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