Clarify stream-only dedupe check

[kashif]

Makes the storage dedupe check explicitly streams-only. Follow up from #6241

This matches the existing comment: fwd_stash gives streams mode a clear release point, while single-stream mode does not.

Selective activation checkpointing (SAC) for SFT

Also adds selective activation checkpointing to SFT. When gradient checkpointing is on, the attention output is saved during the forward pass instead of being recomputed in the backward pass. At long context this recovers most of the checkpointing slowdown for one extra hidden-state sized tensor per layer, and it runs fully eager with no torch.compile needed.

Turn it on with selective_activation_checkpointing=True on SFTConfig (needs gradient_checkpointing=True, which SFT already defaults to). It works by wrapping the model's gradient_checkpointing_enable to inject a policy that saves attention and recomputes the rest, and it forces non-reentrant checkpointing which SAC relies on. This is the same eager SAC approach torchtitan uses under FSDP2. PEFT plus DeepSpeed ZeRO-3 is rejected up front since that combo needs reentrant checkpointing.

Checked that SAC gradients match full checkpointing exactly and that attention is no longer recomputed in backward. Added a test for both.

This is the SFT-only first step. DPO, KTO, reward and GRPO can follow the same two lines in a later PR.

---

Note

Medium Risk
SAC changes the gradient-checkpointing path for all SFT runs that opt in; incorrect policy could affect training correctness, though tests compare gradients to full checkpointing. The offloading dedupe tweak is narrow but touches hot memory paths.

Overview
Adds selective activation checkpointing (SAC) for SFT: SFTConfig.selective_activation_checkpointing wraps the model’s gradient_checkpointing_enable so checkpointing saves attention outputs and recomputes other ops (non-reentrant, eager policy via PyTorch’s selective checkpoint contexts). SFTTrainer enables this before the base trainer init, requires gradient_checkpointing=True, and rejects SAC together with PEFT + DeepSpeed ZeRO-3.

Exports enable_selective_activation_checkpointing from trl.models. New tests assert SAC gradients match full checkpointing and that SDPA runs fewer times under SAC.

Separately, activation offloading only skips duplicate storage offloads when use_streams is true, matching the intended fwd_stash release behavior in single-stream mode.

^{Reviewed by Cursor Bugbot for commit 600adf1. Bugbot is set up for automated code reviews on this repo. Configure here.}

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR makes the activation offload storage-deduplication guard in OffloadActivations.pack_tensor explicitly apply only in streams mode, aligning the condition with the intended lifecycle management of GPU references via fwd_stash.

Changes:

• Restricts the “already-offloaded storage” dedupe fast-path to use_streams=True by gating the membership check with self.use_streams.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[qgallouedec]

cc @winglian

[bot-ci-comment]

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.

[cursor]

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 aa6b8d0. Configure here.}