Repository navigation
Make a non-finite training step visible in log_history - #6862
Closed
behroozazarkhalili wants to merge 24 commits into
Closed
behroozazarkhalili wants to merge 24 commits into
behroozazarkhalili wants to merge 24 commits into
Conversation
`transformers` replaces a NaN or Inf loss with the average of the previously logged losses before logging it (`logging_nan_inf_filter`, on by default), but `training_step` has already run and `optimizer.step()` still runs on the non-finite value. The reported loss curve stays plausible while the weights are being corrupted, and nothing in `log_history` points at the step that went wrong. This costs more in GRPO than in plain SFT, because several GRPO paths reach a NaN from data rather than from divergence: an unscorable token logprob returned by vLLM (#6166), a degenerate reward group, an overflowing `exp` of a KL log-ratio (#3015). A user has no reason to suspect the learning rate in those cases, and a single poisoned group or token is not something a loss curve can point at. `compute_loss` now appends `frac_nonfinite_loss` and warns once. The indicator is gathered across ranks before it is recorded, since a NaN confined to one rank would otherwise read as 0.0 from the main process, which reproduces the same invisibility one layer up. The warning text carries no step index on purpose: `accelerate`'s `warning_once` is an `lru_cache` keyed on its arguments, so an interpolated step number misses the cache and warns on every step. `compute_loss` is the single funnel for both the Liger and the non-Liger path, so one site per trainer covers both. `GRPOWithReplayBufferTrainer` and the `gspo_token` variant subclass `GRPOTrainer` and override only `_compute_loss`, so they inherit the change. The regression tests poison the second of two steps and leave the first as the finite control, then assert on `log_history`: 0.0 then 1.0 for the new metric, with the reported `loss` still finite for the poisoned step. Both fail on main with `KeyError: 'frac_nonfinite_loss'`. Resolves #6702
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
3 of 8 tasks
Checkpoint so the work survives. Not the final message. `transformers` swaps a non-finite loss for a substitute value before logging it, while the backward pass still runs on the real one, so a diverging run keeps a plausible loss curve. Each trainer now gathers the non-finite indicator across ranks at the end of its loss method, records `frac_nonfinite_loss` where a metrics dict exists, and warns once. Every block is emitted from one generator so the copies cannot drift; all 20 are byte-identical modulo indent. grpo, rloo and async_grpo are re-synced onto the same text. Known open, from a full-diff review: the warning's account of what transformers substitutes is wrong (it reads the tr_loss accumulator, not logged history), the fp16/bf16 sentence has counterexamples on CPU and DeepSpeed, the filter claim does not hold on XLA, the metric ignores metric_key_prefix, and predict() leaks into the next eval window. None of those are fixed here.
`transformers` substitutes a non-finite training loss before logging it (`logging_nan_inf_filter`, on by default) while the backward pass still runs on the non-finite value. The substitution is not an average of previously logged losses, as the option's name and the surrounding comment suggest: it adds a fraction of the current accumulator to itself and never reads logged history. The reported curve therefore stays plausible while a run degrades, and nothing in `log_history` identifies the step that went wrong. Each trainer's loss path now appends `frac_nonfinite_loss` and warns once. The indicator is gathered across ranks first, since a NaN confined to one rank would otherwise read as zero from the main process and reproduce the same invisibility one layer up. The warning text carries no step index because `warning_once` caches on its arguments, so an interpolated step number would warn on every step instead of once. Trainers here are self-contained by design, so the block is duplicated rather than abstracted, and kept byte-identical across every copy apart from indentation and the four metric mechanisms the trainers already use. Two pre-existing defects surfaced while adding it, both fixed because the new metric is the first whose value a reader would notice is wrong. Eval metrics were filed under a hardcoded `eval_` prefix rather than the caller's `metric_key_prefix`, which is `test` for `predict()`. And `predict()` never calls `log()`, so the batches it scored survived into the next evaluation window and skewed its averages. Adding a collective to the loss path also turned any rank-local `raise` upstream of it into a hang rather than a clean failure, because the raising rank leaves while its peers block in the gather until the NCCL timeout. Twelve such raises across eight trainers now agree across ranks before they fire. The per-rank predicate stays inside the pure loss helpers, which are unit-tested on objects that have no accelerator. That agreement has to precede the operation it guards. In CPO and ORPO it sat below the cross-entropy that raises on the very shape mismatch it checks for, so a rank failing there never reached the gather and its peers waited on it anyway. Verifying the gather needs more than one process: with a single rank it is the identity, so removing it entirely leaves every single-process test green. `tests/distributed/nonfinite_loss.py` poisons one rank and asserts on all of them that the logged fraction covers the whole world, and `test_distributed.py` launches it. Four of the per-trainer tests inject an infinity rather than a NaN, which an implementation checking `isnan` in place of `isfinite` does not catch. Resolves #6702.
…ecks `iw_opd` and `server_distillation` each tested `missing_teacher[:, 1]` twice in a row and raised in both branches. The first raise always fires when the condition holds, so the second was dead from the moment it was written: it survived a rewrite that added the count-carrying message above it without removing the flag-only one below. An AST sweep over `trl/` now reports no other `if` that raises on a condition a preceding sibling already tested, and it finds exactly these two when run against the previous revision. While counting for that sweep, the rank-agreeing raise total in the previous commit message proved wrong. Measured against `origin/main`, the branch takes that count from 6 to 22, so it adds 16 rather than the twelve claimed there. The count of eight trainers is right.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Every non-finite loss test injected NaN only, so an `isnan` implementation of the guard passed the whole suite while the guard itself tests `~isfinite`. The 18 tests whose poison lands on the loss scalar now run for both NaN and Inf. That set now includes a2po: mutating its guard to `isnan` fails the Inf case and to `isinf` fails the NaN case, so both values reach its loss distinctly. The other 7 poison logits or log-probabilities, where an injected Inf arrives at the loss as NaN, so they keep injecting NaN and record why. Four comments described code that does something else. The stated reason for the NaN-only limit named a log-softmax the poison does not pass through. The GRPO and RLOO comments still called the injection multiplicative after it became additive, and claimed the poisoned step corrupts the weights, which adding a constant does not. The distributed test claimed a rank-local result differs from the expected fraction by at least 0.5, which is true at two ranks and false above them. `_open_eval_window` is defined in 16 trainers and no test touched it. A new test evaluates under three different `metric_key_prefix` values and asserts the metric arrives under each. Hardcoding the prefix back to "eval" fails two of the three.
…alse claims The two branches of the distillation-loss metric in SDFT and SDPO gathered different dtypes. One gathered `torch.zeros(())`, which is always float32; the other gathered `mean_distill_loss`, which follows the autocast dtype and so is bfloat16 or float16 under mixed precision. `accelerator.gather` requires every rank to contribute the same dtype and NCCL does not promote, so a run where one rank took each branch could not complete the collective. Measured under a bfloat16 autocast: the peer's value arrives as bfloat16 while the zero stays float32. Both sides are now pinned to float32. The mismatch was introduced by 13a0666 on this branch and is not present on main. The docstring in tests/distributed/nonfinite_loss.py claimed the assertion "would be false" with a single rank. It would not. With one rank the sole rank is also the poisoned one, so the assertion reads 1.0 against an expected 1.0 and passes whether or not the gather is there, which is exactly the blindness the file exists to remove. The docstring now says that. Five doc pages head the metric list with "While training and evaluating", but CPO and ORPO record `frac_nonfinite_loss` with `train_eval="train"`, and NashMD, OnlineDPO and XPO append it inside `training_step`. Each of those five bullets now states that the trainer records the metric on the training path only, so readers do not look for it among the evaluation metrics.
`PRMTrainer` defines no loss of its own, so it inherits `transformers.Trainer.compute_loss` and the standard training loop, which is exactly where `logging_nan_inf_filter` applies. It was left out of this PR on the reasoning that it owns no loss site, but owning one and being subject to the filter are different questions. Measured before the fix: with a NaN injected on step 2, `log_history` reported `loss: 0.0` and carried no `frac_nonfinite_loss` key at all, which is the symptom this PR exists to remove. `PPOTrainer` stays exempt for a reason that does hold: it writes its own `train()` and never enters that loop. PRM now carries the same guard block as the other 25, byte-identical, and a test parametrized over NaN and Inf. Both mutants die: neutering the guard fails both cases, and swapping `~isfinite` for `isnan` fails the Inf case alone. The warning was asserted by none of the 27 non-finite tests, so deleting every `warning_once` call left the suite green. The regression that permits is specific: `warning_once` keys its cache on the message text, so collapsing the train and eval messages into one lets an evaluation warning suppress every later training warning. A test now drives both paths and asserts both messages. It clears the cache first, because that cache is process-wide and an earlier test emitting the same text would otherwise suppress the messages here and make the result depend on test order. The metric is documented as a rate averaged over every loss computation in the window, but all 27 tests ran at `gradient_accumulation_steps=1`, where each window holds exactly one computation and the averaging is invisible. Changing `sum(val) / len(val)` to `max(val)` passed everything. A test now poisons one micro-batch of two and asserts the metric reads 0.5. Two comments asserted mechanisms that are not true. `test_distributed.py` said the child's assertion "would be false" under one rank; with one rank the sole rank is also the poisoned one, so it reads 1.0 against an expected 1.0 and passes, which is the blindness the file exists to remove and what its own sibling docstring already said. The async-GRPO test explained its NaN-only injection by `inf - inf`, but only the model's log-probabilities are poisoned and they are differenced against rollout-time values no forward pass touches, so an injected Inf gives `exp(-inf)`, a finite 0.0, and never reaches the guard at all. `sdft` spelled the eval-prefix comprehension with `k` and `v` where the other fifteen copies use `key` and `val`. Same behaviour, but the duplicated blocks are required to match, so it is now spelled like its siblings.
The guard reduced the loss with `loss.detach().mean()`. A low-precision dtype has no `mean` kernel, so a `float8_e5m2` loss raised `NotImplementedError: "sum_cpu" not implemented for 'Float8_e5m2'` rather than reporting a rate. That turned a diagnostic into a crash on the dtypes most likely to produce a non-finite loss in the first place. Casting to float32 before the reduction fixes both float8 variants and leaves the finiteness of every other dtype unchanged. Applied at all 31 sites. The guard also ends each loss path in a collective, and ten raises upstream of it ran on one rank only. A rank that raised left its peers waiting inside a gather that never completed: measured on two ranks, the peer blocked 3 times out of 3. Those ten now agree across ranks before raising, so the run fails together instead of hanging. Four config-time raises stay rank-local on purpose, because a gather inside a branch only the raising rank enters would create the hang it prevents everywhere else. BCO stored the metric under `is_main_process`, discarding every other rank's contribution to a value the gather had already averaged across ranks. Its store line now matches CPO's and ORPO's. Two tests were passing against a mutation. The SFT rate test split one non-finite micro-batch against one finite one, so an inverted `~torch.isfinite` still averaged to 0.5 and the assertion could not tell the two apart; it now uses four micro-batches and asserts 0.25. The DPO Liger loss path had no test at all, so deleting its guard changed nothing observable. Both mutations now fail. The 15 doc pages said the filter substitutes a finite value unconditionally. It does that only when `is_torch_xla_available()` is false; under XLA the step is dropped instead.
The `doc-builder-style` pre-commit hook wraps docstring prose to fill `--max_len 119`, and its scope is `^(trl|tests|docs/source)/`, so the three docstrings added with these tests were short of the column it packs to and CI's code-quality job rejected them. This is the hook's own output rather than a hand rewrap, so a later run is a no-op. No wording changes, only line breaks.
The guard reduced the loss with `loss.detach().float().mean()`. Narrowing first is lossy: a finite `float64` loss above `float32`'s maximum becomes `inf`, so the guard reported a healthy step as non-finite and emitted the warning. Measured over 13 dtype and edge cases, the narrowing form is right in 11 of them and `.double()` in all 13. Widening also keeps both `float8` dtypes working, where testing finiteness elementwise instead raises `NotImplementedError: "isfinite" not implemented for 'Float8_e4m3fn'`. The DPO test added here fails against the old cast and passes against this one. DPO and KTO translate an image-token mismatch into an actionable `max_length` message, and that raise ran on one rank only. The guard's gather sits downstream of it, so a rank whose own batch failed left its peers inside a collective it never entered: measured on two ranks, the survivor blocked for the full collective timeout instead of failing with the message. SFT already agreed across ranks before raising; these two now do too. Their guard moves out of the two loss methods and into `compute_loss`, which is where the agreement can sit upstream of the guard's gather. Keeping the guard inside those methods instead pairs one rank's guard with another rank's failure flag: measured, the survivor read the peer's flag as its own loss value and then blocked anyway. That placement differs from the other 24 trainers, which have no rank-local raise upstream of their guard. Guard sites drop from 31 to 29, because DPO and KTO each had one per loss method and now have one per trainer. Coverage is unchanged at 26 trainers.
0cab401 put the failure agreement in `compute_loss`, downstream of the sixteen metric gathers that `_compute_loss` runs after its forward. A rank whose forward failed never reached those gathers, so the survivor's first gather paired with the failing rank's agreement flag and the run blocked in the next collective. The agreement now sits inside each loss method, directly after the forward and before any collective. KTO's reference forward gets an agreement of its own because it runs below the batch-size check's gather. Both `compute_loss` methods return to plain delegation and the non-finite guard moves into the callees, so it stays the last collective on every path. The Liger DPO test hooked the causal-LM wrapper's `logits`, but `_compute_loss_liger` reads `last_hidden_state` from the backbone, so the hook never fired and three CI lanes failed on it. The test now poisons the fused loss itself and covers Inf as well as NaN. The distributed check fails the forward from a pre-hook rather than raising in an override: an override that raised before delegating would skip the agreement it exists to test. It refuses to run on one rank, where every assertion passes whether or not the ranks agree, and its watchdog is a daemon thread: a stalled rank sits inside a C-level collective and never returns to Python bytecode, so a `SIGALRM` handler is delivered but never runs. Measured on two CPU ranks over MPI: the shipped code passes, the code with both agreements removed ends at the watchdog. GOLD's missing-field check agrees across ranks and so reads `accelerator`; the bare-trainer test stubs it, as its sibling in the same file already does. At fifteen guard sites the `mode` assignment had split the guard's comment in two; it moves above the comment.
The forward-failure agreement caught only `ValueError`, the class it translates into the image-placeholder message. A `RuntimeError` or an out-of-memory error in one rank's forward, the plainest rank-local failure there is, bypassed the agreement and left the peers waiting in the first metric gather. A probe with a failing pre-hook shows it: `ValueError` reached the agreement gather, `RuntimeError` and `OutOfMemoryError` reached none. Every agreement now catches `Exception` and keeps the `ValueError` translation as a special case. KTO's Liger path called the fused loss after its only agreement, so a rank-local failure in the kernel skipped every metric gather that follows. DPO keeps its fused loss inside the forward's `try`; KTO's depends on the KL gather in between, so it gets an agreement of its own, the same shape as the one around its reference forward. The tails of all six agreements are byte-identical. The DPO test for a finite loss above `float32`'s maximum proved nothing: it replaced the loss after `_compute_loss` returned, so the guard had already run on the real value, and mutating the guard's cast to `float32` left it green. The loss is now produced inside the trainer, with the model in `float64` and the IPO loss at `beta=1e-20`, about `2.5e39` whatever the log-ratios are. The `float32` mutant fails it. The distributed child's watchdog comment records the measured margin: the healthy two-rank CPU run finishes in under 90 seconds against a 300 second deadline. On the two-rank MPI harness the shipped code passes, and the code with the agreements removed stalls until the watchdog ends it.
The previous commit widened the DPO and KTO agreements from `ValueError` to `Exception`, so a `RuntimeError` or an out-of-memory error on one rank reaches the same gather as a failed batch. SFT's agreement, the one GOLD, SDFT, SSD and the other SFT subclasses inherit, still caught `ValueError` only, so any other exception on one rank skipped the agreement and left the peers waiting in the next collective. The catch block is now the same in all seven sites (SFT, DPO x2, KTO x4).
… dataset comments `SFTTrainer.evaluate` prepared each split of a dict and then handed the whole dict to `Trainer.evaluate`, which calls `self.evaluate` again per split and so re-entered the override on data it had just prepared. The tokenizing steps skip an already-tokenized split, but `packing` packs it a second time. The override now runs `super().evaluate` once per split itself, opening the eval window with that split's prefix, the same shape DPO's `evaluate` already has in this branch. A test counts one `_prepare_dataset` call per split and checks the per-split metric keys. The `evaluate` comments in DPO, KTO and TPO said `_prepare_dataset` is idempotent and skips already-tokenized data. None of the three has such a skip: each tokenizes from the text columns, and TPO refuses input without them. The comments now say the method has to receive the raw dataset, which is why a `str` is left untouched.
# Conflicts: # tests/experimental/test_async_distillation_trainer.py
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 4db11db. Configure here.
…ing count The distributed child now takes a case name and injects a rank-local failure at each of the seven agreement sites across DPO, KTO and SFT (standard and Liger paths, ValueError, RuntimeError and OutOfMemoryError). Before, only one DPO site was reached. Failure hooks register on the module DDP wraps, so the DDP buffer broadcast completes before the hook raises; a hook on the wrapper itself raises ahead of that collective. The non-finite warning test asserted message presence only, which an uncached warning satisfies; it now requires exactly one training and one evaluation warning.
Member
|
Thanks @behroozazarkhalili but I don't think I get the purpose here: we're adding a metric that only lights up once the run is already dead? In what real training did you end up finding this metric really useful? |
Collaborator
Author
|
You are right on both points. There was no real run. This came from the regression test for #6166. An assertion on I also agree on the timing. Without a grad scaler, a NaN loss still takes an optimizer step, so the run is already dead by the time anything is logged. Given that, 70 files is far more than this is worth, and I am fine closing it. If you see value in a small version, it would be a warning on the first filtered step in GRPO and RLOO only. It would have to be a new PR from my fork, because this branch now conflicts with #7136 and I can no longer push to it. It is your decision. Noted on the description. I will keep descriptions to a few sentences from now on. |
Member
|
Closing, I think it's important to work from real training. And fix real issue. The amount of edge / unreachable cases that we can find from just reading the codebase is infinite. |
This was referenced Sep 17, 2026
Collaborator
Author
|
That's a fair bar, and I agree with it. Six of my other open PRs had the same origin. I found them by reading the code, not because a run broke or a user reported a problem. I've closed all six (#6977, #7033, #7042, #7044, #7045, #7053) so they're out of your queue. I also closed the issues I had filed for them. From here on, I'll only open a PR when there's a failing run or a user report behind it, and the description will start with that evidence. I appreciate the direct feedback. |
This was referenced Sep 17, 2026
Closed
Closed
4 of 5 tasks
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
2 participants
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.
The training logger could hide a non-finite loss while backward still used it. I added a gathered
frac_nonfinite_lossmetric, warnings, and agreement across ranks for selected failure paths to prevent peers from hanging. Regression tests expose the missing metric on main, and distributed GPU checks pass for the forward-failure agreement, while local async tests remain blocked by an unavailable kernel build.Resolves #6702.
Note
Medium Risk
Touches every trainer’s hot loss path with new collectives and cross-rank failure agreement; incorrect placement could cause deadlocks or change logging semantics, though behavior is heavily tested.
Overview
Adds
frac_nonfinite_lossacross TRL trainers so NaN/Inf steps are visible in logs even when Transformers’ defaultlogging_nan_inf_filterkeeps the reportedlossfinite. Each loss path records a rank-gathered fraction of non-finite micro-batch losses, emits train/evalwarning_oncemessages, and documents the metric on trainer pages.Also fixes eval logging: metrics use the caller’s
metric_key_prefix(not a hardcodedeval_) andevaluate/predictopen a clean eval metrics window sopredict()batches do not leak into the next eval average. SFTevaluateon a dict of splits prepares each split once instead of double-preparing via nestedTrainer.evaluate.For DPO, KTO, and SFT, forward failures are agreed across ranks before metric gathers so one rank’s exception does not hang peers in collectives; distributed child tests cover gathered
frac_nonfinite_lossand peer-failure raises. Broad unit tests poison losses (NaN/Inf where applicable) and assertfrac_nonfinite_lossbehavior.Reviewed by Cursor Bugbot for commit f096119. Bugbot is set up for automated code reviews on this repo. Configure here.