diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3a0ad7dd3f..eb3bbc145c 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -146,6 +146,80 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def test_operation_fuser_caches_plans_by_grad_requirement(monkeypatch) -> None: + """Cache and restore fusion plans for checkpoint forward and recompute.""" + + # Count fusion-plan construction without depending on any particular real + # fusion implementation. Each distinct fusion configuration invokes this + # hook once, while a cache hit must bypass it entirely. + fusion_calls = 0 + + def track_fusion(ops, *, recipe): # pylint: disable=unused-argument + nonlocal fusion_calls + fusion_calls += 1 + # Preserve the operation list so this hook observes plan construction + # without changing the topology under test. + return ops + + # The fusion registries are class attributes shared by every OperationFuser. + # pytest's monkeypatch fixture restores all three after the test, preventing + # this synthetic fusion function from leaking into other tests. Keep only a + # joint forward-backward fusion hook so each plan build has one countable + # callback and no registered TE fusion can affect the result. + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + # One Identity op is enough to exercise the cache. With one basic op, + # first_op_requiring_backward has an intentionally simple interpretation: + # 0: backward starts at the Identity op; + # 1: the boundary is past the only op, so no backward work is required. + fuser = OperationFuser([te_ops.Identity()]) + x = torch.ones(1, requires_grad=True) + # maybe_fuse_ops expects one extra-input collection per basic op. Identity + # has no extra inputs, so its collection is an empty tuple. + extra_inputs = [()] + + # Phase 1: the original checkpointed forward runs with grad disabled. This + # is the first invocation, so the fuser must construct and cache the no-grad + # configuration. The runtime backward boundary is past the only op. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 1 + no_grad_forward_ops = fuser._forward_ops + no_grad_backward_ops = fuser._backward_ops + + # Phase 2: backward replays the checkpointed region with grad enabled. The + # backward boundary is part of the fusion key, allowing future fusion rules + # to choose a training-specific topology. The first grad-enabled invocation + # therefore constructs and caches a second configuration. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + grad_forward_ops = fuser._forward_ops + grad_backward_ops = fuser._backward_ops + assert grad_forward_ops is not no_grad_forward_ops + assert grad_backward_ops is not no_grad_backward_ops + + # Phase 3: the next checkpointed forward must select the exact no-grad lists + # cached in phase 1. Before the cache was added, every boundary transition + # rebuilt the fused operations and called track_fusion again. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 1 + assert fuser._forward_ops is no_grad_forward_ops + assert fuser._backward_ops is no_grad_backward_ops + + # Phase 4: another recomputation must likewise restore the grad-enabled + # lists from phase 2. The full alternating sequence has built only the two + # configurations represented by its two fusion keys. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + assert fuser._forward_ops is grad_forward_ops + assert fuser._backward_ops is grad_backward_ops + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 6b3f53fbd9..bb625aa09e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1312,7 +1312,7 @@ def fuser_forward( alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) @@ -1979,7 +1979,7 @@ def fuser_backward( # Kernel scaling factors alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) scales_f32 = None diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..e684384dc8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -48,6 +48,8 @@ def _is_graph_capturing() -> bool: OperationFusionFunction: TypeAlias = ( "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]" ) +_FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]] +_FusionParams: TypeAlias = tuple[type, int, Optional[str]] class _OperationFuserAutogradFunction(torch.autograd.Function): @@ -535,15 +537,22 @@ def __init__( op._lock_extra_tensor_channels() # Ops for forward and backward pass, will be populated in maybe_fuse_ops - self._forward_ops: list[tuple[FusibleOperation, list[int]]] - self._backward_ops: list[tuple[FusibleOperation, list[int]]] + self._forward_ops: _FusedOpList + self._backward_ops: _FusedOpList + + # Fused operation configurations are reusable wrappers around the basic + # ops, so cache each configuration by the state that selected it. + self._fused_ops_cache: dict[_FusionParams, tuple[_FusedOpList, _FusedOpList]] = {} # Cache and detect change of state relevant for fusing operations self.recipe_type = None - self.first_op_requiring_backward = 0 self.backward_override = None self._last_amax_history_len = 0 + # Runtime backward boundary. Full activation recompute alternates this + # between the checkpointed forward and the grad-enabled recomputation. + self.first_op_requiring_backward = 0 + # Flatten list of parameters self._basic_op_params = [list(op.parameters()) for op in self._basic_ops] self._basic_op_num_params = list(map(len, self._basic_op_params)) @@ -626,36 +635,52 @@ def maybe_fuse_ops( first_op_requiring_backward = op_idx break - # Early exit if fusion parameters haven't changed - need_reset = False + # Update the runtime backward boundary on every invocation, including + # paths that reuse a cached fused operation configuration. + self.first_op_requiring_backward = first_op_requiring_backward + + # Recipe state belongs to the basic ops and is independent of the + # runtime backward boundary. Only reconfigure it when recipe settings + # that are relevant to operation fusion change. recipe_type = type(recipe) backward_override = recipe.backward_override if recipe is not None else None - fusion_params = (recipe_type, first_op_requiring_backward, backward_override) - if fusion_params != ( + recipe_params = (recipe_type, backward_override) + need_reset = recipe_params != ( self.recipe_type, - self.first_op_requiring_backward, self.backward_override, - ): - # Recipe type, backward override, or grad requirements have changed - need_reset = True - elif ( - recipe is not None + ) + if ( + not need_reset + and recipe is not None and recipe.delayed() and self._last_amax_history_len != recipe.amax_history_len ): - # FP8 delayed scaling has changed amax history length + # Delayed-scaling history affects quantizer state, but it does not + # select a different fused operation configuration. need_reset = True - if not need_reset: - return + if need_reset: + for op in self._basic_ops: + op.reset_recipe_state(recipe=recipe) - # Reset recipe state - for op in self._basic_ops: - op.reset_recipe_state(recipe=recipe) + # Check if this is the first iteration + if self.recipe_type is None: + for op in self._basic_ops: + op.pre_first_fuser_forward() - # Check if this is the first iteration - if self.recipe_type is None: - for op in self._basic_ops: - op.pre_first_fuser_forward() + self.recipe_type, self.backward_override = recipe_params + self._last_amax_history_len = ( + recipe.amax_history_len if isinstance(recipe, DelayedScaling) else 0 + ) + + # Training and inference may support different fusions. Keep the + # backward boundary in the key, but pay construction cost only once for + # each configuration. Full recompute therefore builds at most one + # no-grad plan and one grad-enabled plan for a stable recipe. + fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + cached_ops = self._fused_ops_cache.get(fusion_params) + if cached_ops is not None: + self._forward_ops, self._backward_ops = cached_ops + return # Apply joint forward-backward fusions first joint_ops = OperationFuser._apply_fusions( @@ -682,14 +707,9 @@ def maybe_fuse_ops( self._basic_ops, ) - # Save current fusion params - self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params - - # Save amax history length - if isinstance(recipe, DelayedScaling): - self._last_amax_history_len = recipe.amax_history_len - else: - self._last_amax_history_len = 0 + # The FusedOperation contract excludes parameters and per-invocation + # state, so the mapped lists can be selected directly on cache hits. + self._fused_ops_cache[fusion_params] = (self._forward_ops, self._backward_ops) def __call__( self,