fix(xtoken): keep the offset layout for all-empty text batches

When every text in a batch tokenizes to no tokens, the tokenizer
returns offset_mapping as a (B, 0) tensor without the start/end axis.
Text batches now pad to the longest row, so such a batch is padded to
one sequence-length multiple, and concatenating (B, pad, 2) offsets onto
the (B, 0) tensor raised "Tensors must have same number of dimensions:
got 2 and 3". Reshape the offsets to (B, T, 2) from the token ids'
shape, so the batch collates to all-padding rows as it did when text
batches padded to the full context length. Non-empty batches are
unchanged.

Add a regression test for an all-empty text batch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>

Author: avenkateshha

nemo_rl/data/cross_tokenizer/cross_tokenizer_collate.py

@@ -670,7 +670,10 @@ def _tokenize_batch(
         )
         input_ids: torch.Tensor = encoded["input_ids"]
         attention_mask: torch.Tensor = encoded["attention_mask"]
-        offset_mapping: torch.Tensor = encoded["offset_mapping"]
+        # An all-empty batch returns (B, 0) offsets; keep the (B, T, 2) layout.
+        offset_mapping: torch.Tensor = encoded["offset_mapping"].reshape(
+            *input_ids.shape, 2
+        )
 
         b, t = input_ids.shape
         pad = CrossTokenizerCollator._padded_width(t, make_seq_div_by) - t

tests/unit/data/test_cross_tokenizer_collate.py

@@ -797,6 +797,19 @@ def test_batch_matches_full_context_padding_at_real_positions(self):
             for key in ("alignment_0_pair_valid", "alignment_0_pair_is_correct"):
                 assert torch.equal(new[key][row, pairs], old[key][row, pairs]), key
 
+    def test_all_empty_batch_pads_to_the_divisor(self):
+        collator = _text_collator(ctx_length_student=8, ctx_length_teachers=[16, 8])
+        out = collator([_datum("", idx=0), _datum("", idx=1)])
+        # No row has tokens, so each side pads to one sequence-length multiple.
+        assert out["input_ids"].shape == (2, 4)
+        assert out["teacher_0_input_ids"].shape == (2, 3)
+        assert out["input_lengths"].tolist() == [0, 0]
+        assert out["teacher_0_input_lengths"].tolist() == [0, 0]
+        assert out["token_mask"].eq(0).all()
+        assert out["teacher_0_token_mask"].eq(0).all()
+        assert out["alignment_0_student_chunk_id"].eq(-1).all()
+        assert out["alignment_0_teacher_chunk_id"].eq(-1).all()
+
 
 class TestCollatorPadTokenFallback:
     def test_missing_pad_token_set_from_eos(self):