初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD
Source: Original Platform

slime/backends/megatron_utils/ci_utils.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""CI utilities for Megatron backend testing."""
+
+import logging
+from collections.abc import Sequence
+
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+logger = logging.getLogger(__name__)
+
+
+def check_mtp_only_grad(model: Sequence[DDP], step_id: int) -> None:
+    """Check that only MTP parameters have non-zero gradients.
+
+    This is used for CI testing to verify that when all outputs are truncated,
+    only the MTP layers receive gradients (since only mtp_loss contributes).
+
+    Args:
+        model: Sequence of DDP-wrapped model chunks.
+        step_id: Current step index for logging.
+
+    Raises:
+        AssertionError: If any non-MTP parameter has a non-zero gradient.
+    """
+    non_mtp_nonzero_grads = []
+    mtp_nonzero_grads = []
+
+    for model_chunk in model:
+        for name, param in model_chunk.named_parameters():
+            # Get the main_grad from the distributed optimizer if available
+            grad = getattr(param, "main_grad", None)
+            if grad is None:
+                grad = param.grad
+            if grad is None:
+                continue
+
+            grad_norm = grad.abs().max().item()
+            is_mtp = ".mtp." in name
+
+            if is_mtp:
+                if grad_norm > 0:
+                    mtp_nonzero_grads.append((name, grad_norm))
+            else:
+                if grad_norm > 0:
+                    non_mtp_nonzero_grads.append((name, grad_norm))
+
+    # Log the results
+    logger.info(
+        f"[CI MTP Grad Check] Step {step_id}: "
+        f"MTP params with non-zero grad: {len(mtp_nonzero_grads)}, "
+        f"non-MTP params with non-zero grad: {len(non_mtp_nonzero_grads)}"
+    )
+
+    if non_mtp_nonzero_grads:
+        # Log the first few non-MTP params with non-zero gradients for debugging
+        for name, grad_norm in non_mtp_nonzero_grads[:5]:
+            logger.error(f"[CI MTP Grad Check] Non-MTP param with non-zero grad: {name}, max_grad={grad_norm}")
+
+    assert len(non_mtp_nonzero_grads) == 0, (
+        f"Expected all non-MTP parameters to have zero gradients, "
+        f"but found {len(non_mtp_nonzero_grads)} with non-zero gradients. "
+        f"First few: {non_mtp_nonzero_grads[:5]}"
+    )
+
+    # Also verify that MTP params do have gradients (otherwise the test is not valid)
+    assert len(mtp_nonzero_grads) > 0, (
+        "Expected MTP parameters to have non-zero gradients, but all were zero. "
+        "This may indicate the MTP loss is not being computed."
+    )
+
+
+def check_mtp_loss(mtp_loss: float, max_mtp_loss: float = 1.0) -> None:
+    """Check that MTP loss is within expected bounds.
+
+    Args:
+        mtp_loss: The computed MTP loss value.
+        max_mtp_loss: Maximum allowed MTP loss (default: 1.0).
+
+    Raises:
+        AssertionError: If MTP loss exceeds the maximum allowed value.
+    """
+    assert mtp_loss < max_mtp_loss, (
+        f"MTP loss {mtp_loss} exceeds maximum allowed value {max_mtp_loss}. "
+        "This may indicate an issue with MTP training."
+    )