Support saving/loading checkpoints to fsspec URIs (e.g. `gs://`, `s3://`)

[zhixiangli]

Summary

TorchTitan cannot save or load checkpoints to a cloud object store, even though PyTorch DCP already supports fsspec URIs natively. The only thing standing in the way is a handful of os.path / shutil calls in torchtitan/components/checkpoint.py that assume a local POSIX filesystem. As a result, passing a perfectly valid gs://... path either raises a spurious ValueError or silently mangles the path into a bogus local directory.

I'd like to propose making checkpoint IO fsspec-aware (behaviorally identical for local paths), and I have a PR ready to send if the direction is agreeable.

Motivation

Training on cloud GPUs (GKE/GCP, EKS/AWS, etc.) the compute nodes are usually ephemeral and have small local disks. The durable, canonical place to keep checkpoints is object storage -- gs://bucket/run/checkpoints on GCP, s3://... on AWS. Today a user who points TorchTitan at their bucket cannot:

• write periodic checkpoints straight to GCS for fault tolerance, or
• resume a run (initial_load_path) from a checkpoint that lives in GCS.

The key point: PyTorch already does this

dcp.save / dcp.load auto-select an fsspec-backed writer/reader whenever the checkpoint_id contains a URI scheme. From https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/_storage_utils.py

So the actual bytes-on-the-wire is a first-class, battle-tested torch feature. TorchTitan's dcp_save/dcp_load already pass checkpoint_id straight through to DCP -- the write itself would just work. What breaks is the bookkeeping around the save/load: existence checks, directory listing for step discovery, path joining, and stale-checkpoint purging.

Reproduction

Part A -- torch supports gs:// out of the box (sanity check)

import torch, torch.distributed.checkpoint as dcp
# pip install gcsfs; gcloud auth application-default login
dcp.save({"x": torch.arange(8)}, checkpoint_id="gs://my-bucket/dcp-smoke")
sd = {"x": torch.zeros(8, dtype=torch.long)}
dcp.load(sd, checkpoint_id="gs://my-bucket/dcp-smoke")
print(sd["x"])   # tensor([0,1,2,3,4,5,6,7]) -- works

Part B -- TorchTitan blocks the same path (no cloud/creds needed)

This is pure-Python path handling from checkpoint.py on main; you can see it fail in seconds without a bucket or gcsfs installed:

import os
base_folder = "./outputs"                              # --dump_folder
folder_cfg  = "gs://my-bucket/run/checkpoints"         # --checkpoint.folder

# checkpoint.py:442  self.folder = os.path.join(base_folder, config.folder)
print(os.path.join(base_folder, folder_cfg))
# -> './outputs/gs://my-bucket/run/checkpoints'   (URI mangled into a local dir)

# checkpoint.py:795  if not os.path.exists(self.folder): ...
# -> False, so save/resume silently targets the wrong place

# checkpoint.py:392  initial_load_path validation
ilp = "gs://my-bucket/run/checkpoints/step-1000"
assert ilp.startswith("/"), \
    f"ValueError: initial_load_path must be absolute: {ilp}"   # raises

Output on main:

./outputs/gs://my-bucket/run/checkpoints
ValueError: initial_load_path must be absolute: gs://my-bucket/run/checkpoints/step-1000

Part C -- the real end-to-end use case

# Save periodic checkpoints straight to GCS
NGPU=4 CONFIG=llama3_debugmodel ./run_train.sh \
  --training.steps 20 \
  --checkpoint.enable \
  --checkpoint.interval 10 \
  --checkpoint.folder gs://my-bucket/torchtitan-smoke/checkpoints

# ...and later resume from GCS
NGPU=4 CONFIG=llama3_debugmodel ./run_train.sh \
  --training.steps 40 \
  --checkpoint.enable \
  --checkpoint.folder gs://my-bucket/torchtitan-smoke/checkpoints

On main this fails during checkpoint setup / load-step discovery (os.path.exists, os.listdir, os.path.isdir over the mangled path). With the proposed change it saves to GCS every 10 steps and resumes from step-10/step-20 on the next launch.

Proposed change (small, local path unchanged)

Route the filesystem bookkeeping in checkpoint.py through a thin helper that dispatches on whether the path is a remote fsspec URI ("://" in path -- the same heuristic DCP uses in FileSystem.validate_checkpoint_id):

• local paths keep using os/shutil verbatim -- behaviorally identical, so the converged/battle-tested local code path does not change;
• remote paths use fsspec (exists, isdir, isfile, ls, rm), with the backend driver (gcsfs, s3fs, ...) imported lazily so pure-local installs never need fsspec.

Concretely the touched call sites are: os.path.join(base_folder, folder) (treat a remote folder as absolute, don't prefix dump_folder), the initial_load_path "must be absolute" check (accept remote URIs), os.path.exists/isdir/isfile existence checks, os.listdir step discovery, and shutil.rmtree purging.

Auth is delegated entirely to the fsspec backend (ADC / GOOGLE_APPLICATION_CREDENTIALS for GCS), so TorchTitan takes on no credential handling.

Scope / non-goals

• Native DCP format only. HF-safetensors save/load (last_save_in_hf / initial_load_in_hf) to a remote URI is out of scope for the first cut.
• No new hard dependency: fsspec/gcsfs are only needed if you actually use a remote URI; local heckpointing is untouched.
• Other outputs (TensorBoard, profiler traces) keep writing locally; only the checkpoint folder honors a remote URI.

[tianyu-l]

Is this related to #3730 which is being reviewed?

cc @fegin

[zhixiangli]

Is this related to #3730 which is being reviewed?

cc @fegin

Ah, yes! I wasn't aware of the previous PR before submitting mine. I would appreciate it if you could take a quick look at #3887 to see if any parts of this implementation might be useful. Thanks!