transformer_engine/pytorch/module/base.py · 1258bbe022a54f14256123ed51c5e5b648dc0f4c · OpenDAS / TransformerEngine

"...git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "e5452ddfd6e9a08d5e15bd81a010934550b9b507"

Find file Blame History Permalink

Manually launch wgrad accumulation and reduce in backward_dw() instead of backward() (#1976) · 1258bbe0

Hongbin Liu authored Aug 01, 2025



* disable wgrad accumulation and reduce in backward() And manually launch it in backward_dw()
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* format
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* refactor
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* refactor
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* set skip_backward_post_hook to True only if delay_wgrad_compute is True
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* format
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

---------
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

1258bbe0

base.py 63.8 KB

Replace base.py