Pre-download only relevant files per rank.

[2015aroras]

This builds on #123 and tries to change it in ways listed below. Completely untested and just an idea.

• Instead of pre-downloading the whole checkpoint, download a checkpoint file when a range query is about to made on it for the first time. resource_path will (hopefully) prevent any concurrency issues. As a bonus, this also splits up the work between the various processes.
• Instead of saving each non-sharded data item on the rank with the smallest save payload (thus spreading the items out across a bunch of files/ranks), save them all to the lowest rank (empirically, rank 0). This will limit how many files we have to download/read from, but makes rank 0 bigger (I think only by like 1 MB though).
  • This is the solution for [Distributed Checkpoint] When loading FSDP sharded checkpointing each rank needs all the checkpointing files  pytorch/pytorch#125740

[epwalsh]

This is good, thanks for implemented it

[epwalsh]

@2015aroras ready to merge into the other branch?

[2015aroras]

c99da44 I need to cache the GCS client in cached_path, otherwise we get an error we saw in the old trainer until we did a similar fix.