[jemitche1]

This PR replaces hard-coded torch.cuda calls with torch.accelerator, torch.Stream and torch.Event so EMA CPU offload, RL weight staging, the optimizer first-step sync and graph_trainer precompile no longer assume CUDA

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[tianyu-l]

@claude review this PR

[claude]

Claude finished @tianyu-l's task in 2m 7s —— View job

---

Review: Use torch.accelerator for device, stream and memory APIs

• Gather context / diff
• Review changed files
• Post review

The direction looks good. torch.Stream / torch.Event / torch.accelerator.* are the right device-agnostic replacements, and torchtitan targets PyTorch nightly, so the APIs exist. I found no correctness bugs in the CUDA path. Each change is a 1:1 swap:

• torch.cuda.stream(s) → with s:. The torch.Stream context manager also switches device.
• torch.cuda.current_stream(dev) → torch.accelerator.current_stream(dev).
• memory_stats(self.device) with an index-less device still resolves to the current device.

Leftovers are also clean. training_engine.py has no remaining device_type/device_module references after the local was removed. nullcontext is still used in rl/trainer/trainer.py. torchft/trainer.py already imports torch.

I could not run Python or the tests in this environment (the commands need approval), so these notes come from reading the code.

Main concern: out-of-tree / privateuse1 backends

The goal is to "work on any accelerator". But this PR moves the dependency from per-backend module functions to the torch.accelerator hooks and DeviceAllocator interface. Those are a newer contract, and some backends may not implement it yet:

1. tools/utils.py:62-66 get_device_info()
   • Before: _get_available_device_type() probed getattr(torch, privateuse1_name).is_available().
   • Now: torch.accelerator.current_accelerator() only reports a privateuse1 backend if it registered PrivateUse1HooksInterface.
   • Neuron: tools/utils.py:188-191 has a dedicated Neuron (torch.neuron) branch. If that backend doesn't register the hooks, device_type silently falls back to "cuda".
   • Default changed: with check_available=False (the default), a build compiled with a non-CUDA accelerator but no visible device now reports that accelerator instead of "cuda". Arguably more correct, but it is a behavior change.
2. observability/metrics.py:60-73: torch.accelerator.memory_stats / reset_peak_memory_stats / empty_cache require the backend's allocator to be registered as a DeviceAllocator. A backend that only provides torch.<backend>.memory_stats will now raise when DeviceMemoryMonitor is built. This sits on the default training path for every run.

It would help to note in the PR description which backends were validated (XPU? ROCm? Neuron?). If Neuron hasn't been checked, please check it, or ping the people who added the Neuron peak-flops branch.

Inconsistencies and leftovers

• components/optim/optim.py:113-115: the first-step torch.cuda.is_initialized() / torch.cuda.synchronize() / torch.cuda.empty_cache() is on the core training path and wasn't converted. This is the most visible miss given the PR's stated goal. A natural replacement is torch.accelerator.synchronize() + torch.accelerator.empty_cache(), but keep the guard so CPU-only runs don't break. Fix this →
• observability/metrics.py:53,55 still uses device_module.get_device_name / get_device_properties. That's fine since torch.accelerator has no equivalent, but the class now mixes two API families. A one-line comment would stop the next reader from "fixing" it halfway.
• distributed/flex_shard/_optimizer_reshard_runtime.py:266-275 still builds streams and events through torch.get_device_module(device).Stream/Event. Not wrong, but by Core Principle 3 (reuse over duplication) it's the same pattern this PR unifies elsewhere. Worth converting here or in a follow-up.
• distributed/cuda_graph.py: keeping torch.cuda.Stream there is correct, since it's CUDA-graph-specific.

Smaller notes

• rl/generator/vllm/model_wrapper.py:51,416: this imports device_type from tools.utils only for the NUMA call. Other callers (rl/trainer/trainer.py:120, graph_trainer/trainer.py:208) pass device.index, device.type from a real device. Taking the type from the same place as the index (e.g. torch.accelerator.current_accelerator().type) keeps the two consistent and avoids the new import.
• tools/utils.py:95-97: dropping the hasattr(device_module, "device_count") guard is fine, because torch.accelerator.device_count() returns 0 when there is no accelerator, which keeps the previous behavior.
• experiments/graph_trainer/paged_stash_memory_policy.py: the stash buffers are Triton/CUDA-graph based, so these changes are mostly cosmetic. That's fine, but the PR description says paged stash now "works on any accelerator", which is an overstatement.
• tests/rl/unit_tests/cpu/test_trainer.py:163: patch("torch.Event", ...) replaces torch.Event process-wide for the whole push_model_state_dict() call. It works, but patch("torchtitan.rl.trainer.trainer.torch.Event") makes the intent clearer. Same idea for torch.accelerator.current_stream at line 513. The test name ..._stages_on_dedicated_cuda_stream is also now slightly stale.
• tests/unit_tests/cpu/tools/test_tools_utils.py: patching torch.accelerator.device_count is fine with monkeypatch.

Verification

None of this should change numerics, but the device, stream and event setup touches every training run. Before merging:

• A short llama3 debug-model run (seed 42, deterministic) on CUDA showing identical loss and grad_norm.
• The GPU EMA CPU-offload test (tests/unit_tests/gpu/components/optim/test_ema.py) passing; it directly exercises the with stream: change.
• If possible, one run on a non-CUDA backend (XPU) to back up the "any accelerator" claim.

Overall LGTM in direction. Before landing, please (1) confirm or validate privateuse1/Neuron behavior and (2) convert optim.py.

[jemitche1]

Thanks @tianyu-l, the review was helpful. Good catch on out-of-tree backends: device_module already covers every backend, so converting it to torch.accelerator added risk for neuron without adding anything. I've reverted those changes, so the PR now only replaces hard-coded torch.cuda calls. I also converted the first-step sync in optim.py that I'd missed and fixed up the tests. On XPU, llama3 debug run matches main exactly for 10 steps; for cuda, I will rely on CI.

[tianyu-l]

@claude review again

[claude]

Claude finished @tianyu-l's task in 1m 31s —— View job

---

Re-review: Use torch.accelerator for device, stream and memory APIs

• Gather diff
• Review changed files
• Post review

The revised scope is much safer. device_module, get_device_info() and DeviceMemoryMonitor are back to how they were, so the privateuse1/Neuron concern from my first review no longer applies. Every remaining change replaces a hard-coded torch.cuda call, and the CUDA path is a 1:1 swap:

• torch.cuda.Stream() → torch.Stream() and torch.cuda.Event() → torch.Event(). Both default to the current accelerator and current device, as before.
• with torch.cuda.stream(s): → with s:. The torch.Stream context manager also switches the device, which matches what torch.cuda.stream does.
• torch.cuda.current_stream(dev) → torch.accelerator.current_stream(dev).
• Paged stash: _launch_stash and _launch_reload run inside eager custom_op bodies, not dynamo-traced code, so with stream: is fine there.
• No leftover references: _staging_device_module is gone everywhere, and nullcontext is still used in rl/trainer/trainer.py. tools/utils.py only imports torch, so the new import in optim.py adds no import cycle.

As before, I could not run the tests here, so these notes come from reading the code.

One behavior change worth fixing: components/optim/optim.py:114-116

if torch.accelerator.is_available():
    torch.accelerator.synchronize()
    device_module.empty_cache()

1. The guard is now looser. The old code synchronized only if torch.cuda.is_initialized(). torch.accelerator.is_available() is true on any host with a visible GPU, even when nothing has touched it. torch.accelerator.synchronize() then lazily creates a context on the current device, which defaults to device 0. This affects:

   • CPU-only unit tests (e.g. tests/unit_tests/cpu/components/optim/test_optim.py) run on a GPU host.
   • Any CPU-device run where GPUs happen to be visible.

   Normal training is unaffected, because the device is already set and initialized by the time step() runs. It is still a quiet change on the core training path.

2. It mixes two API families. Per your own reasoning that device_module already covers every backend, device_module.synchronize() would be consistent with empty_cache() here and with distributed/utils.py:490. Right now synchronize goes through torch.accelerator and empty_cache goes through device_module.

Suggestion: keep the lazy-init guard and use one API family, e.g.

if torch.accelerator.is_available() and device_module.is_initialized():
    device_module.synchronize()
    device_module.empty_cache()

I haven't checked that every backend's device module exposes is_initialized(). torch.cuda and torch.xpu do. Fix this →

Leftovers (optional, follow-up is fine)

• rl/generator/vllm/model_wrapper.py:415: maybe_apply_numa_binding(torch.cuda.current_device(), "cuda"), and :712: torch.cuda.empty_cache(). These are the remaining hard-coded calls on the RL weight path this PR touches. generator.py now uses torch.accelerator.current_device_index, so these two stand out.
• rdma_microbatch.py:198, vllm/attention.py:119 and paged_stash_memory_policy.py:1089 (is_current_stream_capturing) are genuinely CUDA-specific. Keeping them as they are is correct.

Nits

• precompile_main.py:106: torch.accelerator.set_device_index(device) accepts a torch.device, so this works. Since the function name says index, passing device.index (or 0) would read better.
• tests/rl/unit_tests/cpu/test_trainer.py: correction to my earlier note. patch("torchtitan.rl.trainer.trainer.torch.Event") still patches the global torch.Event, because trainer.torch is the torch module, so it doesn't narrow the scope. It's harmless, and the test rename is good.
• PR description: it still says EMA, paged stash and the rest "no longer assume CUDA". Paged stash is still Triton/CUDA-graph based, so "no longer hard-code torch.cuda for streams and events" would be more accurate.

Verification

The XPU 10-step parity result is good to have. For CUDA, the changes affect stream ordering in EMA offload and RL staging, so please confirm these CI jobs actually run green, since the ciflow/rl and ciflow/fake-pg workflows were still awaiting approval:

• tests/unit_tests/gpu/components/optim/test_ema.py, which tests the race in TestEMACpuOffload
• the RL GPU tests
• the paged-stash tests

LGTM once the optim.py guard is fixed. Everything else is optional.

[jemitche1]

Thanks @tianyu-l. I fixed the optim.py guard: it now checks torch.accelerator.is_available() and device_module.is_initialized() and uses device_module.synchronize(), matching the original behavior.

[tianyu-l]

@claude check again

[claude]

Claude finished @tianyu-l's task in 44s —— View job

---

Third review: Use torch.accelerator for device, stream and memory APIs

• Gather diff (origin/main...HEAD, 1 commit b45d7f4)
• Re-check the optim.py guard and earlier suggestions
• Check CI status
• Post review

Verdict: LGTM. The one blocking item from last time is fixed, and I found nothing new. I could not run the tests here, so this comes from reading the code.

components/optim/optim.py:115-117 is fixed

if torch.accelerator.is_available() and device_module.is_initialized():
    device_module.synchronize()
    device_module.empty_cache()

• Lazy-init guard is back. A CPU-only run or a CPU unit test on a GPU host no longer creates a device context.
• On a CPU-only host, device_module falls back to torch.cuda. is_available() is false there, so the condition stops before calling is_initialized(). That is safe.
• One API family. Both calls now go through device_module, which matches distributed/utils.py.
• Small behavior change: empty_cache() used to run outside the guard and now runs inside it. For CUDA this makes no difference, because torch.cuda.empty_cache() already does nothing when CUDA isn't initialized.
• No import cycle: tools/utils.py only imports torch.

Earlier suggestions

• precompile_main.py:104-105: now uses torch.device(utils.device_type, 0) and set_device_index(device.index). The utils import is already there (line 47).
• The test is renamed to ..._stages_on_dedicated_stream. _staging_device_module no longer appears anywhere in the repo. nullcontext is still used in rl/trainer/trainer.py:450 and is correctly removed from the test file.
• In the test, with stream: on a MagicMock stream works because MagicMock supports the context-manager methods.
• The remaining CUDA calls in rl/generator/vllm/model_wrapper.py (torch.cuda.current_device() and torch.cuda.empty_cache()) are still there. As noted before, a follow-up is fine.

CI

All runs were queued or in progress when I checked. Nothing had passed or failed yet. GraphTrainer H100 was skipped. Before merging, please confirm these come back green, since they directly test the stream changes:

• GPU Unit Tests, which includes test_ema.py's TestEMACpuOffload
• RL GPU/CPU Unit Tests
• GraphTrainer 8 GPU Integration Tests, which covers paged stash

[tianyu-l]

@jemitche1 CPU failure seems only happening on this PR, could you help investigate?

[jemitche1]

Thanks for flagging! test_optim.py was still mocking the old torch.cuda calls in the first-step block. I updated it to use torch.accelerator and device_module instead, and moved empty_cache() back outside the guard to match the original behavior. The CPU unit tests pass locally against nightly.

[tianyu-l]

thanks, there seems a conflict -- could you rebase