fix: don't hang every rank when a micro-batch has no valid anchors - #2
Open
Boreas618 wants to merge 1 commit into
Open
fix: don't hang every rank when a micro-batch has no valid anchors#2Boreas618 wants to merge 1 commit into
Boreas618 wants to merge 1 commit into
Conversation
…ank hang Raising when a micro-batch has no valid anchors skips the loss-denominator all-reduce and hangs peer ranks until NCCL timeout. Factor candidate selection into `_anchor_candidates` so DSpark can override eligibility without duplicating the sampler.
Boreas618
force-pushed
the
fix/anchor-sampling-multi-rank-hang
branch
from
August 1, 2026 16:51
e765a03 to
c161b72
Compare
This was referenced Aug 1, 2026
Closed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
_sample_anchor_positionsraisesValueErrorwhen no sample in a micro-batch has a usable anchor:Under distributed training the loss denominator is all-reduced later in the same forward:
The raising rank exits before that collective, so every peer blocks on it until the NCCL timeout. One degenerate micro-batch on one rank hangs the whole job rather than failing it, and the traceback names the data rather than the rank.
When this fires
Whenever the target emits non-finite hidden states for some rows. Those tokens get their loss mask zeroed, and the sampler then finds no supervised adjacent pair. Hit on a 32-rank Kimi-K3 DSpark run where one target replica poisoned its batches — the other 24 ranks hung with no error, and the visible symptom was a stall, not a crash.
Change
Degrade to a fully-masked micro-batch instead of raising.
keep_maskis allFalse, so the rank contributes 0 to every numerator and 0 to both the local and global loss denominator, and still reaches the collective. Training continues on the other ranks' samples.If no rank has supervision, the existing post-all-reduce check on
global_loss_denstill raises — collectively, on every rank, which is safe.A
logger.warningkeeps the condition visible instead of silently dropping supervision.Scope
OnlineDFlashModel._sample_anchor_positions— fixedOnlineDSparkModel._sample_anchor_positions— fixedOnlineDominoModel._sample_anchor_positions— already degrades viamax(1, ...), unchangedTests
tests/test_utils/test_anchor_sampling_degrade.py— 3 cases: DFlash with nothing supervised, DSpark with only isolated supervised tokens (the post-sanitization shape), and a normal fully-supervised batch that must still produce anchors.Pre-existing failures in
test_dflash_losses/test_dflash_mask(18/17) are identical on unmodifiedmain— this environment lacks CUDA, triton and flash_attn. No regressions.