The variable dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"] is misleadingly named. By default both layers are placed in layers_all_gather_overlap with method="bulk". They are only moved to layers_reduce_scatter_overlap under two specific conditions:
The user explicitly provides a non-bulk method for these layers via ub_cfgs
with_cublasmp=True is set, since cuBLASMp doesn't support bulk overlap
Why the name is misleading:-
The name implies these dgrad layers always use reduce-scatter overlap, but the default is the opposite -- they default to all-gather overlap. What actually makes them special is that they are the only two GEMMs in the transformer block that support communication on both sides (all-gather on input, reduce-scatter on output), making their overlap strategy switchable -- unlike every other GEMM which has only one comm direction.
The variable
dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"]is misleadingly named. By default both layers are placed inlayers_all_gather_overlapwith method="bulk". They are only moved tolayers_reduce_scatter_overlapunder two specific conditions:The user explicitly provides a non-bulk method for these layers via ub_cfgs
with_cublasmp=True is set, since cuBLASMp doesn't support bulk overlap
Why the name is misleading:-
The name implies these dgrad layers always use reduce-scatter overlap, but the default is the opposite -- they default to all-gather overlap. What actually makes them special is that they are the only two GEMMs in the transformer block that support communication on both sides (all-gather on input, reduce-scatter on output), making their overlap strategy switchable -- unlike every other GEMM which has only one comm direction.