fix(megatron): stage common state for cloud resume

[bvolpato]

Summary

Multi-node Megatron resume from S3/GCS stages only each node's own rank shards. The pinned Megatron version stores common state inside a DCP shard and reads it independently on every rank before distributed loading. Nodes without that shard fail with FileNotFoundError, even when the cloud checkpoint is complete and the training topology is unchanged.

Have each node's local rank 0 resolve the common-state payload through DCP metadata and download its backing file(s) before the first barrier. Rank-specific downloads reuse files already staged. Legacy common.pt checkpoints keep their existing behavior. This also handles alternate writer/bucket names without assuming __0_0.distcp.

The format contract is in the pinned Megatron common-state loader.

Testing

Eight CPU regressions use real PyTorch DCP serialization, metadata, and common-state reads with simulated cloud transport and Megatron dependencies. They cover zero/nonzero node ranks, alternate and nested shard paths, legacy common.pt, two concurrent ranks sharing node-local staging, download deduplication, and cleanup.

All eight regressions and three existing checkpoint-loading tests pass. Against the unchanged strategy, six regressions fail while the rank-0 and legacy controls pass. Executed with PyTorch 2.11.0+cpu through isolated uv; independent review also passed the new tests. Required pre-commit hooks pass. Full multi-node GPU restore and live S3/GCS were not run.

The GitHub GPU job exited before running tests: ANYSCALE_CLI_TOKEN was empty and Anyscale rejected the credentials.

Repro command

From the repository root:

uv run --isolated --extra skyrl-train --extra dev pytest --noconftest \
  tests/backends/skyrl_train/distributed/test_megatron_cloud_checkpoint.py -q

Downsides

A node may download and temporarily store additional common-state backing file(s), once per node. Those files can also contain model tensors, so the added disk/network cost is their complete file size, not just the common-state object's size. The loader still downloads only its assigned rank shards and the required common-state files.

Risk and rollback

The change affects cloud checkpoint staging and does not alter saved files. Additional staging can fail if local disk space is insufficient. Revert the commit to roll back; affected checkpoints can still be recovered by staging the complete checkpoint on a shared filesystem and resuming from that local path. Flexible resharding remains outside this loader's supported scope.