From 29147b2d5837cf034670cbc0259b8c81674d3a19 Mon Sep 17 00:00:00 2001
From: taking-lying-flat <1615405@qq.com>
Date: Sun, 23 Aug 2026 16:57:17 +0800
Subject: [PATCH] Validate vocab-parallel CE labels

---
 torchtitan/components/loss.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/torchtitan/components/loss.py b/torchtitan/components/loss.py
index bcc0fbe907..b3a7ce704c 100644
--- a/torchtitan/components/loss.py
+++ b/torchtitan/components/loss.py
@@ -142,6 +142,14 @@ def forward(
                 "_LossParallelCrossEntropy does not support empty vocab shards."
             )
 
+        torch._assert_async(
+            torch.all(
+                (labels == IGNORE_INDEX)
+                | ((labels >= 0) & (labels < global_vocab_size))
+            ),
+            f"labels must be {IGNORE_INDEX} or in [0, {global_vocab_size})",
+        )
+
         # All-reduce max for numerically stable distributed log-softmax.
         local_max = torch.amax(logits, dim=-1, keepdim=True)
         local_max = funcol.all_reduce(