Skip to content

[SFT] Treat past_key_values and attentions as optional on the model output - #6686

Open
Hakureirm wants to merge 2 commits into
huggingface:mainfrom
Hakureirm:fix-chunked-ce-optional-cache-fields
Open

[SFT] Treat past_key_values and attentions as optional on the model output#6686
Hakureirm wants to merge 2 commits into
huggingface:mainfrom
Hakureirm:fix-chunked-ce-optional-cache-fields

Conversation

@Hakureirm

@Hakureirm Hakureirm commented Aug 8, 2026

Copy link
Copy Markdown

What does this PR do?

_chunked_ce_forward reads outputs.past_key_values and outputs.attentions
unconditionally when it rebuilds the output, so a model whose forward returns
neither raises AttributeError on the first step. loss_type defaults to
"chunked_nll", so SFTTrainer fails out of the box for those models.

Neither field is ever populated here — the backbone is called with
use_cache=False and attentions are never requested — so they are passed as
None. Measured on trl-internal-testing/tiny-Qwen2ForCausalLM-2.5, at the
point the output is rebuilt, both are already None today. last_hidden_state
and hidden_states, the fields the chunked path actually reads, are present on
every output type involved.

Ten entries of MODEL_FOR_CAUSAL_LM_MAPPING_NAMES (transformers 5.14.1) declare
a return type without past_key_valuesmamba, mamba2, falcon_mamba,
xlstm and rwkv keep the recurrence under cache_params or state, xlnet
under mems, and openai-gpt, recurrent_gemma, reformer and xlm have no
state field at all.

Test. trl-internal-testing/tiny-FalconMambaForCausalLM is added to
TestSFTTrainer::test_train, which uses a default SFTConfig and therefore the
default loss_type; it fails on main and passes here, on CPU and on an A800.

No behaviour change for models that carry the fields. Training
trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 for 3 epochs under a fixed seed
on CPU gives byte-identical per-step losses and parameter sums with and without
this change.

tests/test_sft_trainer.py on this branch: 112 passed / 153 skipped on CPU,
187 passed / 78 skipped on an A800, no failures or errors either way.

Fixes #6685

Before submitting

AI writing disclosure

  • No AI usage: the PR was written entirely by a human.
  • AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
  • AI-generated: the PR was mostly or fully generated by an AI tool.

…utput

`_chunked_ce_forward` reads `outputs.past_key_values` and `outputs.attentions`
unconditionally when rebuilding the output. Models whose forward returns neither
raise `AttributeError` on the first step, and `loss_type` defaults to
`"chunked_nll"`, so `SFTTrainer` fails out of the box for them.

Ten model families in the causal-LM auto-mapping declare a return type without
`past_key_values` (mamba, mamba2, falcon_mamba, xlstm and rwkv keep the
recurrence under `cache_params` or `state`; xlnet uses `mems`); none of them
carries `attentions` either. `last_hidden_state` and `hidden_states`, the two
fields the chunked path reads elsewhere, are present on all of them, so reading
the other two with `getattr` is enough.

Covered by adding `trl-internal-testing/tiny-FalconMambaForCausalLM` to
`TestSFTTrainer::test_train`, which uses the default `loss_type` and so fails
before this change and passes after.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit eb54dda. Configure here.

Comment thread trl/trainer/sft_trainer.py Outdated
Review feedback: `getattr` conflicts with the repo's own rule against
`hasattr`/`getattr` (`.ai/AGENTS.md`), and the two fields are never populated
here anyway — the backbone is called with `use_cache=False` and never asked for
attentions. Measured on `trl-internal-testing/tiny-Qwen2ForCausalLM-2.5`: at the
point the output is rebuilt, both are `None`.

Passing them as `None` keeps this path independent of whether a given model type
declares the fields at all, which is the property the FalconMamba test covers.
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.

SFTTrainer's default chunked_nll path fails on models whose output has no past_key_values

1 participant