fix(sft): drop rows with phantom media tokens; fail fast on fatal rank errors

src/prime_rl/trainer/sft/data.py

@@ -312,6 +312,29 @@ def should_mask(message: dict) -> bool:
                 "Set [model.vlm] to train on multimodal samples."
             )
 
+        # Literal media-marker text (e.g. "<image>" inside a code comment or tool-call
+        # arguments) tokenizes to the model's image-placeholder id with no pixel data
+        # behind it. One phantom token in a packed batch fails the scatter-time
+        # token/feature check on whichever rank draws it — and that rank's teardown then
+        # wedges every other rank in a collective until the NCCL timeout, which blames a
+        # victim collective instead of the row. mm_token_type_ids marks only
+        # renderer-emitted placeholder runs, so any placeholder id at a type-0 position —
+        # or anywhere in a row that produced no multimodal data — is phantom text.
+        mm_map = getattr(self.renderer, "mm_token_type_id_map", None)
+        if mm_map:
+            placeholder_ids = set(mm_map)
+            if mm_token_type_ids is None:
+                phantom = sum(1 for t in input_ids if t in placeholder_ids)
+            else:
+                phantom = sum(1 for t, tt in zip(input_ids, mm_token_type_ids) if tt == 0 and t in placeholder_ids)
+            if phantom:
+                self.logger.warning(
+                    f"Dropping example {example.get('__index', '')} ({example.get('id', '?')}): "
+                    f"{phantom} media placeholder token(s) outside any renderer-emitted "
+                    f"placeholder run (literal marker text in content)"
+                )
+                return None
+
         # Causal shift: model predicts next token from current.
         target_ids = input_ids[1:]
         loss_mask = loss_mask[1:]

src/prime_rl/utils/utils.py

@@ -2,7 +2,6 @@
 import functools
 import importlib
 import os
-import sys
 from collections import defaultdict
 from contextlib import contextmanager
 from pathlib import Path
@@ -50,18 +49,21 @@ def clean_exit(func: Callable) -> Callable:
         async def async_wrapper(*args, **kwargs):
             try:
                 ret = await func(*args, **kwargs)
-                wandb.finish()
-                return ret
             except Exception:
                 get_logger().opt(exception=True).error(f"Fatal error in {func.__name__}")
                 wandb.finish(exit_code=1)
-                # sys.exit raises SystemExit so the finally block still runs.
-                # raise alone doesn't terminate the process in an async context —
-                # the event loop swallows it and the process hangs indefinitely.
-                sys.exit(1)
-            finally:
-                if dist.is_initialized():
-                    dist.destroy_process_group()
+                # Do NOT destroy the process group on the fatal path: peer ranks are
+                # typically still blocked in a collective this rank will never join, so
+                # a graceful shutdown hangs in ProcessGroup.shutdown() and the job only
+                # dies at the collective timeout — with the watchdog blaming a victim
+                # collective on a healthy rank. Hard-exit instead so the launcher's
+                # failure propagation (torchrun / srun --kill-on-bad-exit) tears the
+                # world down immediately.
+                os._exit(1)
+            wandb.finish()
+            if dist.is_initialized():
+                dist.destroy_process_group()
+            return ret
 
         return async_wrapper
     else:
@@ -70,16 +72,21 @@ async def async_wrapper(*args, **kwargs):
         def sync_wrapper(*args, **kwargs):
             try:
                 ret = func(*args, **kwargs)
-                wandb.finish()
-                return ret
             except Exception:
                 get_logger().opt(exception=True).error(f"Fatal error in {func.__name__}")
                 wandb.finish(exit_code=1)
-                # sys.exit raises SystemExit so the finally block still runs.
-                sys.exit(1)
-            finally:
-                if dist.is_initialized():
-                    dist.destroy_process_group()
+                # Do NOT destroy the process group on the fatal path: peer ranks are
+                # typically still blocked in a collective this rank will never join, so
+                # a graceful shutdown hangs in ProcessGroup.shutdown() and the job only
+                # dies at the collective timeout — with the watchdog blaming a victim
+                # collective on a healthy rank. Hard-exit instead so the launcher's
+                # failure propagation (torchrun / srun --kill-on-bad-exit) tears the
+                # world down immediately.
+                os._exit(1)
+            wandb.finish()
+            if dist.is_initialized():
+                dist.destroy_process_group()
+            return ret
 
         return sync_wrapper