[QUESTION] What is the principle of pre-communication-optimization?

[LitLeo]

Your question

--no-pre-communication-optimization
By default zb runtime dispatches a tiny communication before the real communication to optimize computation
Why?

code

zero-bubble-pipeline-parallelism/megatron/core/pipeline_parallel/zerobubble/runtime.py

Line 756 in 7e03eac

# Cannot fuse "pre_send" with other send kernels, or they will get stuck

# Cannot fuse "pre_send" with other send kernels, or they will get stuck
# possibly as there will be 2 send-recv with the same source and target.
with nvtx_range_ctx("pre_send"):
    pre_send, _ = multi_pipeline_ops(
        pre_sp_tensors, [],
        pre_sn_tensors, [],
        batch_p2p,
    )
with nvtx_range_ctx(send_fused_name):
    send_reqs, _ = multi_pipeline_ops(
        sp_tensors, [],
        sn_tensors, [],
        batch_p2p,
    )

[ufotalent]

Hi @LitLeo , thanks for interest our work. This small op act as something like a cross-node cuda event. As a small nccl kernel only occupy a small number of GPU SMs, it prevents the main NCCL communication, which overlaps with compute, to occupy too many SMs which delays the computation. In other words, the main NCCL communication is launched only after both send/recv endpoint becomes ready, which is done by inserting theses tiny communications.