use correct PG when collecting metrics with HYBRID shard - Githubissues

allenai / OLMo

Modeling, training, eval, and inference code for OLMo

https://allenai.org/olmo

Apache License 2.0

4.24k stars 400 forks source link

use correct PG when collecting metrics with HYBRID shard #551

Closed epwalsh closed 3 months ago

epwalsh commented 3 months ago

Followup to #540. Fixes how we collect per-param optim metrics when using hybrid sharding. The process group we're using is the same process group that FSDP uses during hybrid sharding when reducing the grad norms, for example, so it should be the right one. See https://github.com/pytorch/pytorch/blob/cb17721899d4d6a55d66d4f7188e36c20a078231/torch/distributed/fsdp/fully_sharded_data_parallel.py#L1149.