diff --git a/tests/model/test_glm53_text_moe.py b/tests/model/test_glm53_text_moe.py index fd8c233b50..ff8e3bad52 100644 --- a/tests/model/test_glm53_text_moe.py +++ b/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 diff --git a/tests/model/test_logits_metrics.py b/tests/model/test_logits_metrics.py new file mode 100644 index 0000000000..d0e00652d4 --- /dev/null +++ b/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)) diff --git a/xtuner/_testing/logits.py b/xtuner/_testing/logits.py new file mode 100644 index 0000000000..9ca217aa7f --- /dev/null +++ b/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