docstr: Syncs the tensor parallel parameters' gradients following the data
  parallel paradigm where gradients are averaged over ranks (in this
  case, the ones in the tensor parallel process group).
docstr_tok: ["Syncs","the","tensor","parallel","parameters","'","gradients","following","the","data","parallel","paradigm","where","gradients","are","averaged","over","ranks","(","in","this","case",",","the","ones","in","the","tensor","parallel","process","group",")","."]
code: def _sync_tp_grads(
      self,
      tp_fsdp_model: FSDP,
      tp_pg: dist.ProcessGroup,
      param_name_to_numel: Dict[str, int],
      non_sharded_param_names: List[str],
  ) -> None:
      """
      Syncs the tensor parallel parameters' gradients following the data
      parallel paradigm where gradients are averaged over ranks (in this
      case, the ones in the tensor parallel process group).
      """
      tp_world_size = tp_pg.size()
      fsdp_world_size = self.world_size // tp_world_size
      assert (
          type(tp_fsdp_model) is FSDP
          and len(list(tp_fsdp_model.parameters())) == 1
      ), (
          "The following logic assumes a single top-level-only FSDP wrapping "
          "the model with TP already applied"
      )
      flat_param = tp_fsdp_model.params[0]
      splits = tuple(param_name_to_numel.values())
      # Create a mask over the gradient elements to manually reduce
      unsharded_size = torch.Size(
          [flat_param.numel() * fsdp_world_size]
      )
      unsharded_zeros = torch.zeros(
          unsharded_size, device=flat_param.device
      )
      per_param_masks = unsharded_zeros.split(splits)
      for param_idx, param_name in enumerate(
          param_name_to_numel.keys()
      ):  # assumes fixed order
          if param_name not in non_sharded_param_names:
              per_param_masks[param_idx][:] = 1
      unsharded_mask = (
          torch.cat(per_param_masks).contiguous().type(torch.BoolTensor)
      )
      sharded_mask = unsharded_mask.chunk(fsdp_world_size)[
          self.rank // tp_world_size
      ]
      grad_device = flat_param.grad.device
      grad = flat_param.grad.detach().clone().cuda(self.rank)
      dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=tp_pg)
      grad = grad.to(grad_device)
      flat_param.grad[~sharded_mask] = grad[~sharded_mask]
      # Average *all* gradient elements to match the FSDP only semantics
      flat_param.grad /= tp_world_size
  
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 3 for calling the TestTPFSDPIntegration obj's underlying member method code with example usage: obj._sync_tp_grads(tp_fsdp_model, tp_pg, param_name_to_numel, non_sharded_param_names) without return types