[jackychris]

What does this PR do?

The three torch.load calls in FSDPCheckpointManager.load_checkpoint pass no
map_location, so every tensor is revived on the device it was saved from. On a
colocated run where the rollout engine already holds most of the card, resuming dies
before the first step:

torch.OutOfMemoryError: Tried to allocate 4.64 GiB.
GPU 0 has a total capacity of 79.33 GiB of which 1.11 GiB is free.

Loading to CPU is consistent with what the surrounding code already declares: the
enclosing get_fsdp_state_ctx is entered with
ShardedStateDictConfig(offload_to_cpu=True) whenever CUDA is available, so this
manager treats its sharded state dicts as CPU-resident on the save side. The load side
is the only place that does not say so.

What this changes is where the tensors land on their way in: staging them through host
memory bounds the peak by the shard size rather than by whatever the card happens to
have free. See the note at the end of the Test section for the part of that I have
observed and the part I have only reasoned about.

Checklist Before Starting

• Search for similar PRs:
  https://github.com/verl-project/verl/issues?q=repo%3Averl-project%2Fverl+is%3Aopen+map_location+checkpoint

  One hit, Critical Security Vulnerabilities in Mathematical Evaluation and Serialization Code #5331, which mentions weights_only and map_location in general terms but
  does not concern fsdp_checkpoint_manager.py or the resume path, so it is not a
  duplicate of this. Worth noting for a different reason: a future weights_only=True
  hardening pass would edit these same three calls, and whoever does that should not be
  surprised to find them already carrying a map_location.

  Also open on this file: [ckpt] fix: normalize PEFT keys in HF save and dump LoRA adapter with live config #4770 ([ckpt] fix: normalize PEFT keys in HF save and dump LoRA adapter with live config, last touched 2026-07-13) edits
  fsdp_checkpoint_manager.py but not these three torch.load calls, so the two do not
  overlap in content.

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

Test

Found in the field, on a single A800-80G with a colocated SGLang engine, resuming a run
whose model_world_size_1_rank_0.pt was 34 GiB. The traceback above is from that run.

CI cannot reach it: reproducing the OOM needs a resume from a sharded checkpoint onto a
device that is already near-full, and no current workflow sets that up. The underlying
behaviour is not in question, though — torch.load restores tensors to their saved
device unless told otherwise, which is documented and is exactly why map_location
exists.

A standalone script (below) isolates it on a single card: save a shard from CUDA,
allocate ballast until the card is as full as it is mid-run, then load it both ways. On
the A800 that hit the original failure:

device: NVIDIA A800 80GB PCIe
memory: 79.25 GiB total, 78.84 GiB free
memory is hard-limited; the OOM in step 3 is reachable.

1. saving a 2.00 GiB shard from CUDA ...
2. allocating 75.83 GiB of ballast (a resident rollout engine) ...
   1.00 GiB free, shard needs 2.00 GiB
3. torch.load(path, weights_only=False)          # what main does
   -> torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB.
      GPU 0 has a total capacity of 79.25 GiB of which 1023.94 MiB is free.
4. torch.load(path, weights_only=False, map_location='cpu')   # the patch
   -> loaded onto cpu, GPU untouched

Same shape as the traceback above, at a size chosen to be quick rather than to match the
34 GiB shard. The placement is the whole mechanism: with no map_location a checkpoint
saved from CUDA is revived in device memory, and once an engine is resident there is
nowhere to put it.

Also run on a consumer RTX 5070 under WSL, where step 3 does not raise — WDDM pages
device allocations out to host RAM instead of refusing them, and a single 20 GiB
allocation succeeds on an 11.9 GiB card. The placement difference is identical there;
only the OOM needs a card that hard-limits. The script probes for this and reports which
case it is in, so it does not read as a pass when the platform simply papered over it.

Lint: ruff==0.12.2 check and ruff format --check clean.

Not covered by the above: a resume that completes with the patch applied. The run
that hit this died at the load and was not restarted, so what is established is the
failure and the mechanism — that torch.load without map_location claims device
memory, and that naming CPU does not — but not that load_state_dict then proceeds
normally on CPU-resident tensors. The surrounding code expects them to be CPU-resident
(ShardedStateDictConfig(offload_to_cpu=True) on the save side), which is why I think
this is the right change rather than a workaround, but expecting is not the same as
observing and I would rather say so than imply otherwise.

tests/special_distributed/test_fsdp_ckpt.py covers save/load round trips and needs
multiple GPUs; it has not been run here and is the natural place to settle this if you
would like it settled before merging.

Repro script (single CUDA device, no verl import needed)

#!/usr/bin/env python3
"""Minimal repro: `torch.load` without `map_location` revives checkpoint
tensors in device memory, so resuming OOMs when a rollout engine already holds the card.

Needs one CUDA device. Sizes itself to whatever card it finds, so it runs on a 12 GiB
consumer GPU as happily as on an 80 GiB A800:

    python pr02_map_location_gpu.py

It reproduces the failure rather than merely describing it:

1. save a checkpoint shard from CUDA, as `FSDPCheckpointManager.save_checkpoint` does;
2. free it, then allocate ballast so the card is as full as it is mid-run with a
   colocated inference engine resident;
3. `torch.load(path)` -- what `load_checkpoint` does today -- and watch it OOM;
4. `torch.load(path, map_location="cpu")` -- the patch -- and watch it succeed.
"""

import sys
import tempfile
from pathlib import Path

import torch

GIB = 1024**3


def _fmt(nbytes: float) -> str:
    return f"{nbytes / GIB:.2f} GiB"


def oversubscribes(total: int) -> bool:
    """Whether the platform lets a single allocation exceed physical VRAM.

    WSL2 drives the GPU through WDDM, which pages device memory out to host RAM, so
    `cudaMalloc` succeeds well past the card's capacity and no amount of ballast will
    force an OOM. Native Linux with a datacenter card does not do this.
    """
    try:
        t = torch.empty(int(total * 1.5) // 2, dtype=torch.float16, device="cuda")
    except torch.OutOfMemoryError:
        return False
    del t
    torch.cuda.empty_cache()
    return True


def main() -> int:
    if not torch.cuda.is_available():
        print("no CUDA device -- this repro is about device placement and needs one.")
        return 2

    torch.cuda.init()
    name = torch.cuda.get_device_name(0)
    free_at_start, total = torch.cuda.mem_get_info()
    print(f"device: {name}")
    print(f"memory: {_fmt(total)} total, {_fmt(free_at_start)} free")

    paged = oversubscribes(total)
    if paged:
        print("note:   this platform oversubscribes GPU memory (allocations exceed VRAM),")
        print("        so step 3 will not OOM here however full the card is. Placement is")
        print("        still measured; the OOM needs a card that refuses to page out.\n")
    else:
        print("memory is hard-limited; the OOM in step 3 is reachable.\n")

    # A shard big enough to matter but comfortably saveable on any card. The real one
    # that hit this was 34 GiB; the mechanism does not depend on the size.
    shard_bytes = min(2 * GIB, int(free_at_start * 0.25))
    n_elems = shard_bytes // 4  # fp32
    side = int(n_elems**0.5)
    shard_bytes = side * side * 4

    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "model_world_size_1_rank_0.pt"

        print(f"1. saving a {_fmt(shard_bytes)} shard from CUDA ...")
        state = {"weight": torch.randn(side, side, device="cuda")}
        torch.save(state, path)
        del state
        torch.cuda.empty_cache()
        print(f"   wrote {_fmt(path.stat().st_size)} to disk")

        # Occupy the card the way a resident SGLang engine does, leaving less headroom
        # than the shard needs.
        free_now, _ = torch.cuda.mem_get_info()
        headroom = int(shard_bytes * 0.5)
        ballast_bytes = max(0, free_now - headroom)
        print(f"\n2. allocating {_fmt(ballast_bytes)} of ballast (a resident rollout engine) ...")
        ballast = torch.empty(ballast_bytes // 2, dtype=torch.float16, device="cuda")
        free_after_ballast, _ = torch.cuda.mem_get_info()
        print(f"   {_fmt(free_after_ballast)} free, shard needs {_fmt(shard_bytes)}")

        print("\n3. torch.load(path, weights_only=False)          # what main does")
        try:
            loaded = torch.load(path, weights_only=False)
        except torch.OutOfMemoryError as e:
            first_line = str(e).strip().splitlines()[0]
            print(f"   -> torch.OutOfMemoryError: {first_line}")
            oomed = True
        else:
            why = "allocation paged out to host RAM" if paged else "card had room after all"
            print(f"   -> loaded onto {loaded['weight'].device}, no OOM ({why})")
            oomed = False
            del loaded
            torch.cuda.empty_cache()

        print("\n4. torch.load(path, weights_only=False, map_location='cpu')   # the patch")
        loaded = torch.load(path, weights_only=False, map_location="cpu")
        print(f"   -> loaded onto {loaded['weight'].device}, GPU untouched")
        assert loaded["weight"].device.type == "cpu"
        assert loaded["weight"].shape == (side, side)

        del ballast, loaded
        torch.cuda.empty_cache()

    print()
    print("Placement confirmed: the default reload claims device memory before FSDP has")
    print("a say; naming CPU keeps the card untouched.")
    if oomed:
        print("OOM reproduced: on this card that placement is fatal once the engine is resident.")
        return 0
    if paged:
        print("OOM not reachable here: the platform paged the allocation out to host RAM")
        print("rather than refusing it. Run on a card that hard-limits to see it fail.")
        return 0
    print("OOM not reproduced, and memory is hard-limited -- the card had room after all.")
    print("Re-run with a shard sized closer to the card.")
    return 1


if __name__ == "__main__":
    sys.exit(main())

API and Usage Example

No API change.

Design & Code Changes

verl/utils/checkpoint/fsdp_checkpoint_manager.py — add map_location="cpu" to the
model, optimizer and extra-state torch.load calls in load_checkpoint.

Checklist Before Submitting

• Read the Contribute Guide
• ruff and ruff-format clean (pinned to the repo's v0.12.2)
• Docs — no user-facing change, so nothing to update
• Tests — explained above why CI cannot cover this
• ci-request Slack message
• Not related to the recipe submodule, so no submodule reference to update

AI assistance disclosure

AI assistance (Claude) was used to write the patch and the repro and to draft this
description.

---

Note on scope

AGENTS.md warns against one-off PRs for tiny edits. This is three words in three
places, so it is worth saying why it is not busywork: it is a hard crash on the
resume path, not a style preference, and the fix has exactly one reasonable spelling.
If you would rather see it folded into a broader checkpoint-hygiene change (e.g.
alongside weights_only), say so and I will resubmit it that way.

[cla-assistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

fuhao seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}