fix: assert no duplicate starting bos (#835)

Signed-off-by: Zhiyu Li <[email protected]>

Author: ZhiyuLi-Nvidia

examples/run_grpo_math.py

@@ -89,15 +89,11 @@ def hf_data_processor(
         add_generation_prompt=True,
         add_special_tokens=False,
     )
-    # add bos token if not in chat template
-    bos_token_in_chat_template = tokenizer.chat_template.startswith(
-        "{{- bos_token }}"
-    ) or tokenizer.chat_template.startswith("{{ bos_token }}")
 
     user_message["token_ids"] = tokenizer(
         message,
         return_tensors="pt",
-        add_special_tokens=not bos_token_in_chat_template,
+        add_special_tokens=False,
     )["input_ids"][0]
     user_message["content"] = message
     message_log.append(user_message)

nemo_rl/data/datasets.py

@@ -125,7 +125,7 @@ def __getitem__(self, idx: int) -> DatumSpec:
                 assert isinstance(token_ids, torch.Tensor), (
                     f"token_ids must be a torch.Tensor, got {type(token_ids)}"
                 )
-                assert_start_with_single_bos(token_ids, self.tokenizer)
+                assert_no_double_bos(token_ids, self.tokenizer)
             self._bos_checked = True
 
         return datum_spec
@@ -301,25 +301,20 @@ def dpo_collate_fn(
     return train_data
 
 
-def assert_start_with_single_bos(
-    token_ids: torch.Tensor, tokenizer: TokenizerType
-) -> None:
-    """Assert that the first token is a BOS token and the second token is not a BOS token.
+def assert_no_double_bos(token_ids: torch.Tensor, tokenizer: TokenizerType) -> None:
+    """Assert that there are no double starting BOS tokens in the message.
 
     Args:
         token_ids: List of token IDs
         tokenizer: Tokenizer
     """
     if tokenizer.bos_token_id is not None:
         token_ids_list = token_ids.tolist()
-        if len(token_ids_list) > 0:
-            assert token_ids_list[0] == tokenizer.bos_token_id, (
-                f"Expected BOS token at the start of the message, but got {token_ids_list[0]}"
-            )
         if len(token_ids_list) > 1:
-            assert token_ids_list[1] != tokenizer.bos_token_id, (
-                f"Expected non-BOS token at the second position of the message, but got {token_ids_list[1]}"
-            )
+            assert not (
+                token_ids_list[0] == tokenizer.bos_token_id
+                and token_ids_list[1] == tokenizer.bos_token_id
+            ), "Found double BOS token in the first two positions of the message."
     else:
         print(
             f"skip assert_start_single_bos since Tokenizer {tokenizer.name_or_path} has no BOS token"

tests/unit/data/test_data_processor.py

@@ -73,6 +73,8 @@ def test_math_data_processor():
         "Qwen/Qwen2.5-1.5B-Instruct",  # no bos token
         "google/gemma-3-1b-it",
         "Qwen/Qwen3-0.6B",  # no bos token
+        "deepseek-ai/DeepSeek-V3",
+        "moonshotai/Moonlight-16B-A3B-Instruct",
     ],
 )
 @pytest.mark.parametrize(