[PyTorch] Bugfix for wgrad bulk overlap conflict when dgrad overlap is reduce-scatter #1341
Conversation
|
/te-ci pytorch L0 L1 |
denera
force-pushed
the
rs-dgrad-overlap-bugfix
branch
from
ab9e05f to
2ca29de
Compare
|
/te-ci pytorch L0 L1 |
| if name in ub_cfgs and "method" in ub_cfgs[name] and ub_cfgs[name]["method"] != "bulk": | ||
| # Loop over user configs and disable dgrad and wgrad bulk overlaps for every layer that has a | ||
| # reduce-scatter dgrad overlap. | ||
| ub_cfg = {} if ub_cfg is None else ub_cfg |
There was a problem hiding this comment.
local variable ub_cfg referenced before assignment.
ub_cfg --> ub_cfgs?
| methods[new_method].append(name) | ||
|
|
||
| ub_cfg[name] = final_cfg |
| for name in methods["ring_exchange"] + methods["pipeline"] + methods["bulk"]: | ||
| ub_cfg = get_default_config(name) | ||
| if ub_cfgs is not None and name in ub_cfgs: |
There was a problem hiding this comment.
if ub_cfgs is not None and name in ub_cfgs ==> if name in ub_cfgs
| fp8_buf = (name in layers_all_gather_overlap) or ( | ||
| final_cfg = get_default_config(name) | ||
| final_cfg.update(ub_cfgs[name]) | ||
| final_cfg["fp8_buf"] = (name in layers_all_gather_overlap) or ( |
There was a problem hiding this comment.
if I am using 'bf16' dtype for training, and name in layers_all_gather_overlap is true for some pattern (e.g., fc2_dgrad), then final_cfg["fp8_buf"] will be set to True, which is unexpected.
There was a problem hiding this comment.
Okay line 361 has already using use_fp8 to guard the effectiveness of fp8_buf, but it would be better if we can advance this guard when we set the default value of cfg['fp8_buf'] to avoid confusion.
There was a problem hiding this comment.
We can't advance this guard because it's not only for all-gather overlaps. I'll run through the logic to explain.
use_fp8 tells the initialization whether the user intends to invoke comm+GEMM overlap under the te.fp8_autocast() context (i.e. with FP8 GEMM inputs).
fp8_buf tells the initialization whether the communication buffer should be allocated as FP8.
For all-gather overlaps, the buffer has to match the GEMM input type, so it will always be allocated in FP8 when use_fp8 == True (i.e. GEMM inputs are FP8) regardless of what fp8_buf is set to in the user's layer configuration. In other words, the user does not get a choice here.
For reduce-scatter overlaps, the GEMM output has to match the buffer type, which can be either FP8 or BF16 when GEMM inputs are FP8. In this scenario, setting fp8_buf = True means that we communicate FP8 data between devices, and then fuse the BF16 upcast into the sum-reduce.
Advancing this guard to the default config options means that the user is denied the option to set fp8_buf for reduce-scatter overlaps, and that RS overlaps always communicate BF16 data, which is not always the optimal choice.
On a side note, the name in method["pipeline"] part of this is oudated and needs to be removed because we support optional FP8 GEMM outputs/buffers in all reduce-scatter overlaps now, not just collective/pipeline methods.
denera
force-pushed
the
rs-dgrad-overlap-bugfix
branch
from
aad8294 to
d2d9938
Compare
…dgrad overlap is enabled Signed-off-by: Alp Dener <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Tim Moon <[email protected]>
Signed-off-by: Alp Dener <[email protected]>
…er logic Signed-off-by: Alp Dener <[email protected]>
Signed-off-by: Alp Dener <[email protected]>
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Signed-off-by: Alp Dener <[email protected]>
denera
force-pushed
the
rs-dgrad-overlap-bugfix
branch
from
1036168 to
1797f14
Compare
Description
When Userbuffers config dictionary sets overlap method to
ring-exchangeorpipelinefor any*_dgradlayer, that layer's*_wgradoverlap needs to be disabled in order forub_overlap_rs_dgrad=Trueoption for related TE modules to function correctly.This PR fixes a bug where the "*_wgrad" overlap was persisting in the Userbuffer configuration and the corresponding UB object was being initialized even when it was not needed.
Type of change
Changes
Please list the changes introduced in this PR:
*_wgradoverlap is now removed frommethods["bulk"]list when the same layer's*_dgradoverlap has its method set to eitherring-exchangeorpipeline.add_ub(name, **ub_cfg)is now only called ifnameis in the original user-providedub_cfg. This avoids creating UB objects with default configs that may conflict with the user's intended TP overlap use.Checklist: