[kahlun]

What does this PR do?

FSDPCheckpointManager gates FSDP's sharded-state-dict CPU offload on is_cuda_available:

state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)

is_cuda_available is torch.cuda.is_available() evaluated once at import
(verl/utils/device.py). On any non-CUDA accelerator it is permanently False, so the gathered
model and optimizer shards stay resident in device memory for the whole torch.save window
instead of being staged through host memory. CUDA gets the memory headroom; every other backend
pays for it, with no way to opt in.

This switches the four sites to is_device_available() — the platform-abstraction check
(get_platform().is_available()) that this same class already uses to decide whether accelerator
RNG state is reachable in BaseCheckpointManager.get_rng_state.

Why this is a per-device-class oversight rather than a deliberate CUDA-only choice: the
offload_to_cpu lines predate the platform abstraction layer (present at least as of
2026-04-20; is_device_available arrived 2026-06-04 in #6086), and the migration to it has been
incremental — #7871 applied the same is_device_available() treatment to the Megatron engine's
seed path in September. offload_to_cpu is a plain PyTorch FSDP feature with no CUDA-specific
code behind it, so there was never a technical reason to exclude other backends. The measurement
below confirms it works on a non-CUDA backend today.

Checklist Before Starting

• Search for similar PRs:

  • https://github.com/verl-project/verl/pulls?q=offload_to_cpu
  • https://github.com/verl-project/verl/pulls?q=is_device_available
  • https://github.com/verl-project/verl/issues?q=offload_to_cpu

  I enumerated every open PR in the repo and checked its changed-file list: [ckpt] fix: load FSDP checkpoint shards onto CPU instead of the saved device #7286 is the only
  open PR that touches verl/utils/checkpoint/fsdp_checkpoint_manager.py.

  Not a duplicate of [ckpt] fix: load FSDP checkpoint shards onto CPU instead of the saved device #7286 ([ckpt] fix: load FSDP checkpoint shards onto CPU instead of the saved device, open). That PR is the closest neighbour — same file, same "treat sharded state
  dicts as CPU-resident" theme — but a disjoint change: it adds map_location="cpu" to the three
  torch.load calls in load_checkpoint, which is wrong on CUDA too and fixes where tensors are
  revived from disk. This PR changes the ShardedStateDictConfig / ShardedOptimStateDictConfig
  gather configs, which is CUDA-privileged by construction and governs where tensors land during
  the FSDP gather, on both the save and load paths. The two touch different lines and neither
  subsumes the other; they merge independently in either order. Worth noting for whoever lands the
  second one: together they make the manager's CPU-residency story consistent end to end, which is
  the stated motivation of [ckpt] fix: load FSDP checkpoint shards onto CPU instead of the saved device #7286.

  Related precedent, not overlapping: [hardware] fix: add device_available check for devices other than GPU/NPU #7871 (merged) added an is_device_available() check for
  non-GPU/NPU devices in verl/workers/engine/megatron/utils.py. Same class of fix, different
  file; it did not touch the checkpoint managers.

• PR title formatted as [{modules}] {type}: {description}

Test

Measured on real hardware: 2× Intel Arc Pro B60 (XPU), torch 2.11.0, xccl backend.

The standalone repro below wraps an 8-layer 4096-wide model in FSDP across 2 ranks, runs one
optimizer.step() so optimizer state exists, then performs the same sharded gather twice —
once with offload_to_cpu=False (today's behaviour on XPU) and once with True (this patch) —
reporting the resulting tensor device and the peak device memory during each gather:

device=xpu world_size=2 free/total=22057000960/25669140480

[offload_to_cpu=False] sample tensor 'net.0.weight' device=xpu:0  peak_device_mem_during_gather=2816.8 MiB  optim_sd_keys=16
[offload_to_cpu=True]  sample tensor 'net.0.weight' device=cpu    peak_device_mem_during_gather=2560.8 MiB  optim_sd_keys=16

RESULT offload_False_device=xpu:0 offload_True_device=cpu
RESULT peak_delta_MiB=256.1
PASS: offload_to_cpu mechanism confirmed functional on this device

Two things this establishes:

1. offload_to_cpu=True works correctly on a non-CUDA backend. The gathered shard lands on
   cpu and the optimizer state dict is intact (16 entries, matching the 16 parameters). Nothing
   about this FSDP feature is CUDA-specific, so the current gate is excluding a working code path.
2. The cost of the current False default is real and scales with model size. 256 MiB of extra
   peak device memory held through the gather, on the order of the per-rank sharded state-dict size
   for this model — i.e. proportional to model/optimizer size, not a fixed overhead. On a card
   already near capacity this is the difference between a checkpoint save that completes and one
   that OOMs, which is the failure mode [ckpt] fix: load FSDP checkpoint shards onto CPU instead of the saved device #7286 describes from the load side.

On CUDA this patch is a no-op by construction, which is why no CUDA regression run is included:
PlatformCUDA.is_available() is return torch.cuda.is_available() — the identical expression the
old module-level is_cuda_available captured. The only semantic difference is import-time vs.
call-time evaluation, which cannot differ in practice since CUDA availability is fixed for a
process. On a CPU-only host both spellings yield False, so the *_on_cpu checkpoint tests are
unaffected.

Lint: ruff==0.12.2 check and ruff format --check clean (pinned to the repo's
.pre-commit-config.yaml rev; note ruff format reflows two now-shorter ternaries, included in
the diff).

Not covered: tests/special_distributed/test_fsdp_ckpt.py — the repo's own end-to-end
FSDPCheckpointManager save/load round trip — has not been run with this patch. It needs 2
accelerators and a Qwen2.5 tokenizer, and on XPU additionally a verl image with
flash_attention_2 available. What is established here is the mechanism and its memory cost on
real XPU silicon, not a full verl checkpoint round trip through the patched manager. That test is
the natural place to settle it if you would like it settled before merging, and I am happy to run
it on the B60s on request.

Known trade-off, stated explicitly: this makes non-CUDA backends stage shards through host
memory, so it trades device-memory pressure for host-memory pressure — exactly the trade CUDA has
always made here. On nodes where host RAM is the scarcer resource this is not free. I considered
exposing it as a checkpoint.offload_to_cpu config override instead, and deliberately did not:
that adds a user-facing knob to work around what is currently a one-platform exclusion from an
existing default, and the right first step is to make the default consistent across backends. If
maintainers would prefer the knob (either instead of or on top of this), say so and I will follow up
separately.

Repro script (2 accelerators, torchrun, no verl import needed)

#!/usr/bin/env python3
"""Repro: FSDP sharded state_dict() offload_to_cpu behavior on XPU vs CUDA.

Run with: torchrun --standalone --nproc_per_node=2 repro_offload_to_cpu.py
"""

import argparse
import os

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType


def get_device_module():
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        return "xpu", torch.xpu
    if torch.cuda.is_available():
        return "cuda", torch.cuda
    raise RuntimeError("no accelerator available")


class Toy(nn.Module):
    def __init__(self, hidden=4096, layers=8):
        super().__init__()
        self.net = nn.Sequential(*[nn.Linear(hidden, hidden) for _ in range(layers)])

    def forward(self, x):
        return self.net(x)


def gather_once(model, optimizer, offload_to_cpu, device_mod, rank):
    device_mod.reset_peak_memory_stats()
    cfg = ShardedStateDictConfig(offload_to_cpu=offload_to_cpu)
    optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=offload_to_cpu)
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT, cfg, optim_cfg):
        model_sd = model.state_dict()
        optim_sd = FSDP.optim_state_dict(model, optimizer)
    dist.barrier()
    peak = device_mod.max_memory_allocated()

    sample_key = next(iter(model_sd))
    sample_tensor = model_sd[sample_key]
    sample_device = getattr(sample_tensor, "device", None)
    if sample_device is None and hasattr(sample_tensor, "local_tensor"):
        sample_device = sample_tensor.local_tensor().device

    if rank == 0:
        print(
            f"[offload_to_cpu={offload_to_cpu}] sample tensor '{sample_key}' device={sample_device} "
            f"peak_device_mem_during_gather={peak / 1024**2:.1f} MiB "
            f"optim_sd_keys={len(optim_sd.get('state', {}))}"
        )
    return sample_device, peak


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hidden", type=int, default=4096)
    parser.add_argument("--layers", type=int, default=8)
    args = parser.parse_args()

    device_name, device_mod = get_device_module()
    local_rank = int(os.environ["LOCAL_RANK"])
    device_mod.set_device(local_rank)
    backend = "xccl" if device_name == "xpu" else "nccl"
    dist.init_process_group(backend=backend)
    rank = dist.get_rank()
    device = torch.device(f"{device_name}:{local_rank}")

    model = Toy(args.hidden, args.layers).to(device)
    model = FSDP(model, device_id=device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    x = torch.randn(8, args.hidden, device=device)
    loss = model(x).sum()
    loss.backward()
    optimizer.step()
    dist.barrier()

    if rank == 0:
        free, total = None, None
        try:
            free, total = device_mod.mem_get_info()
        except Exception:
            pass
        print(f"device={device_name} world_size={dist.get_world_size()} free/total={free}/{total}")

    dev_false, peak_false = gather_once(model, optimizer, False, device_mod, rank)
    dev_true, peak_true = gather_once(model, optimizer, True, device_mod, rank)

    if rank == 0:
        print(f"RESULT offload_False_device={dev_false} offload_True_device={dev_true}")
        print(f"RESULT peak_delta_MiB={(peak_false - peak_true) / 1024**2:.1f}")
        assert str(dev_true).startswith("cpu"), "offload_to_cpu=True did not land the shard on cpu"
        assert not str(dev_false).startswith("cpu"), "offload_to_cpu=False unexpectedly landed on cpu"
        print("PASS: offload_to_cpu mechanism confirmed functional on this device")

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

API and Usage Example

No API change. No config change. Behaviour change is limited to non-CUDA accelerators, which now
get the same CPU-offload default CUDA already had.

Design & Code Changes

verl/utils/checkpoint/fsdp_checkpoint_manager.py:

• Import is_device_available instead of is_cuda_available.
• load_checkpoint: ShardedStateDictConfig / ShardedOptimStateDictConfig now take
  offload_to_cpu=is_device_available().
• save_checkpoint: same for both configs.
• Drops the redundant True if … else False wrapper; ruff format then reflows the two
  surrounding ternaries onto single lines.

Checklist Before Submitting

• Read the Contribute Guide
• ruff / ruff-format clean at the pinned v0.12.2
• Docs — no user-facing surface changes, nothing to update
• Tests — see the Test section for what was measured on 2× B60 and what was not run
• ci-request Slack message
• Not related to the recipe submodule

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: ce61f5d5-609d-4ca4-8067-78c9b81b1172

📥 Commits

Reviewing files that changed from the base of the PR and between 2bc6ac4 and 235a6d7.

📒 Files selected for processing: 1

• verl/utils/checkpoint/fsdp_checkpoint_manager.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Checkpoint saving and loading now determine whether to offload data to the CPU based on device availability, instead of treating CPU offloading as always enabled. This keeps checkpoint handling aligned with the devices available at runtime and avoids applying an unsuitable offloading setting.

FSDP checkpoint loading and saving now set offload_to_cpu using is_device_available() instead of testing the is_cuda_available function object.

Changes: FSDP checkpoint offloading

Layer / File(s)	Summary
Set offloading during checkpoint operations verl/utils/checkpoint/fsdp_checkpoint_manager.py	The checkpoint manager imports is_device_available() and uses its result to set offload_to_cpu during model and optimizer checkpoint loading and saving.

Suggested reviewers: wuxibin89

Priority: ⬇️ Low

Merge Risk: ⚪ Minimal · up to 235a6

Checkpoint CPU offloading now follows device availability across save and load, including non-CUDA platforms. The supplied evidence identifies no concrete regression or merge blocker.

Security architecture risk: 🔵 Low · up to 235a6

This changes where checkpoint tensors are staged, not who can access them. No new access or privilege path was identified. Compatibility with older non-CUDA checkpoints and recovery under host-memory pressure remain unverified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The demonstrated change affects checkpoint operations for the training job's model and optimizer shards on available non-CUDA backends. It moves staging pressure toward host memory; the inspected change establishes no new tenant, service, credential, or network authority.

Trust Boundaries and Controls

• observed — Checkpoint loading continues to derive rank-specific paths from the caller's directory, copy files locally, and deserialize them with torch.load(weights_only=False). These calls and their exposure are unchanged in the inspected comparison. The PR therefore does not establish a new checkpoint-input attack path; it also does not remove the pre-existing need to trust checkpoint contents.

Pre-merge checks: 5

Passed checks: 5 passed

Check name	Status	Explanation
Title check		The title clearly and concisely identifies the main change: making checkpoint CPU offload device-agnostic across hardware backends.
Description check		The description directly explains the FSDP checkpoint offload change, its hardware rationale, measured XPU results, trade-offs, and test coverage.
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}