[FSDP] Option to summon on rank 0 only (#73903)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73903
Add an option to summon full parameter on rank0_only. Note that as documented, when using this flag, writes are only guaranteed to persist for rank 0.
ghstack-source-id: 151311233
Test Plan: CI
Reviewed By: zhaojuanmao
Differential Revision: D34706449
fbshipit-source-id: 034d76cdae6916a225bd6fcee092db48359f8a0b
(cherry picked from commit 9cef01d2c5cbeda4ea02d22d7dd6ef17fb4e7250)