Skip to content

fix: don't hang every rank when a micro-batch has no valid anchors - #2

Open
Boreas618 wants to merge 1 commit into
mainfrom
fix/anchor-sampling-multi-rank-hang
Open

fix: don't hang every rank when a micro-batch has no valid anchors#2
Boreas618 wants to merge 1 commit into
mainfrom
fix/anchor-sampling-multi-rank-hang

Conversation

@Boreas618

Copy link
Copy Markdown
Owner

Problem

_sample_anchor_positions raises ValueError when no sample in a micro-batch has a usable anchor:

max_n = min(self.num_anchors, int(valid_counts.max().item()) - 1)
if max_n <= 0:
    raise ValueError("should preprocess the data.")

Under distributed training the loss denominator is all-reduced later in the same forward:

dist.all_reduce(global_loss_den, op=dist.ReduceOp.SUM)   # dflash_family_model.py

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_mask is all False, 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_den still raises — collectively, on every rank, which is safe.

A logger.warning keeps the condition visible instead of silently dropping supervision.

Scope

  • OnlineDFlashModel._sample_anchor_positions — fixed
  • OnlineDSparkModel._sample_anchor_positions — fixed
  • OnlineDominoModel._sample_anchor_positions — already degrades via max(1, ...), unchanged

Tests

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.

3 passed

Pre-existing failures in test_dflash_losses / test_dflash_mask (18/17) are identical on unmodified main — this environment lacks CUDA, triton and flash_attn. No regressions.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant