AI in A number of GPUs: Level-to-Level and Collective Operations

• Habits: Once you name a synchronous communication technique, the host course of stops and waits till the NCCL kernel is efficiently enqueued on the present energetic CUDA stream. As soon as enqueued, the operate returns. That is normally easy and dependable. Notice that the host is just not ready for the switch to finish, only for the operation to be enqueued. Nonetheless, it blocks that particular stream from shifting on to the following operation till the NCCL kernel is executed to completion.

• Habits: Once you name an asynchronous communication technique, the decision returns instantly, and the enqueuing operation occurs within the background. It doesn’t enqueue into the present energetic stream, however moderately to a devoted inside NCCL stream per machine. This permits your CPU to proceed with different duties, a way referred to as overlapping computation with communication. The asynchronous API is extra complicated as a result of it may result in undefined habits should you don’t correctly use .wait() (defined beneath) and modify knowledge whereas it’s being transferred. Nonetheless, mastering it’s key to unlocking most efficiency in large-scale distributed coaching.

Level-to-Level (One-to-One)

These operations are usually not thought-about collectives, however they’re foundational communication primitives. They facilitate direct knowledge switch between two particular ranks and are basic for duties the place one GPU must ship particular info to a different.

• Synchronous (Blocking): The host course of waits for the operation to be enqueued to the CUDA stream earlier than continuing. The kernel is enqueued into the present energetic stream.
  • torch.distributed.ship(tensor, dst): Sends a tensor to a specified vacation spot rank.
  • torch.distributed.recv(tensor, src): Receives a tensor from a supply rank. The receiving tensor should be pre-allocated with the proper form and dtype.
• Asynchronous (Non-Blocking): The host course of initiates the enqueue operation and instantly continues with different duties. The kernel is enqueued right into a devoted inside NCCL stream per machine, which permits for overlapping communication with computation. These operations return a request(technically a Work object) that can be utilized to trace the enqueuing standing.
  • request = torch.distributed.isend(tensor, dst): Initiates an asynchronous ship operation.
  • request = torch.distributed.irecv(tensor, src): Initiates an asynchronous obtain operation.
  • request.wait(): Blocks the host solely till the operation has been efficiently enqueued on the CUDA stream. Nonetheless, it does block the at present energetic CUDA stream from executing later kernels till this particular asynchronous operation completes.
  • request.wait(timeout): Should you present a timeout argument, the host habits modifications. It can block the CPU thread till the NCCL work completes or occasions out (elevating an exception). In regular instances, customers don’t must set the timeout.
  • request.is_completed(): Returns True if the operation has been efficiently enqueued onto a CUDA stream. It might be used for polling. It doesn’t assure that the precise knowledge has been transferred.

When PyTorch launches an NCCL kernel, it robotically inserts a dependency (i.e. forces a synchronization) between your present energetic stream and the NCCL stream. This implies the NCCL stream gained’t begin till all beforehand enqueued work on the energetic stream finishes — guaranteeing the tensor being despatched already holds the ultimate values.

Equally, calling req.wait() inserts a dependency within the different path. Any work you enqueue on the present stream after req.wait() gained’t execute till the NCCL operation completes, so you’ll be able to safely use the obtained tensors.

Whereas ship and recv are labeled “synchronous,” their habits in NCCL might be complicated. A synchronous name on a CUDA tensor blocks the host CPU thread solely till the info switch kernel is enqueued to the stream, not till the info switch completes. The CPU is then free to enqueue different duties.

There’s an exception: the very first name to torch.distributed.recv() in a course of is really blocking and waits for the switch to complete, doubtless because of inside NCCL warm-up procedures. Subsequent calls will solely block till the operation is enqueued.

Contemplate this instance the place rank 1 hangs as a result of the CPU tries to entry a tensor that the GPU has not but obtained:

rank = torch.distributed.get_rank()
if rank == 0:
   t = torch.tensor([1,2,3], dtype=torch.float32, machine=machine)
   # torch.distributed.ship(t, dst=1) # No ship operation is carried out
else: # rank == 1 (assuming solely 2 ranks)
   t = torch.empty(3, dtype=torch.float32, machine=machine)
   torch.distributed.recv(t, src=0) # Blocks solely till enqueued (after first run)
   print("This WILL print if NCCL is warmed-up")
   print(t) # CPU wants knowledge from GPU, inflicting a block
   print("This can NOT print")

The CPU course of at rank 1 will get caught on print(t) as a result of it triggers a host-device synchronization to entry the tensor’s knowledge, which by no means arrives.

Should you run this code a number of occasions, discover that This WILL print if NCCL is warmed-up won’t get printed within the later executions, because the CPU remains to be caught at print(t).

Collectives

Each collective operation operate helps each sync and async operations by way of the async_op argument. It defaults to False, that means synchronous operations.

These operations contain one rank sending knowledge to all different ranks within the group.

• torch.distributed.broadcast(tensor, src): Copies a tensor from a single supply rank (src) to all different ranks. Each course of finally ends up with an similar copy of the tensor. The tensor parameter serves two functions: (1) when the rank of the method matches the src, the tensor is the info being despatched; (2) in any other case, tensor is used to avoid wasting the obtained knowledge.

rank = torch.distributed.get_rank()
if rank == 0: # supply rank
  tensor = torch.tensor([1,2,3], dtype=torch.int64, machine=machine)
else: # vacation spot ranks
  tensor = torch.empty(3, dtype=torch.int64, machine=machine)
torch.distributed.broadcast(tensor, src=0)

• torch.distributed.scatter(tensor, scatter_list, src): Distributes chunks of information from a supply rank throughout all ranks. The scatter_list on the supply rank incorporates a number of tensors, and every rank (together with the supply) receives one tensor from this listing into its tensor variable. The vacation spot ranks simply move None for the scatter_list.

# The scatter_list should be None for all non-source ranks.
scatter_list = None if rank != 0 else [torch.tensor([i, i+1]).to(machine) for i in vary(0,4,2)]
tensor = torch.empty(2, dtype=torch.int64).to(machine)
torch.distributed.scatter(tensor, scatter_list, src=0)
print(f'Rank {rank} obtained: {tensor}')

These operations collect knowledge from all ranks and consolidate it onto a single vacation spot rank.

• torch.distributed.cut back(tensor, dst, op): Takes a tensor from every rank, applies a discount operation (like SUM, MAX, MIN), and shops the ultimate end result on the vacation spot rank (dst) solely.

rank = torch.distributed.get_rank()
tensor = torch.tensor([rank+1, rank+2, rank+3], machine=machine)
torch.distributed.cut back(tensor, dst=0, op=torch.distributed.ReduceOp.SUM)
print(tensor)

• torch.distributed.collect(tensor, gather_list, dst): Gathers a tensor from each rank into an inventory of tensors on the vacation spot rank. The gather_list should be an inventory of tensors (appropriately sized and typed) on the vacation spot and None in every single place else.

# The gather_list should be None for all non-destination ranks.
rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
gather_list = None if rank != 0 else [torch.zeros(3, dtype=torch.int64).to(device) for _ in range(world_size)]
t = torch.tensor([0+rank, 1+rank, 2+rank], dtype=torch.int64).to(machine)
torch.distributed.collect(t, gather_list, dst=0)
print(f'After op, Rank {rank} has: {gather_list}')

The variable world_size is the full variety of ranks. It may be obtained with torch.distributed.get_world_size(). However don’t fear about implementation particulars for now, crucial factor is to know the ideas.

In these operations, each rank each sends and receives knowledge from all different ranks.

• torch.distributed.all_reduce(tensor, op): Similar as cut back, however the result’s saved on eachrank as an alternative of only one vacation spot.

# Instance for torch.distributed.all_reduce
rank = torch.distributed.get_rank()
tensor = torch.tensor([rank+1, rank+2, rank+3], dtype=torch.float32, machine=machine)
torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
print(f"Rank {rank} after all_reduce: {tensor}")

• torch.distributed.all_gather(tensor_list, tensor): Similar as collect, however the gathered listing of tensors is out there on each rank.

# Instance for torch.distributed.all_gather
rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
input_tensor = torch.tensor([rank], dtype=torch.float32, machine=machine)
tensor_list = [torch.empty(1, dtype=torch.float32, device=device) for _ in range(world_size)]
torch.distributed.all_gather(tensor_list, input_tensor)
print(f"Rank {rank} gathered: {[t.item() for t in tensor_list]}")

• torch.distributed.reduce_scatter(output, input_list): Equal of performing a cut back operation on an inventory of tensors after which scattering the outcomes. Every rank receives a special a part of the diminished output.

# Instance for torch.distributed.reduce_scatter
rank = torch.distributed.get_rank()
world_size = torch.distributed.get_world_size()
input_list = [torch.tensor([rank + i], dtype=torch.float32, machine=machine) for i in vary(world_size)]
output = torch.empty(1, dtype=torch.float32, machine=machine)
torch.distributed.reduce_scatter(output, input_list, op=torch.distributed.ReduceOp.SUM)
print(f"Rank {rank} obtained diminished worth: {output.merchandise()}")

Synchronization

The 2 most incessantly used operations are request.wait() and torch.cuda.synchronize(). It’s essential to grasp the distinction between these two:

• request.wait(): That is used for asynchronous operations. It synchronizes the at present energetic CUDA stream for that operation, guaranteeing the stream waits for the communication to finish earlier than continuing. In different phrases, it blocks the at present energetic CUDA stream till the info switch finishes. On the host facet, it solely causes the host to attend till the kernel is enqueued; the host does not look forward to the info switch to finish.
• torch.cuda.synchronize(): This can be a extra forceful command that pauses the host CPU thread till all beforehand enqueued duties on the GPU have completed. It ensures that the GPU is totally idle earlier than the CPU strikes on, however it may create efficiency bottlenecks if used improperly. Every time it’s essential to carry out benchmark measurements, you need to use this to make sure you seize the precise second the GPUs are carried out.

Conclusion

Congratulations on making it to the top! On this put up, you discovered about:

• Level-to-Level Operations
• Sync and Async in NCCL
• Collective operations
• Synchronization strategies

Within the subsequent weblog put up we’ll dive into PCIe, NVLink, and different mechanisms that allow communication in a distributed setting!

References