Dcp.save straight to cloud storage

Sam_Shleifer: <p>Is there an example somewhere of doing dcp.save straight to cloud storage, like s3 or google cloud storage?</p>

yf225: <p>Curious would these help?</p>
<ul>
<li>AWS S3: <a href="https://aws.amazon.com/about-aws/whats-new/2024/11/amazon-s3-connector-pytorch-distributed-checkpoint/" class="inline-onebox" rel="noopener nofollow ugc">Amazon S3 Connector for PyTorch now supports Distributed Checkpoint - AWS</a></li>
<li>Google Cloud Storage: <a href="https://cloud.google.com/storage/docs/pytorch-connector" class="inline-onebox" rel="noopener nofollow ugc">Connector for PyTorch &nbsp;|&nbsp; Cloud Storage &nbsp;|&nbsp; Google Cloud</a></li>
</ul>

tianyu: <p>Not sure but this PR may help you</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/torchtitan/pull/789">
  <header class="source">

      <a href="https://github.com/pytorch/torchtitan/pull/789" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/torchtitan</a>
  </header>

  <article class="onebox-body">
    <div class="github-row" data-github-private-repo="false">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/torchtitan/pull/789" target="_blank" rel="noopener nofollow ugc">Make CheckpointManager friendlier to custom StorageWriter/StorageReader</a>
      </h4>

    <div class="branches">
      <code>main</code> ← <code>dimdi-y:dimdi-y/extendable-checkpointing</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2025-01-12" data-time="23:05:33" data-timezone="UTC">11:05PM - 12 Jan 25 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/dimdi-y" target="_blank" rel="noopener nofollow ugc">
            <img alt="" src="https://avatars.githubusercontent.com/u/191106290?v=4" class="onebox-avatar-inline" width="20" height="20">
            dimdi-y
          </a>
        </div>

        <div class="lines" title="4 commits changed 2 files with 332 additions and 28 deletions">
          <a href="https://github.com/pytorch/torchtitan/pull/789/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+332</span>
            <span class="removed">-28</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">The goal of this PR is to make using custom `torch.distributed.checkpoint.storag<span class="show-more-container"><a href="https://github.com/pytorch/torchtitan/pull/789" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">e.StorageReader/StorageWriter`s with torchtitan easier.
For example, using `S3StorageWriter` from [s3-connector-for-pytorch](https://github.com/awslabs/s3-connector-for-pytorch/#distributed-checkpoints) would allow storing checkpoints in AWS S3.

Two changes needed to be made to allow this:
1. Allow passing `storage_reader` and `storage_writer` into the `CheckpointManager`.
2. Isolate the usage of `os.path` module under overridable methods (as this logic is specific to the default `torch.distributed.checkpoint.filesystem.FileSystemReader/FileSystemWriter`).

I also added basic tests for CheckpointManager to make sure I didn't break the original behavior</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Sam_Shleifer: <p>That looks very promising, thanks for the link!</p>

Sam_Shleifer: <p>We tried this and it seems the writer is quite fast and the reader is quite slow, so then we switched back to a method where we download files from gs:// to node local /tmp storage before loading.</p>
<p>More specifically, each <code>global_rank</code> (according to device_mesh), downloads</p>
<pre><code class="lang-auto">f"__{global_rank}_0.distcp" 
</code></pre>
<p>from gs to /tmp, and every node downloads a copy of <code>.metadata</code>.</p>
<p>In this process we realized that for some reason, even if the <code>DeviceMesh</code> of the saving job is identical to that of the loading job, we get an error from global rank 7 like</p>
<pre><code class="lang-auto">file "__8_0.distcp" not found
</code></pre>
<p>(recall files 0-7 were download to this node’s /tmp)</p>
<p>If I understand correctly, rank 7 shouldn’t need this file in the identical mesh case.<br>
It should only need <code>__7_0.distcp</code>.</p>
<p>Do I understand correctly?</p>
<p>Checks:</p>
<ul>
<li>Verified that the metadata saved by <code>GCSDistributedWriter</code> and the baseline are identical!</li>
<li>Not yet clear to me whether this is a <code>GCSDistributedWriter</code> issue or a general dcp issue.</li>
</ul>

fegin: <p>It’s not about DeviceMesh but how the model is sharded/parallelized. If the model is parallelized the same way as saving, then all ranks should only need to access these files when loading, 1) metadata file, 2) the files saved by the same rank.</p>