Make a non-finite training step visible in log_history - #6862
Open
behroozazarkhalili wants to merge 19 commits into
Open
Make a non-finite training step visible in log_history#6862behroozazarkhalili wants to merge 19 commits into
behroozazarkhalili wants to merge 19 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. |
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.
| eval_dataset=dataset, ignore_keys=ignore_keys, metric_key_prefix=f"{metric_key_prefix}_{name}" | ||
| ) | ||
| ) | ||
| return metrics |
There was a problem hiding this comment.
Dict eval from init still re-prepares
Medium Severity
The new per-split super().evaluate loop only runs when a dict is passed in as eval_dataset. trainer.evaluate() with no args uses self.eval_dataset; if that is a dict, Trainer.evaluate re-enters this override per split on already-tokenized data, and _prepare_dataset re-packs when packing is on. The new test only covers the argument path, so the usual init-time dict eval still double-prepares.
Reviewed by Cursor Bugbot for commit 4db11db. Configure here.
Collaborator
Author
There was a problem hiding this comment.
Measured on this head (4db11db, transformers 5.11.0) with _prepare_dataset wrapped in a call counter, eval_dataset={"a": zen_test, "b": zen_test} passed at init, then trainer.evaluate() with no argument:
packing=False prepare_calls: init=3 evaluate()=0 evaluate(dict)=2 rows before/after evaluate(): {'a': 2, 'b': 2} -> {'a': 2, 'b': 2}
packing=True prepare_calls: init=3 evaluate()=0 evaluate(dict)=2 rows before/after evaluate(): {'a': 1, 'b': 1} -> {'a': 1, 'b': 1}
The no-argument call prepares nothing and re-packs nothing. The reason is the override flag in Trainer.evaluate (transformers/trainer.py:2581-2587, same shape at 4.56.2): when eval_dataset is None it recurses with eval_dataset=<split name>, a str, and this override leaves str inputs untouched (the not isinstance(eval_dataset, str) guard at the top of the preparation block). Only the explicit-dict argument path prepares, once per split, which is the case this PR's loop handles.
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.
Resolves #6702.
Where to start reading
The diff is large because the guard is duplicated into 26 trainers on purpose. Five places carry the design; the rest is copies.
trl/trainer/sft_trainer.py, the canonical guard block. Every other trainer's block is this one re-indented, checked byte for byte.trl/trainer/dpo_trainer.py_compute_loss, the forward-failure agreement: why a rank-local exception has to be agreed on before the first collective, and how.trl/trainer/kto_trainer.py_compute_loss, the same agreement where the loss path has two rank-local forwards and a gather between them.tests/distributed/nonfinite_loss.py, the multi-rank test of the gather itself.tests/distributed/rank_local_forward_failure.py, the multi-rank test of the agreement, with a note on why its watchdog is a thread.The problem
A non-finite training loss does not reach
log_history. Withlogging_nan_inf_filterenabled, which is the default, andis_torch_xla_available()false,transformersdiscards the step's own loss and logs a substitute, while the backward pass still runs on the non-finite value. The curve stays plausible and nothing identifies the step that failed.The substitution is worth stating precisely, because the option's name suggests otherwise and I had it wrong myself at first. It is not an average of previously logged losses.
transformerscomputestr_loss = tr_loss + tr_loss / (1 + global_step - _globalstep_last_logged), which adds a fraction of the current accumulator to itself and never reads logged history.This costs more in the RL trainers than in plain SFT, because several paths reach a non-finite loss from data rather than from divergence: an unscorable token logprob returned by vLLM (#6166), a degenerate reward group, an overflowing
expof a KL log-ratio (#3015). A loss curve cannot point at a single poisoned group or token.The change
Every trainer's loss path now appends
frac_nonfinite_lossand warns once per mode. That is 31 guard sites across 26 trainers, gathered across ranks before the value is recorded, because a non-finite loss confined to one rank otherwise reads as0.0from the main process and reproduces the same invisibility one layer up.Trainers here are self-contained by design, so the block is duplicated rather than abstracted. A checker re-renders all 31 from one source and diffs them byte for byte; they differ only in indentation and in which of the four metric mechanisms the trainer already uses.
The warning carries no step index.
warning_oncecaches on its arguments, so an interpolated step number misses the cache and warns every step. Measured onaccelerate1.13.0: the constant message emits once from five calls, the interpolated one emits five times.Two pre-existing metric defects, fixed
Both surfaced because
frac_nonfinite_lossis the first metric whose value a reader would notice is wrong.Eval metrics were filed under a hardcoded
eval_prefix rather than the caller'smetric_key_prefix, which istestforpredict(). Andpredict()never calledlog(), so the batches it scored survived into the next evaluation window and skewed its averages.Rank-local raises
Adding a collective to the loss path turns any rank-local
raiseupstream of it into a hang rather than a clean failure: the raising rank leaves while its peers block in the gather until the NCCL timeout. Measured on two H100s, three repetitions:Arm B is the one that shaped the fix. Having a gather nearby does not rescue the peer; only making the raise itself collective does. The branch takes the count of rank-agreeing raises from 6 to 29, so 23 are added: 16 make an existing rank-local raise collective, and 7 agree on a forward failure before any collective runs (next paragraph).
The agreement has to precede the operation it guards. In CPO and ORPO it initially sat below the cross-entropy that raises on the very shape mismatch it checks for, which left the peer waiting anyway. Reproduced as
ValueError: Expected input batch_size (12) to match target batch_size (14)on torch 2.11.0, then fixed by moving the gather above the call in both.DPO and KTO have a second shape of the same problem. Each translates a forward failure (a truncated image placeholder) into a
ValueErrorthat names the fix, and each loss method runs a long run of metric gathers between its forward and its return, sixteen in_compute_loss. A rank whose forward failed never reaches those gathers, so the surviving rank's first metric gather pairs with nothing, or with the failing rank's later agreement flag, and blocks. My first placement agreed on the failure incompute_loss, downstream of all of them, which Bugbot caught. The agreement now sits inside_compute_lossand_compute_loss_liger, directly after the forward and before any collective; KTO's reference forward runs below its batch-size check's gather and carries an agreement of its own, and so does its fused Liger loss, which runs below the KL gather; DPO's fused loss sits inside the forward's owntry.compute_lossis plain delegation and the non-finite guard is the last collective on every path.The agreement catches
Exception, notValueErroralone. The first version caught only the class it translates, so aRuntimeErroror an out-of-memory error in one rank's forward, the plainest rank-local failure there is, bypassed the agreement and left the peers in the first metric gather. A probe with a failing pre-hook shows the difference:ValueErrorreached the agreement gather,RuntimeErrorandOutOfMemoryErrorreached none.tests/distributed/rank_local_forward_failure.pyfails the forward on the last rank from a pre-hook and asserts on every rank that the run ended in an exception: the failing rank re-raises its own error, the others raise the peer-failure error. It refuses to run on one rank, where every assertion passes whether or not the ranks agree. Its watchdog is a daemon thread rather than a signal handler: a stalled rank sits inside a C-level collective and never returns to Python bytecode, soSIGALRMis delivered but never handled. Measured on two CPU ranks over MPI: the shipped code passes; with both agreements removed the survivor stalls until the watchdog ends it. The same check on two GPUs over NCCL is queued and will be reported when it runs.Verification
Each of the 26 tests poisons the second of two steps and leaves the first as the finite control, then asserts
0.0then1.0onlog_historywith the reportedlossstill finite. All 26 fail onmainwithKeyError: 'frac_nonfinite_loss'.A suite that only injects NaN is also passed by an implementation testing
isnaninstead ofisfinite, which would report a clean rate for an infinite loss. Where the poison can land on the loss scalar, 19 of the tests are now parametrized over NaN and+infboth. The remaining 7 poison logits or log-probabilities upstream of the loss, where an injected+infdoes not survive as a distinct Inf, so they inject NaN and say why in place.Which group a trainer belongs to was measured, not assumed. Substituting
+infand mutating the guard toisnanleaves ssd, tpo, dpo, kto, online_dpo and reward passing, so their loss really is NaN either way. a2po fails that same mutation and fails the complementaryisinfmutation too, so both values reach its loss distinctly and it is parametrized with the other 18. Async-GRPO is the one case where the collapse is not to NaN: it poisons the model's log-probabilities, which are then differenced against rollout-time values no forward pass touches, so an injected+infyieldsexp(-inf), a finite0.0, and never reaches the guard at all. It injects NaN for that reason rather than the shared one._open_eval_windowis defined in 17 trainers and had no test. A new one evaluates undereval,testandholdoutand asserts the metric arrives under each. Hardcoding the prefix back toevalfails two of the three.The warning had no test at all, so deleting every
warning_oncecall left the suite green. The regression that permits is specific:warning_oncekeys 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, and fails when the two messages are merged. It clears that cache first, because it is process-wide and an earlier test emitting the same text would otherwise decide the result by test order.The metric is documented as a rate averaged over every loss computation in the window, but every test ran at
gradient_accumulation_steps=1, where a window holds one computation and the averaging is invisible. Replacingsum(val) / len(val)withmax(val)passed everything. A test now poisons one micro-batch of two and asserts the metric reads 0.5, and fails under that replacement.The gather itself needs more than one process to test, since with a single rank it is the identity and deleting it leaves every single-process test green.
tests/distributed/nonfinite_loss.pypoisons one rank and asserts on all of them that the logged fraction covers the whole world;test_distributed.pylaunches it.One DPO test asserted that a finite loss above
float32's maximum is not reported as non-finite, and proved nothing: it replaced the loss after_compute_lossreturned, so the guard had already run on the real value, and mutating the guard's cast tofloat32left it green. It now produces the large loss inside the trainer, with the model infloat64and the IPO loss atbeta=1e-20, which is about2.5e39whatever the log-ratios are. Thefloat32mutant fails it.A GOLD test builds a trainer with
__new__and never runs__init__, so it has noaccelerator; it now stubs one with agatherthat returns its input, the way the neighbouring bare-trainer tests already do, and still reaches the missing-field error it asserts.Suite run on an H100: 25 suites collected, 23 pass. The two async suites are unevaluable on that node rather than failing on this change: both trainers hardcode
attn_implementation="kernels-community/flash-attn3"(added in #6083, untouched here), the kernel ships notorch211-cu129build, and both tracebacks terminate in__init__atfrom_pretrainedwithout ever reachingcompute_loss. Two pre-existing tests from the same classes fail identically on the same node.ruff check,ruff format --checkand the pinneddoc-buildergate are clean, each with a negative control confirming the gate can still fail.Limits worth knowing before merge
Four rank-local raises remain rank-local:
cpo_trainer.py:774,orpo_trainer.py:731,bco_trainer.py:1165andiw_opd_trainer.py:1253. All four predate this branch, are present onmain, andgit blamecredits them to earlier PRs rather than to this one. Making them collective is the same edit as the 16 above and I am happy to do it, but it widens a diff that is already large and none of them is reachable from a path this PR adds.The forward-failure agreement covers SFT, DPO and KTO, the three trainers whose loss paths translate a forward error into a message of their own. GOLD, GKD, TPO, MiniLLM, server-distillation, Reward, GRPO and RLOO run their forward rank-local and reach the metric gather below it without one, so a rank that fails inside those forwards leaves its peers waiting, as it did on
main, where the backward all-reduce sat below the same forward. Giving those eight the SFT block is the natural next step and is left out of this branch for the same reason as the four raises above.Three Liger branches have no non-finite test: GKD gates on
use_liger_kernel, IW-OPD and SDPO onuse_liger_loss, and none of their non-finite tests references Liger. DPO's Liger guard does have one, GPU-gated, which poisons the fused loss itself; its first version hooked the wrapper'slogits, which the Liger path never reads, and CI failed it in three lanes before it was fixed. Liger needs a GPU that this environment cannot give it, so the other three are measured as a gap rather than covered by a test I could not run.One of the 26 guards has a multi-rank test rather than all 26. The gather block is byte-identical across every copy, which the drift checker confirms at 31 of 31, so the single multi-process test exercises the shared logic rather than one trainer's version of it. The same holds for the evaluation-window machinery:
_open_eval_windowis byte-identical in all 17 copies and the prefix comprehension in all 17, so the one prefix test covers the shared code.Under DeepSpeed ZeRO-3 the agreement reaches less far than elsewhere. The forward itself all-gathers parameters, and the fused-loss paths gather the
lm_headweight, so a rank that fails inside the forward has already left its peers in a parameter collective before the agreement runs. The agreement covers every collective after the forward; the ones inside it belong to the sharding backend and no placement in the loss path can cover them.PPOTrainerdeliberately has no guard. It subclasses_BaseTrainer, defines its owntrain()and has neithercompute_lossnortraining_step, sologging_nan_inf_filternever runs against it and the failure this PR makes visible cannot occur there.PRMTrainerlooked like the same case and is not. It defines no loss of its own either, but it inheritstransformers.Trainer.compute_lossand the standard training loop, which is exactly where the filter applies. I had excluded it on the reasoning that it owns no loss site; owning one and being subject to the filter are different questions. Measured before the fix, with a NaN injected on step 2:log_historyreportedloss: 0.0and carried nofrac_nonfinite_losskey at all. It now carries the same block as the other 25 and a test over both values.The async non-finite test does not run outside CI. On a local H100 it stops at
FileNotFoundErrorforkernels-community/flash-attn3, which ships notorch211-cxx11-cu129build variant, so its behaviour rests on CI rather than on a measurement I can show.What I did not do
Option 3 in the issue was to name the cause rather than the location. That means instrumenting each candidate sub-tensor separately, a much larger change than the observability gap warrants, so this reports the location and leaves the cause to the warnings upstream of it.
I kept this as a warning rather than a raise. A non-finite step is survivable,
transformersdeliberately continues through one, and raising would change behavior for every current user. If you would rather it raise, this diff is where that goes, and the metric is what tells you how often it would fire.cc @qgallouedec @kashif
Note
Medium Risk
Wide changes to shared loss/logging paths and new distributed collectives/agreement logic; regressions could affect metric correctness or multi-GPU hangs, but behavior is heavily tested and mostly additive observability.
Overview
Adds
frac_nonfinite_lossacross TRL trainers so non-finite training steps are visible inlog_historyeven when Transformers’ defaultlogging_nan_inf_filter(non-XLA) replaces the reportedlosswith a finite substitute while backward still runs on the bad value.Each loss path records one flag per rank per micro-batch, gathers across ranks before averaging, checks finiteness in
float64(so huge but finite losses are not misclassified), emits mode-specificwarning_oncemessages (train vs eval), and documents the metric on the major trainer pages.Distributed / eval fixes bundled with the metric: rank agreement before metric gathers when a forward fails (notably DPO/KTO and related paths) so peers do not hang in collectives;
_open_eval_windowso eval metrics use the caller’smetric_key_prefix(testforpredict(), not hardcodedeval_) andpredict()leftovers do not pollute the next eval window; SFT eval dict handling so each split is prepared once.Coverage includes per-trainer poison tests, distributed child scripts for gather and forward-failure behavior, and edge cases (Liger DPO path, gradient-accumulation rate, large finite IPO loss).
Reviewed by Cursor Bugbot for commit 77e5495. Bugbot is set up for automated code reviews on this repo. Configure here.