[QUESTION]

[suzewei]

Hi, I am training my Llama2-7b model with Megatron-LM, using four H20s, 32 GPUs in total. The parallel strategy is set to: TP=8/PP=2/DP=2. Now, I want to know the data capacity of different parallel groups communicating, is there some parameter setting to get these values, or if there is no such parameter, how can I get it from the code? Thank you.

[Advait251206]

Megatron-LM doesn't currently provide a built-in parameter that directly logs the communication volume (bytes transferred) for each parallel group (TP/PP/DP).

However, you can obtain this information in a few ways:

1. Use NCCL profiling (recommended)

   Set the following environment variables before launching your training:

   export NCCL_DEBUG=INFO
   export NCCL_DEBUG_SUBSYS=COLL

   Or for even more detailed logs:

   export NCCL_DEBUG=TRACE

   This will print collective operations (e.g., all_reduce, reduce_scatter, all_gather, broadcast, send, recv) along with the number of elements being communicated. Combined with the tensor datatype (BF16 = 2 bytes, FP16 = 2 bytes, FP32 = 4 bytes), you can compute the communication volume.

2. Instrument the communication wrappers in Megatron-LM

   Megatron-LM wraps most distributed communication calls. Depending on your version, look at files such as:

   megatron/core/tensor_parallel/
   megatron/core/pipeline_parallel/
   megatron/core/distributed/
   

   or search for usages of:

   torch.distributed.all_reduce
   torch.distributed.all_gather
   torch.distributed.reduce_scatter
   torch.distributed.broadcast
   torch.distributed.send
   torch.distributed.recv

   Before each communication call, you can log the tensor size:

   bytes_transferred = tensor.numel() * tensor.element_size()
   rank = torch.distributed.get_rank()
   
   print(
       f"Rank {rank}: {op_name}, "
       f"{bytes_transferred / 1024 / 1024:.2f} MB, "
       f"group={group}"
   )

   Since Megatron creates separate process groups for TP, PP, DP (and sometimes CP/EP), logging the associated process group lets you distinguish which parallel strategy each communication belongs to.

3. Use the PyTorch Profiler

   The PyTorch profiler can also capture distributed communication events:

   with torch.profiler.profile(
       activities=[
           torch.profiler.ProfilerActivity.CPU,
           torch.profiler.ProfilerActivity.CUDA,
       ],
       profile_memory=True,
       with_stack=True,
   ) as prof:
       ...

   Export the trace to Chrome or TensorBoard to inspect NCCL operations, their durations, and communication sizes.

For your configuration (TP=8, PP=2, DP=2, total 32 GPUs), the communication groups are:

• Tensor Parallel (8 GPUs/group): Primarily all_reduce, reduce_scatter, and all_gather.
• Pipeline Parallel (2 GPUs/group): Mainly point-to-point send/recv of activations and gradients.
• Data Parallel (2 GPUs/group): Gradient synchronization using reduce_scatter/all_reduce, along with parameter synchronization depending on the optimizer implementation.

If your goal is to measure the actual communication volume per TP/PP/DP group, instrumenting the communication wrappers or enabling NCCL profiling will provide the most accurate information. The PyTorch profiler is also useful if you want to correlate communication with execution time and identify potential bottlenecks.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206