From b4eb3441fece310794e2f445097085563e0a072f Mon Sep 17 00:00:00 2001 From: zhengshengning Date: Mon, 31 Aug 2026 10:53:25 +0800 Subject: [PATCH] add_loss_print --- .../common/language_module/language_module.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 62199c265a1..fb7b8ef405e 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -222,6 +222,22 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor: # [s b] => [b, s] loss = loss.transpose(0, 1).contiguous() + + if _use_accuracy_compatible(): + # 精度对齐锚点 1（对应 PF language_loss.py forward_impl 里的 per_token_loss）： + # CE 直出、mask/归一化前的 per-token loss，两侧语义唯一。 + # 注意：CP > 1 时这里仍是本 rank 的 sequence shard，需与 PF 侧 + # ContextParallelGatherOp 之后的全量 shape 区分。 + import hashlib as _hashlib + + _l = loss.detach().float().contiguous() + print( + f"\nper_token_loss: rank={torch.distributed.get_rank()} " + f"shape={list(_l.shape)} " + f"md5={_hashlib.md5(_l.cpu().numpy().tobytes()).hexdigest()}", + flush=True, + ) + return loss def setup_embeddings_and_output_layer(self) -> None: