def traced_isend()

in sample_workloads/lit-gpt-demo/utilities/monitor_collectives.py [0:0]

• 5 lines of code
• 2 McCabe index (conditional complexity)

def traced_isend(tensor, dst, group=None, tag=0):
  """Intercepts invocations of torch.distributed.isend.

  Calculate [P2P-B/W] = [Message Size]/[Kernel Time]
  """
  if _should_rank_record_comm(group, peer_rank=dst, is_ring=False):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('send', message_size, group, dst)

  return torch.distributed.untraced_isend(tensor, dst, group, tag)