[qgallouedec]

Resolves #7662

OffloadActivations used torch.cuda for every accelerator other than XPU/NPU, so it crashed on Trainium. It now uses torch.get_device_module (25 fewer lines). The offloading tests pass on H100. On Trainium the losses are identical, and Qwen3-0.6B SFT without gradient checkpointing trains instead of running out of memory.

cc @michaelbenayoun

---

Note

Medium Risk
Changes core stream/sync/offload paths for all accelerators; behavior should match prior CUDA/XPU/NPU logic but any get_device_module gap could affect training stability on less-tested hardware.

Overview
OffloadActivations now routes accelerator-specific work through torch.get_device_module instead of hard-coded CUDA/XPU/NPU branches, so devices like Trainium no longer incorrectly call CUDA APIs.

After resolving the current accelerator type, the manager caches device_module and uses it for comm streams, stream contexts, current_stream, empty_cache, and the “already on CPU / skip offload” device check (compare against accelerator_type instead of a fixed device list). NPU-only storage refcount paths are unchanged.

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

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-13T03:51:00.782069Z	e199ad9	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e199ad92ea

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

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