[Test] Check GLM-5.3 end-to-end logits alongside loss

tests/model/test_glm53_text_moe.py

@@ -38,6 +38,7 @@
 
 from transformers import AutoTokenizer, Glm5NextForConditionalGeneration
 from xtuner._testing import DeterministicDDPTestCase
+from xtuner._testing.logits import check_logits
 from xtuner.v1.config import FSDPConfig
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.loss.ce_loss import CELossConfig
@@ -357,12 +358,16 @@ def test_fsdp_accuracy(self, dispatcher, ep_size):
             "就像老树拥抱归巢的鸟儿，内存管理应该给予每个对象足够的安全感",
         ]
         tokenizer = AutoTokenizer.from_pretrained(GLM_5_3_FLASH_PATH)
+        hf_model.eval()
         expected_losses = []
+        expected_logits = []
         for text in text_list:
             input_ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
             with torch.no_grad():
                 output = hf_model(input_ids=input_ids, labels=input_ids.clone())
             expected_losses.append(output.loss)
+            # Last eight next-token positions, including the full vocabulary.
+            expected_logits.append(output.logits[0, :-1][-8:].detach().cpu())
 
         del hf_model
         torch.cuda.empty_cache()
@@ -381,8 +386,9 @@ def test_fsdp_accuracy(self, dispatcher, ep_size):
         model.fully_shard(fsdp_config=fsdp_config)
         model.from_hf(GLM_5_3_FLASH_PATH, strict=False)
 
+        model.eval()
         losses = []
-        for text in text_list:
+        for sample_index, text in enumerate(text_list):
             input_ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
             shift_input_ids = input_ids[:, :-1]
             shifted_labels = input_ids[:, 1:]
@@ -396,6 +402,11 @@ def test_fsdp_accuracy(self, dispatcher, ep_size):
             with torch.no_grad():
                 output = model(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx})
             losses.append(output["loss"])
+            # Preserve the public loss path above; additionally exercise inference logits.
+            with torch.no_grad():
+                logits = model(seq_ctx=seq_ctx, loss_ctx=None).logits
+            metrics = check_logits(logits[0, -8:], expected_logits[sample_index])
+            print(f"GLM53 text sample={sample_index} logits={metrics}", flush=True)
 
         self._check_loss_curve(
             losses=torch.tensor(losses), losses_ref=torch.tensor(expected_losses), sim_tol=3e-2, rtol=3e-2

tests/model/test_logits_metrics.py

@@ -0,0 +1,37 @@
+"""The logits check must catch errors that cross entropy cannot detect."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from xtuner._testing.logits import check_logits
+
+
+def test_identical_logits():
+    x = torch.arange(40).reshape(4, 10).float()
+    assert check_logits(x, x)["relative_l2"] == 0
+
+
+def test_common_offset_preserves_ce_but_fails_logits():
+    x = torch.arange(40).reshape(4, 10).float() / 10
+    y = torch.tensor([0, 1, 2, 3])
+    torch.testing.assert_close(F.cross_entropy(x, y), F.cross_entropy(x + 10, y))
+    with pytest.raises(AssertionError):
+        check_logits(x + 10, x)
+
+
+def test_position_permutation_is_rejected():
+    x = torch.eye(8)
+    with pytest.raises(AssertionError):
+        check_logits(x.roll(1, 0), x)
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_nonfinite_logits_are_rejected(value):
+    with pytest.raises(AssertionError):
+        check_logits(torch.full((2, 4), value), torch.ones(2, 4))
+
+
+def test_wrong_shape_is_rejected():
+    with pytest.raises(AssertionError):
+        check_logits(torch.ones(2, 4), torch.ones(3, 4))

xtuner/_testing/logits.py

@@ -0,0 +1,27 @@
+"""Diagnostics for sampled, full-vocabulary end-to-end logits."""
+
+import torch
+import torch.nn.functional as F
+
+
+def check_logits(actual, expected, *, max_relative_l2=0.05, min_cosine=0.998):
+    """Check each sample separately; use FP64 only for offline reductions.
+
+    Inputs are [sampled_positions, vocab], not hidden states or target-token logits.
+    Callers must select the same globally indexed, non-padding positions on both sides.
+    The 5% L2 bound is independent of the existing mean loss-curve tolerance.
+    """
+    assert actual.shape == expected.shape and actual.ndim == 2 and actual.numel() > 0
+    assert torch.isfinite(actual).all() and torch.isfinite(expected).all(), "non-finite logits"
+    actual = actual.detach().to(device="cpu", dtype=torch.float64).flatten()
+    expected = expected.detach().to(device="cpu", dtype=torch.float64).flatten()
+    norm = expected.norm()
+    assert norm > 0, "zero reference norm"
+    delta = actual - expected
+    metrics = {
+        "relative_l2": (delta.norm() / norm).item(),
+        "cosine": F.cosine_similarity(actual, expected, dim=0).item(),
+        "max_abs": delta.abs().max().item(),
+    }
+    assert metrics["relative_l2"] < max_relative_l2 and metrics["cosine"] > min_cosine, metrics
+    return metrics