[dist.checkpoint] Cleanup usage of collectives and introduce narrow helper (#81828)
Introduce _DistWrapper class that wraps a process group and provides functional
variants of collectives. It works without c10d enabled and is exception
robust.
Introduce tensor_narrow_n that handle narrowing over multiple dimentions.
Fixes #ISSUE_NUMBER
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81828
Approved by: https://github.com/wanchaol
