Skip to content

Precompute writeback dedup indices in forward to eliminate GPU-CPU sync in backward (#5522) #5522

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Zhihan-Lu wants to merge 1 commit into from
Closed

Precompute writeback dedup indices in forward to eliminate GPU-CPU sync in backward (#5522) #5522

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 26 additions & 9 deletions fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,10 @@
TBEStatsReporterConfig,
)
from fbgemm_gpu.utils.loader import load_torch_module, load_torch_module_bc
from fbgemm_gpu.utils.writeback_util import writeback_gradient
from fbgemm_gpu.utils.writeback_util import (
compute_writeback_indices_dispatch,
writeback_gradient,
)

try:
load_torch_module(
Expand Down Expand Up @@ -158,6 +161,7 @@ class UserEnabledConfigDefinition:
use_rowwise_bias_correction: bool = False
use_writeback_bwd_prehook: bool = False
writeback_first_feature_only: bool = False
precompute_writeback: bool = False


@dataclass(frozen=True)
Expand Down Expand Up @@ -1182,9 +1186,6 @@ def __init__( # noqa C901
extra_optimizer_config.use_writeback_bwd_prehook
)

writeback_first_feature_only: bool = (
extra_optimizer_config.writeback_first_feature_only
)
self.log(f"self.extra_optimizer_config is {extra_optimizer_config}")
if self.use_rowwise_bias_correction and not self.optimizer == OptimType.ADAM:
raise AssertionError(
Expand Down Expand Up @@ -1473,7 +1474,11 @@ def __init__( # noqa C901
# self.log("TBE_V2 Knob is set to True; Using experimental TBE")

self.is_experimental: bool = is_experimental
self._writeback_first_feature_only: bool = writeback_first_feature_only
self._writeback_first_feature_only: bool = (
extra_optimizer_config.writeback_first_feature_only
)
self._writeback_precomputed_index: Optional[Tensor] = None
self._precompute_writeback: bool = extra_optimizer_config.precompute_writeback

# Get a debug function pointer
self._debug_print_input_stats: Callable[..., None] = (
Expand Down Expand Up @@ -2218,15 +2223,14 @@ def _feature_is_enabled(self, feature: FeatureGateName) -> bool:

# pyre-fixme[2]: For 1st argument expected not ANY
def writeback_hook(self, module: Any, grad: tuple[Tensor]) -> tuple[Tensor]:
indices = self._indices
offsets = self._offsets
return writeback_gradient(
grad,
indices,
offsets,
self._indices,
self._offsets,
self.feature_table_map,
self._writeback_first_feature_only,
self.is_nobag,
self._writeback_precomputed_index,
)

def forward( # noqa: C901
Expand Down Expand Up @@ -2416,6 +2420,19 @@ def forward( # noqa: C901
# Storing tensors for linear_cache_indices recomputation
self._indices = indices
self._offsets = offsets
# Precompute writeback dedup indices (training-only, not TorchScript-compatible).
if (
not torch.jit.is_scripting()
and self.use_writeback_bwd_prehook
and self._precompute_writeback
):
self._writeback_precomputed_index = compute_writeback_indices_dispatch(
indices,
offsets,
self.feature_table_map,
self._writeback_first_feature_only,
self.is_nobag,
)
self._vbe_B_offsets = vbe_metadata.B_offsets
self._vbe_max_B = vbe_metadata.max_B

Expand Down
Loading
Loading