fix(megatron): stage common state for cloud resume

skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py

@@ -25,6 +25,7 @@
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from torch import distributed as dist
 from torch import optim
+from torch.distributed.checkpoint import FileSystemReader
 from transformers import PreTrainedTokenizer
 
 from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import (
@@ -520,12 +521,12 @@ def load_checkpoint(
     def _load_dist_checkpoint_from_cloud(self, ckpt_dir: str, sharded_state_dict: dict) -> dict:
         """Download checkpoint shards from cloud storage with per-rank parallelism.
 
-        All ranks on the same node share a single local directory. Each rank
-        downloads only its own shard file(s) into the shared dir, so the total
-        download per node equals one full copy of the checkpoint (instead of one
-        copy per rank).
+        All ranks on the same node share a single node-local directory.
+        Each node downloads its ranks' assigned shards plus files backing the
+        common state, rather than downloading the full checkpoint on every rank.
 
-        Local rank 0 creates the directory and downloads common metadata files.
+        Local rank 0 creates the directory and downloads common metadata files
+        and any shard containing the common state read independently by every rank.
         After a barrier, all shard files are present and every rank can load.
 
         Does not currently support flexible trainer resharding.
@@ -556,6 +557,22 @@ def _load_dist_checkpoint_from_cloud(self, ckpt_dir: str, sharded_state_dict: di
                         continue
                     io.download_file(cloud_entry, os.path.join(local_dir, name))
 
+                if not os.path.isfile(os.path.join(local_dir, "common.pt")):
+                    # Current Megatron checkpoints embed common state in a DCP
+                    # shard. Every rank reads it before distributed shard loading.
+                    common_key = dist_checkpointing.ShardedObject("common_state", None, (1,), (0,)).unique_key
+                    metadata = FileSystemReader(local_dir).read_metadata()
+                    common_shards = {
+                        storage.relative_path
+                        for index, storage in metadata.storage_data.items()
+                        if index.fqn == common_key
+                    }
+                    for relative_path in sorted(common_shards):
+                        local_path = os.path.join(local_dir, relative_path)
+                        if not os.path.isfile(local_path):
+                            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+                            io.download_file(f"{ckpt_dir.rstrip('/')}/{relative_path}", local_path)
+
             # Wait for the directory and common files to be ready.
             dist.barrier()
 
@@ -568,7 +585,8 @@ def _load_dist_checkpoint_from_cloud(self, ckpt_dir: str, sharded_state_dict: di
                 if match and int(match.group(1)) == global_rank:
                     cloud_path = ckpt_dir.rstrip("/") + "/" + name
                     local_path = os.path.join(local_dir, name)
-                    io.download_file(cloud_path, local_path)
+                    if not os.path.isfile(local_path):
+                        io.download_file(cloud_path, local_path)
 
             # Wait for all ranks to finish downloading their shards.
             dist.barrier()

tests/backends/skyrl_train/distributed/test_megatron_cloud_checkpoint.py

@@ -0,0 +1,201 @@
+"""Cloud checkpoint staging with real PyTorch DCP metadata and common-state reads."""
+
+import importlib.util
+import io
+import pickle
+import shutil
+import sys
+import threading
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.distributed.checkpoint as dcp
+
+COMMON_KEY = "common_state/shard_0_1"
+CLOUD_PATH = "s3://test-bucket/global_step_10/policy"
+COMMON_STATE = {"lr_scheduler": {"step": 10}, "rng": {"seed": 42}}
+
+
+@pytest.fixture
+def strategy_module(monkeypatch):
+    """Import the actual strategy while replacing GPU-only Megatron dependencies."""
+
+    def stub(name, **attributes):
+        module = ModuleType(name)
+        module.__path__ = []
+        module.__dict__.update(attributes)
+        monkeypatch.setitem(sys.modules, name, module)
+        if "." in name:
+            parent, leaf = name.rsplit(".", 1)
+            if parent in sys.modules and parent.startswith("megatron"):
+                monkeypatch.setattr(sys.modules[parent], leaf, module, raising=False)
+        return module
+
+    class DistributedOptimizer:
+        load_parameter_state_from_dp_reshardable = Mock()
+
+    stub("megatron")
+    stub("megatron.core")
+    stub("megatron.core.parallel_state", get_data_parallel_group=lambda **kwargs: None)
+    stub(
+        "megatron.core.dist_checkpointing",
+        load=Mock(),
+        ShardedObject=Mock(return_value=SimpleNamespace(unique_key=COMMON_KEY)),
+    )
+    stub(
+        "megatron.core.dist_checkpointing.serialization",
+        get_default_load_sharded_strategy=Mock(),
+        get_default_save_sharded_strategy=Mock(),
+    )
+    stub("megatron.core.dist_checkpointing.strategies")
+    stub("megatron.core.dist_checkpointing.strategies.async_utils", AsyncCallsQueue=Mock())
+    stub(
+        "megatron.core.dist_checkpointing.strategies.fully_parallel",
+        FullyParallelLoadStrategyWrapper=Mock(),
+        FullyParallelSaveStrategyWrapper=Mock(),
+    )
+    stub("megatron.core.optimizer", DistributedOptimizer=DistributedOptimizer)
+    stub("megatron.core.optimizer_param_scheduler", OptimizerParamScheduler=object)
+    stub(
+        "skyrl.backends.skyrl_train.distributed.megatron.megatron_utils",
+        **{
+            name: Mock()
+            for name in (
+                "load_megatron_grads_to_gpu",
+                "load_megatron_model_to_gpu",
+                "load_megatron_optimizer",
+                "offload_megatron_grads_to_cpu",
+                "offload_megatron_model_to_cpu",
+                "offload_megatron_optimizer",
+            )
+        },
+    )
+    stub("skyrl.backends.skyrl_train.workers.megatron.megatron_model_wrapper", MegatronModelWrapper=object)
+    path = Path(__file__).parents[4] / "skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py"
+    spec = importlib.util.spec_from_file_location("_cloud_checkpoint_test_strategy", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def make_cloud_checkpoint(directory, common_shard, *, legacy=False):
+    directory.mkdir()
+    common = io.BytesIO()
+    torch.save([COMMON_STATE], common)
+    # Megatron serializes each ShardedObject as a BytesIO containing a list.
+    dcp.save({COMMON_KEY: common}, checkpoint_id=directory)
+    metadata = dcp.FileSystemReader(directory).read_metadata()
+    original_shard = next(iter(metadata.storage_data.values())).relative_path
+    destination = directory / common_shard
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    (directory / original_shard).rename(destination)
+    metadata.storage_data = {
+        index: replace(storage, relative_path=common_shard) for index, storage in metadata.storage_data.items()
+    }
+    with (directory / ".metadata").open("wb") as file:
+        pickle.dump(metadata, file)
+    (directory / "metadata.json").write_text('{"sharded_backend":"torch_dist","sharded_backend_version":1}')
+    for rank in (1, 2, 9):
+        (directory / f"__{rank}_0.distcp").write_bytes(b"model shard not needed by the common-state read")
+    if legacy:
+        torch.save(COMMON_STATE, directory / "common.pt")
+        destination.unlink()
+
+
+def setup_cloud_transport(monkeypatch, module, root, cloud):
+    downloads = []
+    staged_paths = []
+    local = root / "node-local"
+    local.mkdir()
+    monkeypatch.setattr(module.tempfile, "gettempdir", lambda: str(local))
+    monkeypatch.setattr(module.io, "list_dir", lambda _: [f"{CLOUD_PATH}/{path.name}" for path in cloud.iterdir()])
+
+    def source_path(path):
+        return cloud / path.removeprefix(CLOUD_PATH + "/")
+
+    monkeypatch.setattr(module.io, "isdir", lambda path: source_path(path).is_dir())
+
+    def download(source, destination):
+        downloads.append(source.removeprefix(CLOUD_PATH + "/"))
+        shutil.copyfile(source_path(source), destination)
+
+    monkeypatch.setattr(module.io, "download_file", download)
+
+    def read_common_state(*, checkpoint_dir, **kwargs):
+        staged_paths.append(Path(checkpoint_dir))
+        legacy = Path(checkpoint_dir) / "common.pt"
+        if legacy.exists():
+            return torch.load(legacy, weights_only=True)
+        state = {COMMON_KEY: io.BytesIO()}
+        # Megatron performs this independent local read before its parallel load.
+        dcp.load(state, storage_reader=dcp.FileSystemReader(checkpoint_dir), no_dist=True)
+        state[COMMON_KEY].seek(0)
+        return torch.load(state[COMMON_KEY], weights_only=True)[0]
+
+    module.dist_checkpointing.load.side_effect = read_common_state
+    return downloads, staged_paths
+
+
+@pytest.mark.parametrize("global_rank", [0, 1])
+@pytest.mark.parametrize("common_shard", ["__0_0.distcp", "__7_3.distcp", "objects/__5_2.distcp"])
+def test_cloud_load_stages_common_state_on_every_node(
+    strategy_module, tmp_path, monkeypatch, global_rank, common_shard
+):
+    cloud = tmp_path / "cloud"
+    make_cloud_checkpoint(cloud, common_shard)
+    downloads, staged_paths = setup_cloud_transport(monkeypatch, strategy_module, tmp_path, cloud)
+    monkeypatch.setattr(strategy_module.dist, "get_rank", lambda: global_rank)
+    monkeypatch.setattr(strategy_module.dist, "barrier", lambda: None)
+    strategy = strategy_module.MegatronStrategy(SimpleNamespace(), node_local_rank=0)
+
+    result = strategy._load_dist_checkpoint_from_cloud(CLOUD_PATH, {})
+
+    assert result == COMMON_STATE
+    assert downloads.count(common_shard) == 1
+    if global_rank == 1:
+        assert downloads.count("__1_0.distcp") == 1
+    assert "__2_0.distcp" not in downloads
+    assert "__9_0.distcp" not in downloads
+    assert all(not path.exists() for path in staged_paths)
+
+
+def test_legacy_common_pt_does_not_require_embedded_state_shard(strategy_module, tmp_path, monkeypatch):
+    cloud = tmp_path / "cloud"
+    make_cloud_checkpoint(cloud, "__0_0.distcp", legacy=True)
+    downloads, _ = setup_cloud_transport(monkeypatch, strategy_module, tmp_path, cloud)
+    monkeypatch.setattr(strategy_module.dist, "get_rank", lambda: 1)
+    monkeypatch.setattr(strategy_module.dist, "barrier", lambda: None)
+    strategy = strategy_module.MegatronStrategy(SimpleNamespace(), node_local_rank=0)
+
+    assert strategy._load_dist_checkpoint_from_cloud(CLOUD_PATH, {}) == COMMON_STATE
+
+    assert downloads.count("common.pt") == 1
+    assert "__0_0.distcp" not in downloads
+
+
+def test_node_ranks_share_common_shard_without_duplicate_downloads(strategy_module, tmp_path, monkeypatch):
+    cloud = tmp_path / "cloud"
+    make_cloud_checkpoint(cloud, "__0_0.distcp")
+    downloads, staged_paths = setup_cloud_transport(monkeypatch, strategy_module, tmp_path, cloud)
+    rank_context = threading.local()
+    barrier = threading.Barrier(2)
+    monkeypatch.setattr(strategy_module.dist, "get_rank", lambda: rank_context.rank)
+    monkeypatch.setattr(strategy_module.dist, "barrier", lambda: barrier.wait(timeout=10))
+
+    def load(local_rank):
+        rank_context.rank = local_rank + 1  # A node that does not host global rank 0.
+        strategy = strategy_module.MegatronStrategy(SimpleNamespace(), node_local_rank=local_rank)
+        return strategy._load_dist_checkpoint_from_cloud(CLOUD_PATH, {})
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        results = list(executor.map(load, (0, 1)))
+
+    assert results == [COMMON_STATE, COMMON_STATE]
+    assert Counter(downloads) == Counter([".metadata", "metadata.json", "__0_0.distcp", "__1_0.distcp", "__2_0.distcp"])
+    assert all(not path.exists() for path in staged_paths)