sample_workloads/lit-gpt-demo/utilities/monitor_collectives.py

"""A utility to trace torch.distributed calls.

Traces torch.distributed collectives before dispatch. In particular, logs the
collective kind (all_reduce, all_to_all, ..), message size (10 MB), and which
GPU devices are participating ([0, 1, 6, 7]). These are logged as NVTX markers
by NVIDIA Nsight, as well as printed to stdout. By default, we only log
cross-node collective communications.

To assist with computing the effective bandwidth of a collective, a nominal
expression is provided in the doc string of each 'traced_<collective>'. This
also requires extracting the timings of the corresponding NCCL kernels.

Typical usage example:

  import utilities.monitor_collectives
  utilities.monitor_collectives.shunt_torch_communication()

When running a workload, also define TORCH_DISTRIBUTED_TRACING to be one of
'ALL' or 'CROSSNODE'. See `should_rank_record_comm` for added details.
"""


import functools
import inspect
import io
import json
import os
import pickle
import sys
from datetime import datetime
import calendar
import uuid

import nvtx
import torch.cuda
import torch.distributed


_TRACE_MODE = None


# Note: By default, we only target tracing *cross-node* communications.
# See 'should_rank_record_comm'
def shunt_torch_communication():
  _identify_trace_mode()
  if _TRACE_MODE == 'none':
    if int(os.environ.get("RANK", "0")) == 0:
      print('Tracing torch.distributed collectives disabled.', flush=True)
    return

  _shunt_torch_communication_objects()
  _shunt_torch_communication_calls()

  if int(os.environ.get("RANK", "0")) == 0:
    print('NVTX and print tracing of torch.distributed collectives enabled.',
          flush=True)
    print(f"{_GPU_SERIAL=}, {_VM_ID=}")

    if not _SHOULD_PRINT:
      print('Collectives are traced but will not be printed to stdout', flush=True)


def _identify_trace_mode():
  global _TRACE_MODE
  _TRACE_MODE = os.environ.get('TORCH_DISTRIBUTED_TRACING', 'CROSSNODE')
  _TRACE_MODE = _TRACE_MODE.lower()

  global _SHOULD_PRINT
  _SHOULD_PRINT = os.environ.get('TORCH_DISTRIBUTED_TRACING_PRINT', 'False')
  _SHOULD_PRINT = _SHOULD_PRINT.lower() in ['true', '1', 't', 'y', 'yes']

  global _GPU_SERIAL
  _GPU_SERIAL = os.environ.get("GPU_SERIAL", "unknown")
  global _VM_ID
  _VM_ID = os.environ.get("VM_ID", "unknown")


# Each wrapper should match format 'traced_<collective>'
def _shunt_torch_communication_calls():
  """Replaces torch.distributed.<target_collective> with a traced version.
  """
  target_collectives = [
      'barrier',
      'broadcast_object_list',
      'broadcast',
      'gather',
      'scatter',
      'reduce',
      'reduce_scatter',
      'reduce_scatter_tensor',
      'all_reduce',
      'all_gather',
      'all_gather_into_tensor',
      'all_to_all',
      'all_to_all_single',
      'batch_isend_irecv',
      'isend',
      'irecv',
      'send',
      'recv',
  ]

  this_module = sys.modules[__name__]
  for collective in target_collectives:
    original_fn = getattr(torch.distributed, collective)
    replaced_fn = getattr(this_module, 'traced_' + collective)
    setattr(torch.distributed, 'untraced_' + collective, original_fn)
    setattr(torch.distributed, collective, replaced_fn)


def _shunt_torch_communication_objects():
  original_p2p = torch.distributed.P2POp
  setattr(torch.distributed, 'UntracedP2POp', original_p2p)
  setattr(torch.distributed, 'P2POp', _TracedP2POp)


# Each 'traced_<comm>' defines a 'message_size' to compute B/W.
# Ref https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md

# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_barrier(group=None, async_op=False, device_ids=None):
  """Intercepts invocations of torch.distributed.barrier.
  """
  if _should_rank_record_comm(group):
    _emit_call_description('barrier', message_size=1, group=group)

  return torch.distributed.untraced_barrier(group, async_op, device_ids)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_broadcast_object_list(object_list, src=0, group=None, device=None):
  """Intercepts invocations of torch.distributed.broadcast_object_list.

  Converts objects to tensor data using the pickle library. Then conducts a
  torch.distributed.broadcast call.
  """

  if _should_rank_record_comm(group, root_rank=src):
    message_size = 0
    for obj in object_list:
      # Note: This computation is sadly redundant with underlying call :(
      # For now we don't expect this invocation to be in critical path.
      buf = io.BytesIO()
      pickle.Pickler(buf).dump(obj)
      message_size += buf.getbuffer().nbytes
    _emit_call_description(
        'broadcast_object_list', message_size, group, root_rank=src)

  return torch.distributed.untraced_broadcast_object_list(
      object_list, src, group, device)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_broadcast(tensor, src, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.broadcast.

  Calculate [Ring-B/W] = [Message Size]/[Kernel Time] for large [Message Size]

  https://images.nvidia.com/events/sc15/pdfs/NCCL-Woolley.pdf
  """
  if _should_rank_record_comm(group, root_rank=src):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('broadcast', message_size, group, root_rank=src)

  return torch.distributed.untraced_broadcast(
      tensor, src, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_gather(
    tensor, gather_list=None, dst=0, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.gather.

  Let T := sum([Receive Kernel Time from Rank i] for i != dst)
  Calculate [P2P-B/W] = [Message Size]/T

  Each of (n-1) ranks sends a message to the root.

  Note that any correction factors for the bus bandwidth (e.g. [n-1]/n) depend
  on the *definition* of 'Message Size'. In some cases, such as for 'gather', we
  define 'Message Size' so as to omit the size of data that is already local
  to the destination GPU for the 'gather' operation. In this case, no correction
  factor is needed. In NCCL tests, they assume all ranks send equal sized
  messages and include this size of data already resident on the destination
  GPU. Thus, in there case you see a (n-1)/n correction factor on calculating
  the bus bandwidth. In general, the goal of computing the bus bandwidth is
  to compare data transfer rates on the bus relative to peak bus bandwidth.
  See https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md.

  https://github.com/NVIDIA/nccl-tests/blob/1a5f551ffd6e/src/gather.cu#L54
  https://github.com/pytorch/pytorch/blob/bfd995f0d6bf/torch/csrc/cuda/nccl.cpp#L1040
  """
  if _should_rank_record_comm(group, root_rank=dst, is_ring=False):
    message_size = functools.reduce(
        lambda sz, x: sz + x.nelement() * x.element_size(), gather_list, 0)
    message_size -= tensor.nelement() * tensor.element_size()

    _emit_call_description('gather', message_size, group, root_rank=dst)

  return torch.distributed.untraced_gather(
      tensor, gather_list, dst, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_scatter(
    tensor, scatter_list=None, src=0, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.scatter.

  Let T := sum([Send Kernel Time from Rank i] for i != src)
  Calculate [P2P-B/W] = [Message Size]/T

  Each of (n-1) ranks receives a message from the root.
  There is no (n-1)/n factor as we factor it in [Message Size].

  https://github.com/NVIDIA/nccl-tests/blob/1a5f551ffd6e/src/scatter.cu#L50
  https://github.com/pytorch/pytorch/blob/bfd995f0d6bf/torch/csrc/cuda/nccl.cpp#L1089
  """
  if _should_rank_record_comm(group, root_rank=src, is_ring=False):
    message_size = functools.reduce(
        lambda sz, x: sz + x.nelement() * x.element_size(), scatter_list, 0)
    message_size -= tensor.nelement() * tensor.element_size()

    _emit_call_description('scatter', message_size, group, root_rank=src)

  return torch.distributed.untraced_scatter(
      tensor, scatter_list, src, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_reduce(
    tensor, dst, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.reduce.

  Calculate [Ring-B/W] = [Message Size]/[Kernel Time] for large [Message Size]
  Also see 'traced_broadcast'
  """
  if _should_rank_record_comm(group, root_rank=dst):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('reduce', message_size, group, root_rank=dst)

  return torch.distributed.untraced_reduce(tensor, dst, op, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_reduce_scatter(
    output,
    input_list,
    op=torch.distributed.ReduceOp.SUM,
    group=None,
    async_op=False):
  """Intercepts invocations of torch.distributed.reduce_scatter.

  Let n := [Group Size].
  Calculate [Ring-B/W] = (n-1)/n * [Message Size]/[Kernel Time]
  Assumes equal tensor sizes. It's the same as first half of ring All-Reduce.
  """
  if _should_rank_record_comm(group):
    message_size = output.nelement() * output.element_size()
    _emit_call_description('reduce_scatter', message_size, group)

  return torch.distributed.untraced_reduce_scatter(
      output, input_list, op, group, async_op)


# pylint: disable=redefined-builtin,g-doc-args,g-doc-return-or-yield
def traced_reduce_scatter_tensor(
    output,
    input,
    op=torch.distributed.ReduceOp.SUM,
    group=None,
    async_op=False):
  """Intercepts invocations of torch.distributed.reduce_scatter_tensor.

  Similar to 'traced_reduce_scatter'
  """

  if _should_rank_record_comm(group):
    message_size = output.nelement() * output.element_size()
    _emit_call_description('reduce_scatter', message_size, group)

  return torch.distributed.untraced_reduce_scatter_tensor(
      output, input, op, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_all_reduce(
    tensor, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.all_reduce.

  Let n := [Group Size]
  Calculate [Ring-B/W] = 2(n-1)/n * [Message Size] / [Kernel Time]

  https://images.nvidia.com/events/sc15/pdfs/NCCL-Woolley.pdf
  """
  if _should_rank_record_comm(group):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('all_reduce', message_size, group)

  return torch.distributed.untraced_all_reduce(
      tensor, op, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_all_gather(tensor_list, tensor, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.all_gather.

  Let n := [Group Size]
  Calculate [Ring-B/W] = (n-1)/n * [Message Size] / [Kernel Time]
  Assuming equal tensor sizes.
  """
  if _should_rank_record_comm(group):
    message_size = functools.reduce(
        lambda size, x: size + x.nelement() * x.element_size(), tensor_list, 0)
    _emit_call_description('all_gather', message_size, group)

  return torch.distributed.untraced_all_gather(
      tensor_list, tensor, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_all_gather_into_tensor(
    output_tensor, input_tensor, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.all_gather_into_tensor.

  Similar 'traced_all_gather'
  """
  if _should_rank_record_comm(group):
    message_size = output_tensor.nelement() * output_tensor.element_size()
    _emit_call_description('all_gather', message_size, group)

  return torch.distributed.untraced_all_gather_into_tensor(
      output_tensor, input_tensor, group, async_op)


# Note: The TCP Direct team intends to implement a custom version of AllToAll.
# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_all_to_all(
    output_tensor_list, input_tensor_list, group=None, async_op=False):
  """Intercepts invocations of torch.distributed.all_to_all.

  Let S := sum([Message Size on Rank i] for i = 1..n) where n := [Group Size]
  Let T := [End of last Receive last rank] - [Start of first Send first rank]
  Calculate [Algo B/W] = S / T.

  There is no n/(n-1) correction factor as we factor it in [Message Size].

  https://github.com/NVIDIA/nccl-tests/blob/1a5f551ffd6e/src/alltoall.cu#L57
  https://github.com/pytorch/pytorch/blob/bfd995f0d6bf/torch/csrc/cuda/nccl.cpp#L911
  """
  if _should_rank_record_comm(group):
    message_size = functools.reduce(
        lambda s, x: s + x.nelement() * x.element_size(), input_tensor_list, 0)

    # Omit bytes corresponding to send and receive on the same rank
    self_tensor = input_tensor_list[torch.distributed.get_rank(group)]
    message_size -= self_tensor.nelement() * self_tensor.element_size()

    _emit_call_description('all_to_all', message_size, group)

  return torch.distributed.untraced_all_to_all(
      output_tensor_list, input_tensor_list, group, async_op)


# pylint: disable=g-doc-args,g-doc-return-or-yield,redefined-builtin
def traced_all_to_all_single(
    output,
    input,
    output_split_sizes=None,
    input_split_sizes=None,
    group=None,
    async_op=False):
  """Intercepts invocations of torch.distributed.all_to_all_single.

  Similar to 'traced_all_to_all'
  """
  if _should_rank_record_comm(group):
    self_rank = torch.distributed.get_rank(group)

    if input_split_sizes is not None:
      self_slice = input_split_sizes[self_rank]
    else:
      self_slice = input.size(dim=0) / torch.distributed.get_world_size(group)

    slice_nelement = input.nelement() / input.size(dim=0)
    message_size = input.nelement() * input.element_size()
    message_size -= self_slice * slice_nelement * input.element_size()

    _emit_call_description('all_to_all_single', message_size, group)

  return torch.distributed.untraced_all_to_all_single(
      output, input, output_split_sizes, input_split_sizes, group, async_op)


# Note: Each send and receive occurs on indepenent CUDA streams
# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_batch_isend_irecv(p2p_op_list):
  """Intercepts invocations of torch.distributed.batch_isend_irecv.

  Calculate [P2P-B/W] = [Message Size]/[Kernel Time] for each send and recv.
  """
  correlation_id = str(uuid.uuid4())
  for p2p in p2p_op_list:
    if _SHOULD_PRINT:
      print(f"Num p2p ops in batch: {len(p2p_op_list)}")
    if _should_rank_record_comm(p2p.group, peer_rank=p2p.peer, is_ring=False):
      api = 'send' if p2p.op == torch.distributed.untraced_isend else 'recv'

      message_size = p2p.tensor.nelement() * p2p.tensor.element_size()
      _emit_call_description(api, message_size, group=p2p.group, peer_rank=p2p.peer, correlation_id=correlation_id)

  return torch.distributed.untraced_batch_isend_irecv(p2p_op_list)


# pylint: disable=g-doc-args,g-doc-return-or-yield
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)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_irecv(tensor, src=None, group=None, tag=0):
  """Intercepts invocations of torch.distributed.irecv.
  """
  if _should_rank_record_comm(group, peer_rank=src, is_ring=False):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('recv', message_size, group, src)

  return torch.distributed.untraced_irecv(tensor, src, group, tag)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_send(tensor, dst, group=None, tag=0):
  """Intercepts invocations of torch.distributed.send.
  """
  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_send(tensor, dst, group, tag)


# pylint: disable=g-doc-args,g-doc-return-or-yield
def traced_recv(tensor, src=None, group=None, tag=0):
  """Intercepts invocations of torch.distributed.recv.
  """
  if _should_rank_record_comm(group, peer_rank=src, is_ring=False):
    message_size = tensor.nelement() * tensor.element_size()
    _emit_call_description('recv', message_size, group, src)

  return torch.distributed.untraced_recv(tensor, src, group, tag)


@functools.lru_cache(maxsize=None)
def _should_rank_record_comm(
    group=None, peer_rank=None, root_rank=None, is_ring=True):
  """Decides whether a given torch.distributed collective should be recorded.

  Args:
    group: The torch process group (i.e. participating GPUs) in this collective.
    peer_rank: In direct peer to peer operations, the global rank of the peer.
    root_rank: The global rank of the root GPU, for collectives with a root.
    as_ring: Whether the default NCCL implementation uses a ring algorithm.
    Specifying 'peer_rank' and 'is_ring=True' are incompatible.

  Returns:
    Whether to record a descriptive NVTX marker, and possibly print a log trace.
  """
  if not _is_current_process_in_group(group):
    return False
  if _TRACE_MODE == 'crossnode' and not _is_crossnode_comm(group, peer_rank):
    return False
  if not is_ring and root_rank is not None:
    return torch.distributed.get_rank() == root_rank

  return True


def _is_current_process_in_group(group=None):
  return torch.distributed.get_rank(group) >= 0


@functools.lru_cache(maxsize=None)
def _is_crossnode_comm(group=None, peer_rank=None):
  """Whether this collective involves communication across nodes.

  Args:
    group: The torch process group (i.e. participating GPUs) in this collective.
    peer: In direct peer to peer operations, the global rank of the peer.

  Returns:
    Whether this collective involves communications across nodes.
  """
  count_per_node = torch.cuda.device_count()

  if peer_rank is not None:
    this_node = int(torch.distributed.get_rank() / count_per_node)
    peer_node = int(peer_rank / count_per_node)
    return this_node != peer_node
  else:
    if group is not None:
      ranks = torch.distributed.get_process_group_ranks(group=group)
    else:
      ranks = [*range(torch.distributed.get_world_size())]

    nodes = list(map(lambda rank: int(rank / count_per_node), ranks))
    return any([node != nodes[0] for node in nodes])


def _emit_call_description(
    name, message_size, group=None, peer_rank=None, root_rank=None, correlation_id=None):
  call_description = _TorchDistributedCallDescriptor(
      name, message_size, group, peer_rank, root_rank, correlation_id).to_json()

  nvtx.mark(call_description)
  if _should_rank_print(group, peer_rank, root_rank):
    print(call_description)


class _TorchDistributedCallDescriptor:
  """Description of a torch.distributed comm call to be stored as NVTX marker.
  """

  def __init__(
      self, name, message_size, group=None, peer_rank=None, root_rank=None, correlation_id=None):
    self.name = name
    self.rank = torch.distributed.get_rank()
    self.source_line = _get_call_source_line()
    self.message_size = message_size
    self.device = torch.cuda.current_device()
    self.timestamp = calendar.timegm(datetime.utcnow().utctimetuple())
    self.gpu_serial = _GPU_SERIAL
    self.vm_id = _VM_ID
    if group is not None:
      self.group_ranks = torch.distributed.get_process_group_ranks(group=group)
    if peer_rank is not None:
      self.peer_rank = peer_rank
    if root_rank is not None:
      self.root_rank = root_rank
    if correlation_id is not None:
        self.correlation_id = correlation_id

  def to_json(self):
    return json.dumps(self, default=lambda o: o.__dict__)


def _should_rank_print(group=None, peer_rank=None, root_rank=None):
  if not _SHOULD_PRINT:
    return False
  if root_rank is not None:
    leader = root_rank
  elif group is not None:
    leader = torch.distributed.get_global_rank(group, 0)
  else:
    leader = 0

  return (peer_rank is not None) or torch.distributed.get_rank() == leader


# A fixed depth works for all cases here
def _get_call_source_line(depth=4):
  caller = inspect.getframeinfo(inspect.stack()[depth][0])
  return '{}:{}'.format(caller.filename, caller.lineno)


# We need to un-hide the original type for 'batch_isend_irecv' due to type
# checks performed by torch.distributed. This is not an issue as by then we
# have already recorded the call.
class _TracedP2POp(torch.distributed.P2POp):
  """Used to redirect torch.distributed.i{send,recv} on 'batch_isend_irecv'.
  """

  def __init__(self, op, tensor, peer, group=None, tag=0):
    original_op = _get_original_p2p_op(op)
    torch.distributed.UntracedP2POp.__init__(
        self, original_op, tensor, peer, group, tag)

  def __new__(cls, op, tensor, peer, group=None, tag=0):
    original_op = _get_original_p2p_op(op)
    return torch.distributed.UntracedP2POp.__new__(
        cls, original_op, tensor, peer, group, tag)


def _get_original_p2p_op(op):
  if op == torch.distributed.isend:
    return torch.distributed.untraced_isend
  elif op == torch.distributed.irecv:
    return torch.distributed.untraced_irecv