synchronized in the first forward-backward pass after exiting
      the context. This should only be used on the root FSDP instance
      and will recursively apply to all children FSDP instances.
  Note:

    This likely results in higher memory usage because FSDP will
    accumulate the full model gradients (instead of gradient
    shards) until the eventual sync.

  Note:

    When used with CPU offloading, the gradients will not be
    offloaded to CPU when inside the context manager. Instead,
    they will only be offloaded right after the eventual sync.

  Return type:
     *Generator*

register_comm_hook(state, hook)
  Registers a communication hook which is an enhancement that
  provides a flexible hook to users where they can specify how
  FSDP aggregates gradients across multiple workers. This hook can
  be used to implement several algorithms like GossipGrad and
  gradient compression which involve different communication
