Function annotate_collective

advanced_source/cuda_graph_annotations_tutorial.py:664–696

Annotate a collective with the metadata eager NCCL traces expose. Returns a ``mark_kernels`` context manager. Any kernels launched inside (i.e. the collective) are tagged with the collective type, message sizes, dtype, and the process group's name/description/ranks, and placed on a

(collective_name, input_tensor, output_tensor, group=None)

# annotated trace consistent with non-graphed traces.

def annotate_collective(collective_name, input_tensor, output_tensor, group=None):
    """Annotate a collective with the metadata eager NCCL traces expose.

    Returns a ``mark_kernels`` context manager. Any kernels launched inside
    (i.e. the collective) are tagged with the collective type, message sizes,
    dtype, and the process group's name/description/ranks, and placed on a
    dedicated lane keyed by the process group so comms are visually separated
    from compute.

    The field names match the keys the profiler records for eager collectives
    (``In msg nelems``, ``Group size``, ``Process Group Name``, ...), so an
    annotated graphed collective reads exactly like a non-graphed one.
    """
    pg = group if group is not None else (dist.group.WORLD if dist.is_initialized() else None)
    ranks = dist.get_process_group_ranks(pg) if pg is not None else [0]
    group_name = getattr(pg, "group_name", "default")
    group_desc = getattr(pg, "group_desc", "default")

    # NCCL always uses its own internal stream, so key the lane on the process
    # group (name + description) and give it a stable id (>= 60).
    pg_key = f"{group_name}_{group_desc}"
    annotation = {
        "name": collective_name,
        "In msg nelems": input_tensor.numel(),
        "Out msg nelems": output_tensor.numel(),
        "Group size": len(ranks),
        "dtype": str(input_tensor.dtype).replace("torch.", ""),
        "Process Group Name": group_name,
        "Process Group Description": group_desc,
        "Process Group Ranks": ranks,
        "stream": get_stream_for_pg(pg_key),
    }
    return mark_kernels(annotation)

###############################################################################
# A Block That Mixes Compute and Communication