FSDP hybrid shard should checkpoint in a single node

[carmocca]

Description & Motivation

pytorch/pytorch#104810 adds the recommendation that the save APIs should be called in a single node (shard_group).

pytorch/pytorch#102904 (comment) Also talks about this

Our logic doesn't do this and runs this code in all ranks.

Additional context

Lit-gpt uses hybrid sharding in pretrain/tinyllama.py but full checkpointing. I believe this feature request is only relevant for sharded pointing. @awaelchli Did you try it? Does sharded hybrid checkpointing work?

cc @Borda @awaelchli @carmocca

[awaelchli]

Yes we should fix this. I have never had the chance to test this I think. While yes we can temporarily fix this simply calling the saver only in node 0, in general hybrid shard can span an arbitrary number of processes. We would need to get the right process group from the device mesh, pass it to the saver too and only call the saver from these ranks.

[qsh-zh]

The feature is so useful. Is it supported now?

[zaptrem]

Also interested in this along with async checkpointing.

[llllvvuu]

I checked the behavior here, it seems like torch.distributed.checkpoint.save already dedupes, so there already is no waste in file size.

But if you want fewer files, you can do save on if trainer.model._inter_node_pg.rank() == 0 with torch.distributed.checkpoint.save(process_group=trainer.model.process_group)