fix(training): neutralize sub-threshold loss masks instead of per-ran…

[chungen04]

Close #126

Summary

Multi-rank FSDP training can hang on an NCCL collective timeout because DataFetcher drops samples (empty loss mask, or < min_loss_tokens supervised tokens) per-rank and uncoordinated. When such a sample lands on some (not all) DP ranks in a dispatched step, the dropping rank runs one fewer micro-batch, the ranks fall out of collective lockstep and the next collective blocks until the watchdog aborts.

The bug was triggered with the example, configs/sglang_qwen3_8b_dflash.yaml, per described in #126, where we have min_loss_tokens: 32. The field was not present on all EAGLE-3 training examples.

In this PR, the data fetcher zeroes the loss mask in place and keeps the sample, so every dispatched sample still yields exactly one micro-batch. Collectives stay balanced, and a zero mask contributes 0 loss.

Code Trace

DataFetcher (torchspec/training/data_fetcher.py) filtered in the per-rank, post-partition data path:

def _should_skip_for_loss_mask(self, data, mooncake_key, skip_count):
    mask = self._compute_loss_mask(data)
    if mask is None: ...; return True, ...                 # all-zero
    if self._min_loss_tokens > 0 and mask.sum() < self._min_loss_tokens:
        ...; return True, ...                              # too few supervised tokens
    return False, ...

def __iter__(self):
    ...
    should_skip, skip_count = self._should_skip_for_loss_mask(...)
    if should_skip:
        continue                                           # ← per-rank DROP

The controller deals an equal per_dp_rank_batch_size to every rank each dispatch. Each rank sees different samples, so the drops are uncorrelated across ranks. Under FSDP every rank must issue the same sequence of all-gathers, so one rank dropping a micro-batch desyncs the group and the next collective hangs.

The fix

Modified torchspec/training/data_fetcher.py.

The per-rank data path is made drop-free. Replace the skip with _resolve_and_neutralize_loss_mask, which keeps the sample and zeroes its loss mask:

def _resolve_and_neutralize_loss_mask(self, data, mooncake_key, neutralized_count):
    mask = self._compute_loss_mask(data)                   # None == all-zero
    neutralize = mask is None or (
        self._min_loss_tokens > 0
        and isinstance(mask, torch.Tensor)
        and mask.sum() < self._min_loss_tokens
    )
    if not neutralize:
        return neutralized_count
    if isinstance(mask, torch.Tensor):
        data["loss_mask"] = torch.zeros_like(mask)
    else:                                                  # resolve returned None (no mask set)
        data["loss_mask"] = torch.zeros(self._fallback_mask_len(data), dtype=torch.long)
    ...                                                    # warn + count
    return neutralized_count

__iter__ now always yields. The USP-sharded path (_usp_get_sharded_item) gets the same treatment. It zeroes the local shard's mask instead of skipping. _fallback_mask_len derives the length from input_ids (or hidden_states/last_hidden_states/target) for the all-zero case where resolve_loss_mask returns None without setting a mask.

Why this is safe:

• Count-preserving. Every dispatched sample yields exactly one micro-batch on every rank. The all-gather sequence is identical across ranks by construction.
• Covers all mask paths. Acting in the fetcher (where the mask is always computed) covers the dynamic-mask path too. However, the controller has only packed_loss_mask + tensor shapes ,no input_ids, so it cannot compute the dynamic mask. A controller filter closes the DFlash (packed) hole but leaves the dynamic-mask path's per-rank drop live.

Behavior change: sub-threshold samples are processed (zeroed) rather than dropped. Loss is unaffected, only a small amount of extra compute (inference) on rare samples that satisfies the neutralize (dropping) samples were added.

Alternatives considered

• Substitute-on-skip (pull a replacement to keep the count): Collective-safe in steady state, but needs a non-starving replacement supply and a coordinated end-of-data drain.

Tests

• No regressions.
• Adding new tests. They fail on the old code and pass on the fix.
  • test_subthreshold_samples_neutralized_not_dropped (sub-threshold/empty samples are kept with zeroed masks; count preserved).
  • test_usp_sharded_keeps_local_zero_loss_shard_when_global_loss_exists

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 77e8b2b352

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve device when creating fallback zero masks

When dynamic_loss_mask is enabled and the fetcher is loading directly to CUDA (for example prefetch_depth=0 or eval), valid dynamic masks are returned on the input_ids device, but this all-zero fallback creates a CPU tensor. DataCollatorWithPadding later torch.cats the per-sample loss masks, so a batch that mixes one neutralized all-zero dynamic sample with one valid dynamic sample fails with a CPU/CUDA device mismatch before training can proceed. Create this fallback on the same device as input_ids or another sequence tensor.

Useful? React with 👍 / 👎.