fix: degrade empty anchor sampling to a masked batch to avoid multi-r…

specforge/algorithms/common/dflash_family_model.py

@@ -1,6 +1,7 @@
 # coding=utf-8
 """DFlash-family training models and shared masking helpers."""
 
+import logging
 from typing import Dict, Optional, Tuple
 
 import torch
@@ -11,6 +12,8 @@
 from specforge.modeling.draft.dflash import DFlashDraftModel
 from specforge.modeling.draft.flex_attention_backend import flex_attention_backend
 
+logger = logging.getLogger(__name__)
+
 try:
     from torch.nn.attention.flex_attention import BlockMask, create_block_mask
 
@@ -214,9 +217,18 @@ def _sample_anchor_positions(
             max_valid_anchors = int(valid_counts.max().item())
         width = min(self.num_anchors, max(0, int(max_valid_anchors)))
         if width == 0:
-            raise ValueError(
-                "DFlash-family training requires two consecutive supervised tokens"
+            # Raising here would skip this rank past later loss collectives and
+            # hang peers that did find valid anchors. Keep one fully masked
+            # block so this rank contributes zero supervision and participates.
+            logger.warning(
+                "%s: no valid anchor positions in this micro-batch; "
+                "contributing zero supervision for this step.",
+                type(self).__name__,
             )
+            bsz = loss_mask.shape[0]
+            anchors = torch.zeros((bsz, 1), dtype=torch.long, device=device)
+            keep_mask = torch.zeros((bsz, 1), dtype=torch.bool, device=device)
+            return anchors, keep_mask
 
         random_values = torch.rand(valid.shape, device=device)
         random_values.masked_fill_(~valid, 2.0)

tests/test_utils/test_dflash_losses.py

@@ -837,6 +837,40 @@ def test_dspark_sampler_keeps_sparse_high_index_anchor(self):
         self.assertEqual(anchors[0, 0].item(), 4)
         self.assertEqual(keep[0].tolist(), [True, False])
 
+    def test_sampler_degrades_to_masked_batch_without_supervision(self):
+        # Raising inside anchor sampling would exit before the loss-denominator
+        # all-reduce and hang every peer until the NCCL timeout, so a batch
+        # with no usable anchors must instead yield a fully-masked batch that
+        # still reaches the collective.
+        model = _make_model(self.logits, self.anchors, self.keep_mask)
+        loss_mask = torch.zeros(2, 16)
+        anchors, keep = OnlineDFlashModel._sample_anchor_positions(
+            model,
+            seq_len=16,
+            loss_mask=loss_mask,
+            device=loss_mask.device,
+        )
+        self.assertEqual(anchors.shape, keep.shape)
+        self.assertEqual(anchors.dtype, torch.long)
+        self.assertEqual(keep.dtype, torch.bool)
+        self.assertFalse(bool(keep.any()), "keep_mask must be entirely False")
+
+    def test_dspark_sampler_degrades_without_adjacent_supervised_pair(self):
+        # Supervised tokens exist but are isolated, so no anchor has its first
+        # target supervised -- the shape left behind when non-finite
+        # sanitization zeroes parts of the loss mask.
+        model = _make_dspark_model(self.logits, self.anchors, self.keep_mask)
+        loss_mask = torch.zeros(2, 16)
+        loss_mask[:, ::2] = 1.0
+        anchors, keep = OnlineDSparkModel._sample_anchor_positions(
+            model,
+            seq_len=16,
+            loss_mask=loss_mask,
+            device=loss_mask.device,
+        )
+        self.assertEqual(anchors.shape, keep.shape)
+        self.assertFalse(bool(keep.any()), "keep_mask must be entirely False")
+
     def test_shared_sampler_uses_adjacent_targets_and_partial_tails(self):
         sampler = OnlineDFlashModel._sample_anchor_positions
         model = _anchor_sampler_subject()
@@ -879,16 +913,19 @@ def test_shared_sampler_is_batch_padding_invariant(self):
         self.assertEqual(batch_anchors[0][batch_keep[0]].tolist(), [2])
         self.assertFalse(batch_keep[2].any())
 
-    def test_shared_sampler_rejects_a_batch_without_adjacent_targets(self):
+    def test_shared_sampler_degrades_without_adjacent_targets(self):
         model = _anchor_sampler_subject()
         loss_mask = torch.tensor([[1.0, 0.0, 1.0]])
-        with self.assertRaisesRegex(ValueError, "two consecutive"):
-            OnlineDFlashModel._sample_anchor_positions(
-                model,
-                seq_len=loss_mask.shape[1],
-                loss_mask=loss_mask,
-                device=loss_mask.device,
-            )
+        anchors, keep = OnlineDFlashModel._sample_anchor_positions(
+            model,
+            seq_len=loss_mask.shape[1],
+            loss_mask=loss_mask,
+            device=loss_mask.device,
+        )
+
+        self.assertEqual(anchors.shape, (1, 1))
+        self.assertEqual(keep.shape, (1, 1))
+        self.assertFalse(bool(keep.any()), "keep_mask must be entirely False")
 
 
 if __name__ == "__main__":