diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 3f8f7ce0d..e0b2297dd 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/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) diff --git a/tests/test_utils/test_dflash_losses.py b/tests/test_utils/test_dflash_losses.py index 6001ce336..858afed56 100644 --- a/tests/test_utils/test_dflash_losses.py +++ b/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__":