88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
# 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."
|
|
)
|