docstr: Returns all unsharded gradients as a single flattened tensor. This
  returns the same value on all ranks.
docstr_tok: ["Returns","all","unsharded","gradients","as","a","single","flattened","tensor",".","This","returns","the","same","value","on","all","ranks","."]
code: def _get_grads_as_flattened(
      self,
      model: FSDP,
      uses_tp: bool,
      param_name_to_numel: Dict[str, int],
      param_name_to_sharding_info: Dict[str, Tuple[torch.Size, int]],
      tp_pg: Optional[dist.ProcessGroup],
      fsdp_pg: Optional[dist.ProcessGroup],
      sharded_param_names: Optional[List[str]],
  ) -> torch.Tensor:
      """
      Returns all unsharded gradients as a single flattened tensor. This
      returns the same value on all ranks.
      """
      local_grads_as_flattened = (
          torch.cat(
              [
                  torch.flatten(param.grad)
                  for param in model.parameters()
              ]
          )
          .contiguous()
          .cuda(self.rank)
      )
      all_grads_as_flattened = torch.cat(
          [
              torch.empty_like(local_grads_as_flattened)
              for _ in range(fsdp_pg.size())
          ]
      ).contiguous()
      dist._all_gather_base(
          all_grads_as_flattened,
          local_grads_as_flattened,
          group=fsdp_pg,
      )
      if not uses_tp:
          return all_grads_as_flattened
      splits = tuple(param_name_to_numel.values())
      all_grads_per_param = list(all_grads_as_flattened.split(splits))
      for param_idx, param_name in enumerate(
          param_name_to_numel.keys()
      ):  # assumes fixed order
          if param_name in sharded_param_names:
              local_tensor_size = list(
                  param_name_to_sharding_info[param_name][0]
              )
              sharding_dim = param_name_to_sharding_info[param_name][1]
              local_tensor_size[sharding_dim] //= tp_pg.size()
              local_tensor = all_grads_per_param[param_idx].view(
                  *local_tensor_size
              )
              local_tensors = [
                  torch.empty_like(local_tensor)
                  for _ in range(tp_pg.size())
              ]
              dist.all_gather(local_tensors, local_tensor, group=tp_pg)
              all_grads_per_param[param_idx] = torch.cat(
                  local_tensors, dim=sharding_dim
              ).reshape(-1)
      return torch.cat(all_grads_per_param).contiguous()
  
imports: import copy
  import sys
  from collections import OrderedDict
  from typing import Dict, List, Optional, Tuple
  import torch
  from torch import distributed
  from torch.distributed._tensor import DeviceMesh, distribute_module, DTensor, init_device_mesh, Replicate, Shard
  from torch.distributed._tensor.debug import CommDebugMode
  from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, FullyShardedDataParallel
  from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module, RowwiseParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN
  from torch.testing._internal.distributed._tensor.common_dtensor import MLPModule
label_desc: Use image node_id 4 for calling the TestTPFSDPIntegration obj's underlying member method code with example usage: obj._get_grads_as_flattened(model, uses_tp, param_name_to_numel, param_name_to_sharding_info, tp_pg, fsdp_pg, sharded_param_names) and returns: torch, all_grads_as_flattened