test(data): align packed EOS padding masks

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

tests/unit_tests/data/datasets/test_chat_template.py

@@ -1027,7 +1027,7 @@ def test_eos_index_logic_uses_shape_check(self):
             {
                 "input_ids": np.array([7, 0, 0, 0, 0], dtype=np.int64),
                 "seq_boundaries": [0, 5],
-                "loss_mask": np.ones(5, dtype=np.int64),
+                "loss_mask": np.array([1, 0, 0, 0, 0], dtype=np.int64),
             }
         ]
 
@@ -1056,7 +1056,7 @@ def test_static_metadata_keeps_logical_and_physical_boundaries_aligned(self):
             {
                 "input_ids": np.array([7, 0, 0, 0, 0], dtype=np.int64),
                 "seq_boundaries": [0, 5],
-                "loss_mask": np.ones(5, dtype=np.int64),
+                "loss_mask": np.array([1, 0, 0, 0, 0], dtype=np.int64),
             }
         ]
 
@@ -1068,6 +1068,23 @@ def test_static_metadata_keeps_logical_and_physical_boundaries_aligned(self):
         assert processed["cu_seqlens_q_padded"][0].tolist() == [0, 4, 8, 8, 8]
         assert processed["cu_seqlens_kv_padded"][0].tolist() == [0, 4, 8, 8, 8]
 
+    def test_supervised_terminal_eos_precedes_zero_loss_padding(self):
+        """A supervised EOS remains logical while later zero-loss EOS tokens are padding."""
+        dataset = _create_minimal_packed_dataset()
+        batch = [
+            {
+                "input_ids": np.array([7, 0, 0, 0, 0], dtype=np.int64),
+                "seq_boundaries": [0, 5],
+                "loss_mask": np.array([1, 1, 0, 0, 0], dtype=np.int64),
+            }
+        ]
+
+        processed = dataset.collate_fn(batch)
+
+        assert processed["cu_seqlens_q"][0].tolist() == [0, 2]
+        assert processed["cu_seqlens_q_padded"][0].tolist() == [0, 4]
+        assert processed["padding_mask"].tolist() == [[False, False, True, True]]
+
     def test_without_alignment_padding_omits_physical_variants(self):
         """Identical logical and physical layouts use only the faster logical TE fields."""
         dataset = _create_minimal_packed_dataset()
@@ -1261,14 +1278,14 @@ def test_logical_cu_seqlens_uses_correct_eos(self):
                     dtype=np.int64,
                 ),
                 "seq_boundaries": [0, 5, 10],
-                "loss_mask": np.ones(10, dtype=np.int64),
+                "loss_mask": np.array([1, 0, 0, 0, 0, 1, 0, 0, 0, 0], dtype=np.int64),
             }
         ]
 
         processed = dataset.collate_fn(batch)
         cu_logical = processed["cu_seqlens_q"][0].tolist()
 
-        # Each non-EOS token contributes exactly once despite EOS padding.
+        # Each non-EOS token contributes exactly once despite zero-loss EOS padding.
         assert cu_logical == [0, 1, 2]