fix: sanitize non-finite target hidden states instead of poisoning the run - #18
Closed
Boreas618 wants to merge 2 commits into
Closed
fix: sanitize non-finite target hidden states instead of poisoning the run#18Boreas618 wants to merge 2 commits into
Boreas618 wants to merge 2 commits into
Conversation
`_sample_anchor_positions` raised `ValueError` when no sample in the micro-batch had a usable anchor. Under distributed training the loss denominator is all-reduced across the world further down the same forward (`dist.all_reduce(global_loss_den, ...)`), so the raising rank exits *before* that collective and every peer blocks on it until the NCCL timeout. One degenerate micro-batch on one rank hangs the entire job instead of failing it, with no indication of which rank was at fault. This is reachable in practice whenever the target emits non-finite hidden states for some rows: the loss mask is zeroed for those tokens and the anchor sampler then finds no supervised adjacent pair. Observed on a 32-rank run where a single target replica poisoned its batches and the remaining 24 ranks hung. Degrade to a fully-masked micro-batch instead. `keep_mask` is all False, so the rank contributes 0 to every numerator and 0 to both the local and global loss denominator, while still reaching 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 -- but collectively, on every rank, which is safe. Applies to `OnlineDFlashModel` and `OnlineDSparkModel`. `OnlineDominoModel` already degrades via `max(1, ...)` and is unchanged. A warning is logged so the condition stays visible rather than silent. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e run Target capture can emit non-finite hidden states for individual rows (serving-side numerical issues, kernel bugs, allocator faults). Upstream has no isfinite handling anywhere in the training path, so one bad capture is fatal twice over: the loss goes non-finite, and under data parallelism the gradient all-reduce propagates that single rank's NaN to every rank -- one bad row over millions poisons the whole run. Add `_sanitize_nonfinite_inputs` to the DFlash family base, called at the top of both OnlineDFlashModel.forward and OnlineDSparkModel.forward (covering Domino via inheritance). Per affected token it zeroes the loss mask (excluded from supervision) AND the feature itself, so inf/NaN cannot leak into healthy tokens through the draft's attention over the context. A logger.warning with a per-tensor finiteness report (count, inf/nan flags, finite absmax) keeps the condition visible and diagnosable rather than silent. On-by-default via `sanitize_nonfinite: bool = True` constructor flag; no config-schema changes in this PR. Stacked on the anchor-sampling degrade fix: a fully poisoned micro-batch sanitizes to an all-masked batch, which sampling now tolerates instead of raising past the loss-denominator collective. The diagnostics format is proven: it is what isolated a target-side replica bug on a 32-rank Kimi-K3 run (hidden_states 9719808/22650880 nonfinite, finite_absmax=0.000e+00 on the final states -- unwritten memory, not overflow). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
1 participant
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
Upstream has zero
isfinitehandling anywhere in the training path. Target capture can emit non-finite hidden states for individual rows (serving-side numerical issues, kernel bugs, allocator faults). One bad capture is then fatal twice over:Change
Adds
_sanitize_nonfinite_inputsto the DFlash family base class, called at the top ofOnlineDFlashModel.forwardandOnlineDSparkModel.forward(Domino inherits the former). Per affected token:nan_to_numon the features — inf/NaN cannot leak into healthy tokens through the draft's attention over the context.A
logger.warningemits a per-tensor finiteness report — count, inf/nan flags, andfinite_absmax— so the condition is visible and diagnosable, never silent.On by default (
sanitize_nonfinite: bool = Trueconstructor flag). No config-schema changes in this PR; wiring a YAML knob can follow separately if wanted.Why the diagnostics format matters
This exact report is what isolated a target-side replica bug on a 32-rank Kimi-K3 DSpark run:
finite_absmax=0.000e+00distinguishes unwritten memory from overflow in one line — the difference between chasing data bugs and chasing serving bugs.Stacked on #2
A fully poisoned micro-batch sanitizes to an all-masked batch. Without #2 that raises inside anchor sampling before the loss-denominator all-reduce and hangs every peer — i.e. this PR alone could turn a NaN loss into a hang. With #2 the batch degrades gracefully and the run continues. Merge #2 first; this branch includes its commit.
Tests
tests/test_utils/test_nonfinite_sanitization.py— 5 cases: clean pass-through, NaN/inf token masking + element zeroing, bad final-hidden with clean context preserved, opt-out flag, and the fully-poisoned batch composing with #2's degrade path.