method handles the fact that gradients are sharded across
        ranks.
  Note:

    The total norm returned will have the "largest" dtype across
    all parameters/gradients as defined by PyTorch's type
    promotion semantics. For example, if *all*
    parameters/gradients use a low precision dtype, then the
    returned norm's dtype will be that low precision dtype, but if
    there exists at least one parameter/ gradient using FP32, then
    the returned norm's dtype will be FP32.

  Warning:

    This needs to be called on all ranks since it uses collective
    communications.

static flatten_sharded_optim_state_dict(sharded_optim_state_dict, model, optim)
  The API is similar to "shard_full_optim_state_dict()". The only
  difference is that the input "sharded_optim_state_dict" should
  be returned from "sharded_optim_state_dict()". Therefore, there
  will be all-gather calls on each rank to gather "ShardedTensor"
  s.
