Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Dec 18, 2024
1 parent bb0a330 commit 744a96f
Showing 1 changed file with 7 additions and 8 deletions.
15 changes: 7 additions & 8 deletions transformer_engine/pytorch/module/linear.py
Original file line number Diff line number Diff line change
Expand Up @@ -927,15 +927,14 @@ def __init__(
assert ub_name is not None, f"Comm+GEMM overlap layer '{ub_name}' is not initialized."
self.ub_name = ub_name

assert not (self.ub_overlap_rs_fprop and self.ub_overlap_ag_fprop), (
"Cannot enable AG+GEMM and GEMM+RS overlaps at the same time."
)
assert not (self.ub_overlap_rs_dgrad and self.ub_bulk_dgrad), (
"Cannot enable DGRAD+RS and bulk DGRAD overlaps at the same time."
)
assert not (
self.ub_overlap_ag_dgrad
and (self.ub_overlap_rs_dgrad or self.ub_bulk_dgrad)
self.ub_overlap_rs_fprop and self.ub_overlap_ag_fprop
), "Cannot enable AG+GEMM and GEMM+RS overlaps at the same time."
assert not (
self.ub_overlap_rs_dgrad and self.ub_bulk_dgrad
), "Cannot enable DGRAD+RS and bulk DGRAD overlaps at the same time."
assert not (
self.ub_overlap_ag_dgrad and (self.ub_overlap_rs_dgrad or self.ub_bulk_dgrad)
), "Cannot enable AG+DGRAD and DGRAD+RS or bulk DGRAD overlaps at the same time."

self.get_rng_state_tracker = get_rng_state_tracker
Expand Down

0 comments on commit 744a96f

Please sign in to comment.