[albertvillanova]

This PR restores activation offloading on torch 2.2 and 2.3, which #7670 broke on main.

Motivation

#7670 calls torch.get_device_module in OffloadActivations.__init__. That function only exists from torch 2.4, so #7670 relies on the torch floor from #7642, but it merged before #7642. TRL still allows torch 2.2 and 2.3 (transformers 4.56.2 requires torch>=2.2). On those versions, building the activation offloading context manager raises AttributeError: module 'torch' has no attribute 'get_device_module'. This affects SFT with activation_offloading=True and direct calls to get_act_offloading_ctx_manager.

#7670 is not released yet, but v1.16 is imminent, and #7642 can't merge before it. Without this fix, v1.16 would break activation offloading for users on torch 2.2 and 2.3, which v1.15 supports. On torch 2.3.1, tests/test_activation_offloading.py passes on v1.15.0 and fails on main.

Solution

Below torch 2.4, the accelerator type is always "cuda" (the existing fallback when torch.accelerator is missing), so the device module falls back to torch.cuda, as before #7670. #7642 removes this guard together with the torch.accelerator fallback.

Changes

• Use torch.cuda as the device module in activation offloading on torch < 2.4

---

Note

Low Risk
Single version guard in offloading init; behavior on torch 2.4+ is unchanged and older torch regains the prior cuda device module path.

Overview
Restores activation offloading on PyTorch 2.2 and 2.3 after a recent change unconditionally called torch.get_device_module, which only exists from 2.4 onward and caused AttributeError when constructing OffloadActivations (e.g. SFT with activation_offloading=True).

OffloadActivations.__init__ now sets device_module via is_torch_version(">=", "2.4.0"): use torch.get_device_module on 2.4+, otherwise torch.cuda, matching the pre-change behavior when accelerator_type already falls back to "cuda" without torch.accelerator.

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

[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.