diff --git a/.github/workflows/integration_test_8gpu_graph_trainer.yaml b/.github/workflows/integration_test_8gpu_graph_trainer.yaml index e63a247291..d1acbe54fd 100644 --- a/.github/workflows/integration_test_8gpu_graph_trainer.yaml +++ b/.github/workflows/integration_test_8gpu_graph_trainer.yaml @@ -78,7 +78,8 @@ jobs: # Run precompile unit tests pytest torchtitan/experiments/graph_trainer/tests/test_precompile.py -v - # Run bitwise deterministic guardrail test + # Run bitwise deterministic and SAC peak-memory guardrail tests pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v + pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint diff --git a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml b/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml index 82dc48a9d5..b671e644eb 100644 --- a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml +++ b/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml @@ -68,7 +68,8 @@ jobs: # Run the MoE numerics tests NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "moe" - # Run bitwise deterministic guardrail test (includes H100-only hardcoded-hash tests) + # Run bitwise deterministic and SAC peak-memory guardrail tests pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v + pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*/checkpoint diff --git a/tests/integration_tests/run_tests.py b/tests/integration_tests/run_tests.py index de4a2415ac..7c5c2848d1 100644 --- a/tests/integration_tests/run_tests.py +++ b/tests/integration_tests/run_tests.py @@ -93,7 +93,6 @@ def run_single_test( def run_tests(args, test_list: list[OverrideDefinitions], module=None, config=None): """Run all integration tests to test the core features of TorchTitan""" - exclude_set = set() if hasattr(args, "exclude") and args.exclude: exclude_set = {name.strip() for name in args.exclude.split(",")} diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py index 55a58b1221..4ece5c0050 100644 --- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py +++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py @@ -12,9 +12,6 @@ import torch import torch.nn as nn import torch.utils._pytree as pytree -from torch._functorch._aot_autograd.logging_utils import ( - setup_stacktrace_preservation_hooks, -) from torch._guards import tracing, TracingContext from torch._subclasses import FakeTensorMode from torch.fx.experimental.proxy_tensor import make_fx @@ -200,52 +197,6 @@ def _remove_cpu_shadow_chains(gm: torch.fx.GraphModule) -> None: gm.recompile() -@contextmanager -def _patch_engine_run_backward() -> Generator[None, None, None]: - """Patch _engine_run_backward to install stacktrace preservation hooks. - - Why this is needed: - When make_fx traces a function that calls loss.backward(), the backward - pass is decomposed into primitive ATen ops. Normally (in eager autograd), - ``setup_stacktrace_preservation_hooks`` is called by the autograd engine - to propagate ``seq_nr`` from forward ops to their corresponding backward - ops. Under make_fx tracing, this hook setup doesn't happen automatically - because the engine path differs, so backward FX nodes end up without - ``seq_nr`` metadata. Without ``seq_nr``, we can't correlate backward - nodes back to their forward counterparts (needed by - ``_copy_fwd_metadata_to_bw_nodes``). - - This context manager patches ``_engine_run_backward`` to call - ``setup_stacktrace_preservation_hooks`` before the autograd engine runs, - restoring ``seq_nr`` propagation during tracing. - - We must patch the name in both modules since ``torch.autograd.__init__`` - imports it via ``from .graph import``. - """ - import torch.autograd - import torch.autograd.graph - - _orig_fn = torch.autograd.graph._engine_run_backward - - def _patched(t_outputs, *args, **kwargs): # type: ignore[no-untyped-def] - roots = [ - t.grad_fn - for t in t_outputs - if isinstance(t, torch.Tensor) and t.grad_fn is not None - ] - if roots: - setup_stacktrace_preservation_hooks(roots) - return _orig_fn(t_outputs, *args, **kwargs) - - torch.autograd.graph._engine_run_backward = _patched # type: ignore[assignment] - torch.autograd._engine_run_backward = _patched # type: ignore[assignment] - try: - yield - finally: - torch.autograd.graph._engine_run_backward = _orig_fn # type: ignore[assignment] - torch.autograd._engine_run_backward = _orig_fn # type: ignore[assignment] - - def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None: """Copy forward node metadata (custom) to later nodes sharing the same seq_nr. @@ -402,10 +353,8 @@ def fn_with_subclass_handling(*plain_args: Any) -> list: state_for_fn = dict(zip(state_fqns, state_wrapped, strict=True)) user_list = pytree.tree_unflatten(list(user_flat), user_args_spec) - - with _patch_engine_run_backward(): + with torch.compiler._patch_autograd_grad(): result = fn(state_for_fn, *user_list) - flat_outs, output_spec = pytree.tree_flatten(result) num_flat_outputs = len(flat_outs) unwrapped_outs, output_layouts = _unwrap_subclasses(flat_outs) @@ -413,9 +362,8 @@ def fn_with_subclass_handling(*plain_args: Any) -> list: ctx = TracingContext(fake_mode) # preserve_node_meta propagates fx.traceback.annotate metadata to traced nodes - # Disable autograd multithreading so that backward tracing - # runs on the calling thread. Without this, the C++ autograd + # runs on the calling thread. Without this, the C++ autograd # engine dispatches backward to a worker thread that has a # fresh contextvars.Context, making the compile_on_one_rank # ContextVar invisible and causing _sym_get_coordinate to @@ -423,12 +371,18 @@ def fn_with_subclass_handling(*plain_args: Any) -> list: # TODO: Move set_multithreading_enabled(False) to global init. # Forcing backward onto the main CPU thread is a good default # for both tracing and runtime, not just the tracing path. + # _skip_nested_compile lets the current make_fx trace inline through + # torch.compile'd FlexAttention kernels instead of erroring. + # _non_strict_tracing_context is required by _patch_autograd_grad() and + # marks this make_fx pass as the non-strict tracing flow, distinct from + # other make_fx-based entry points such as non-strict export. with ( fake_mode, tracing(ctx), preserve_node_meta(), _skip_nested_compile(), torch.autograd.set_multithreading_enabled(False), + torch.compiler._non_strict_tracing_context(), ): traced = make_fx( fn_with_subclass_handling, @@ -470,6 +424,10 @@ def run_traced( This is a reference implementation of traced-graph execution. It keeps the state handling, subclass unwrapping, and output reconstruction logic explicit instead of baking those semantics into ``TracedResult`` itself. + Runs under ``torch.no_grad()`` because the graph already contains explicit + backward ops (from ``torch.autograd.grad`` traced by make_fx). Without + this, PyTorch would build a redundant autograd graph on top, keeping all + forward intermediates alive via ``grad_fn`` references. """ state_flat = list(state.values()) user_args_flat, _ = pytree.tree_flatten(list(args)) @@ -481,7 +439,8 @@ def run_traced( all_args = list(state_flat) + list(user_args_flat) flat_inputs, _ = _unwrap_subclasses(all_args) - flat_outputs = traced_result.gm(*flat_inputs) + with torch.no_grad(): + flat_outputs = traced_result.gm(*flat_inputs) wrapped = _wrap_subclasses( flat_outputs, traced_result.num_flat_outputs, diff --git a/torchtitan/experiments/graph_trainer/passes.py b/torchtitan/experiments/graph_trainer/passes.py index 0bdb971c5f..f47c4e634d 100644 --- a/torchtitan/experiments/graph_trainer/passes.py +++ b/torchtitan/experiments/graph_trainer/passes.py @@ -50,6 +50,10 @@ from torchtitan.tools.logging import logger +def _is_backward_node(node: torch.fx.Node) -> bool: + return node.meta.get("autograd_backward", False) + + def compile_time_passes( traced_result: "TracedResult", ) -> list[Callable]: @@ -69,6 +73,7 @@ def compile_time_passes( remove_detach_pass, remove_identity_view_pass, remove_identity_slice_pass, + selective_activation_remat_pass, # FlexAttention HOPs must be compiled (via regional_inductor) to # produce bitwise identical results to the eager Trainer path. # When left uncompiled, flex_attention still runs correctly but @@ -452,6 +457,11 @@ def apply_sac_pass( if node.op != "call_function": continue + # Skip backward nodes — they must not carry recompute tags, + # otherwise the remat pass would try to duplicate backward ops. + if _is_backward_node(node): + continue + if node.target in ( operator.getitem, torch.ops._c10d_functional.wait_tensor.default, @@ -469,7 +479,6 @@ def apply_sac_pass( node.meta["recompute"] = parent.meta["recompute"] node.meta["ac_graph_id"] = parent.meta.get("ac_graph_id", 0) continue - custom_meta = node.meta.get("custom", {}) ac_region_id = custom_meta.get(_AC_REGION_ID, 0) node.meta["ac_graph_id"] = ac_region_id @@ -504,6 +513,27 @@ def apply_sac_pass( return gm +def selective_activation_remat_pass( + gm: torch.fx.GraphModule, example_inputs: tuple | None = None +) -> torch.fx.GraphModule: + """Apply graph-based SAC to a traced fwd+loss+bwd graph. + + Tags forward nodes with recompute policy via apply_sac_pass (backward + nodes are skipped automatically via ``node.meta["autograd_backward"]``), then + applies remat_using_tags_for_fwd_loss_bwd_graph to duplicate + PREFER_RECOMPUTE forward ops before backward and DCE originals. + + The model must have been annotated with annotate_ac_regions before + tracing so that nodes have custom["ac_region_id"] metadata. + """ + from torch._functorch._activation_checkpointing.remat_using_tags_for_fwd_loss_bwd_graph_pass import ( + remat_using_tags_for_fwd_loss_bwd_graph, + ) + + apply_sac_pass(gm) + return remat_using_tags_for_fwd_loss_bwd_graph(gm) + + # Apply activation checkpointing on joint graph before partitioner def fsdp_reshard_after_fwd_pass( gm: torch.fx.GraphModule, @@ -604,7 +634,7 @@ def inductor_decomposition_pass( f"Placeholder count mismatch: {len(orig_placeholders)} vs {len(decomp_placeholders)}" ) - for orig, decomp in zip(orig_placeholders, decomp_placeholders): + for orig, decomp in zip(orig_placeholders, decomp_placeholders, strict=True): # Copy all metadata from original to decomposed for key, value in orig.meta.items(): if key not in decomp.meta: diff --git a/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py b/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py new file mode 100644 index 0000000000..5c4dbec013 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/_trainer_test_utils.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from torchtitan.components.loss import cross_entropy_loss +from torchtitan.config import ActivationCheckpointConfig +from torchtitan.distributed.utils import get_train_context +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer +from torchtitan.trainer import Trainer + + +def build_minimal_trainer( + model: nn.Module, + model_config, + trainer_cls: type[Trainer], + *, + activation_checkpoint_mode: str = "none", + compile_enable_passes: bool = True, + compile_passes: list[str] | None = None, + compile_joint_passes: list[str] | None = None, + tokenizer=None, +) -> Trainer: + """Build the minimal Trainer/GraphTrainer needed for single-GPU test steps.""" + trainer = object.__new__(trainer_cls) + trainer.model_parts = [model] + trainer.loss_fn = cross_entropy_loss + trainer.parallel_dims = SimpleNamespace(pp_enabled=False, cp_enabled=False) + trainer.train_context = get_train_context(False) + trainer.model_config = model_config + trainer.device = torch.device("cuda") + trainer.tokenizer = tokenizer + + if trainer_cls is GraphTrainer: + trainer.config = SimpleNamespace( + compile=SimpleNamespace( + mode="aot_fx_trace", + enable_passes=compile_enable_passes, + passes=[] if compile_passes is None else list(compile_passes), + joint_passes=[] + if compile_joint_passes is None + else list(compile_joint_passes), + precompile_artifact_dir="", + ), + activation_checkpoint=ActivationCheckpointConfig( + mode=activation_checkpoint_mode + ), + ) + trainer._fwd_bwd_step_module = None + trainer._traced_step = None + else: + trainer.config = SimpleNamespace() + + return trainer diff --git a/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py b/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py index 5c0ba400d6..252fec7f36 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py +++ b/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py @@ -17,7 +17,6 @@ import tempfile import unittest from collections.abc import Callable -from types import SimpleNamespace import torch import torch.nn as nn @@ -26,7 +25,6 @@ from torchtitan.components.loss import cross_entropy_loss from torchtitan.components.tokenizer import HuggingFaceTokenizer -from torchtitan.distributed.utils import get_train_context from torchtitan.experiments.graph_trainer.deepseek_v3 import ( model_registry as dsv3_model_registry, ) @@ -37,6 +35,9 @@ model_registry as llama3_model_registry, ) from torchtitan.experiments.graph_trainer.llama3.parallelize import annotate_llama +from torchtitan.experiments.graph_trainer.tests._trainer_test_utils import ( + build_minimal_trainer, +) from torchtitan.experiments.graph_trainer.trainer import GraphTrainer from torchtitan.tools.utils import has_cuda_capability from torchtitan.trainer import Trainer @@ -58,43 +59,6 @@ def _set_deterministic(seed: int = SEED) -> None: _TOKENIZER_PATH = "./tests/assets/tokenizer" -def _build_trainer( - model: nn.Module, - model_config, - trainer_cls: type, - *, - enable_passes: bool = True, -) -> Trainer: - """Build a minimal Trainer/GraphTrainer for single-GPU non-distributed testing. - - Uses object.__new__ to bypass __init__ because the full Trainer constructor - requires a distributed environment, job config, and checkpoint manager that - are unnecessary for single-GPU numerical verification. The attributes set - below are the minimal set required by forward_backward_step(). - """ - trainer = object.__new__(trainer_cls) - trainer.model_parts = [model] - trainer.loss_fn = cross_entropy_loss - trainer.parallel_dims = SimpleNamespace(pp_enabled=False, cp_enabled=False) - trainer.train_context = get_train_context(False) - trainer.model_config = model_config - trainer.device = torch.device("cuda") - trainer.tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH) - - if trainer_cls is GraphTrainer: - trainer.config = SimpleNamespace( - compile=SimpleNamespace( - mode="aot_fx_trace", - enable_passes=enable_passes, - precompile_artifact_dir="", - ) - ) - trainer._fwd_bwd_step_module = None - trainer._traced_step = None - - return trainer - - @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") class BitwiseDeterministicBase(unittest.TestCase): """Base class for bitwise determinism tests. @@ -133,8 +97,12 @@ def _run_steps( # Annotate after deepcopy: annotate_fn wrappers capture bound methods # that don't rebind correctly through copy.deepcopy. self.annotate_model(model) - trainer = _build_trainer( - model, self.model_config, trainer_cls, enable_passes=enable_passes + trainer = build_minimal_trainer( + model, + self.model_config, + trainer_cls, + compile_enable_passes=enable_passes, + tokenizer=HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH), ) global_valid_tokens = torch.tensor( BATCH_SIZE * SEQ_LEN, dtype=torch.float, device="cuda" @@ -436,11 +404,6 @@ def test_eager_self_deterministic(self): """8bb6e647c3edaa229cc65872086ccc5c4e1b7f1647bb01da4506ab777a64a0db""", ) - # TODO: OOMs during flex_attention compilation on A100 GPUs. - # Revisit when GraphTrainer addresses peak memory during compilation. - @unittest.skipUnless( - has_cuda_capability(9, 0), "OOMs during flex_attention compilation on A100" - ) def test_aot_fx_trace_vs_eager(self): """aot_fx_trace with passes and eager produce bitwise identical results.""" run_eager = self._run_steps(copy.deepcopy(self.model), Trainer) diff --git a/torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py b/torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py new file mode 100644 index 0000000000..fac0e02749 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py @@ -0,0 +1,161 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import unittest +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from torchtitan.config import ActivationCheckpointConfig +from torchtitan.distributed.activation_checkpoint import apply_ac +from torchtitan.experiments.graph_trainer.llama3 import ( + model_registry as llama3_registry, +) +from torchtitan.experiments.graph_trainer.tests._trainer_test_utils import ( + build_minimal_trainer, +) +from torchtitan.experiments.graph_trainer.trainer import GraphTrainer +from torchtitan.trainer import Trainer + +DTYPE = torch.bfloat16 +BATCH_SIZE = 2 +SEQ_LEN = 2048 +MAX_PEAK_MEMORY_RATIO = 1.10 +DEBUGMODEL = "debugmodel" + + +def _set_deterministic() -> None: + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + torch.use_deterministic_algorithms(True) + + +def _build_model(model_flavor: str) -> nn.Module: + model_spec = llama3_registry(model_flavor) + with torch.device("meta"): + model = model_spec.model.build() + model.to_empty(device="cuda") + with torch.no_grad(): + model.init_states(buffer_device=None) + model.to(dtype=DTYPE) + model.train() + return model + + +@dataclass(frozen=True) +class StepResult: + loss: torch.Tensor + grads: list[torch.Tensor] + reserved_gib: float + active_gib: float + + +def _measure_step( + trainer: Trainer, tokens: torch.Tensor, labels: torch.Tensor +) -> StepResult: + model = trainer.model_parts[0] + model.zero_grad(set_to_none=True) + global_valid_tokens = torch.tensor(labels.numel(), dtype=torch.float, device="cuda") + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + loss = trainer.forward_backward_step( + input_dict={"input": tokens}, + labels=labels, + global_valid_tokens=global_valid_tokens, + ) + torch.cuda.synchronize() + + stats = torch.cuda.memory_stats() + grads = [param.grad.detach().clone() for param in model.parameters()] + return StepResult( + loss=loss.detach().clone(), + grads=grads, + reserved_gib=torch.cuda.max_memory_reserved() / 1e9, + active_gib=stats["active_bytes.all.peak"] / 1e9, + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class TestGraphSACPeakMemory(unittest.TestCase): + def setUp(self): + _set_deterministic() + model = _build_model(DEBUGMODEL) + self.state_dict = { + key: value.detach().cpu().clone() + for key, value in model.state_dict().items() + } + del model + torch.cuda.empty_cache() + self.tokens = torch.randint(0, 2048, (BATCH_SIZE, SEQ_LEN), device="cuda") + self.labels = torch.randint(0, 2048, (BATCH_SIZE, SEQ_LEN), device="cuda") + + def tearDown(self): + torch.use_deterministic_algorithms(False) + + def test_llama3_debugmodel_peak_memory_matches_eager_selective_ac(self): + eager_model = _build_model(DEBUGMODEL) + eager_model.load_state_dict(copy.deepcopy(self.state_dict)) + apply_ac(eager_model, ActivationCheckpointConfig(mode="selective")) + eager_trainer = build_minimal_trainer( + eager_model, + llama3_registry(DEBUGMODEL).model, + Trainer, + ) + + traced_model = _build_model(DEBUGMODEL) + traced_model.load_state_dict(copy.deepcopy(self.state_dict)) + traced_trainer = build_minimal_trainer( + traced_model, + llama3_registry(DEBUGMODEL).model, + GraphTrainer, + activation_checkpoint_mode="selective", + ) + + # Warm up both paths so allocator and one-time tracing setup do not skew + # the measured peak memory. + _measure_step(eager_trainer, self.tokens, self.labels) + _measure_step(traced_trainer, self.tokens, self.labels) + torch.cuda.empty_cache() + + eager = _measure_step(eager_trainer, self.tokens, self.labels) + traced = _measure_step(traced_trainer, self.tokens, self.labels) + + self.assertTrue( + torch.equal(eager.loss, traced.loss), + f"loss mismatch: eager={eager.loss.item()} traced={traced.loss.item()}", + ) + for idx, (eager_grad, traced_grad) in enumerate( + zip(eager.grads, traced.grads, strict=True) + ): + self.assertTrue( + torch.equal(eager_grad, traced_grad), f"grad[{idx}] mismatch" + ) + + reserved_ratio = traced.reserved_gib / eager.reserved_gib + active_ratio = traced.active_gib / eager.active_gib + self.assertLessEqual( + reserved_ratio, + MAX_PEAK_MEMORY_RATIO, + "graph SAC reserved peak memory too high: " + f"traced={traced.reserved_gib:.3f} GiB, " + f"eager={eager.reserved_gib:.3f} GiB, " + f"ratio={reserved_ratio:.3f}", + ) + self.assertLessEqual( + active_ratio, + MAX_PEAK_MEMORY_RATIO, + "graph SAC active peak memory too high: " + f"traced={traced.active_gib:.3f} GiB, " + f"eager={eager.active_gib:.3f} GiB, " + f"ratio={active_ratio:.3f}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py index 8f28ab9d3a..246812f0b0 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_trace_module.py +++ b/torchtitan/experiments/graph_trainer/tests/test_trace_module.py @@ -16,7 +16,6 @@ ) from torchtitan.experiments.graph_trainer.make_fx_tracer import ( _copy_fwd_metadata_to_bw_nodes, - _patch_engine_run_backward, extract_module_state, minimal_fx_tracer, run_traced, @@ -391,7 +390,7 @@ def test_dtensor_train_step(self): @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestMetadataPropagation(unittest.TestCase): - """Tests for _patch_engine_run_backward and _copy_fwd_metadata_to_bw_nodes.""" + """Tests for _copy_fwd_metadata_to_bw_nodes.""" DEVICE = "cuda" DTYPE = torch.float32 @@ -400,7 +399,7 @@ def setUp(self): torch.manual_seed(42) def test_backward_nodes_have_seq_nr(self): - """Verify that backward FX nodes get seq_nr metadata via the patched engine.""" + """Verify that backward FX nodes get seq_nr metadata via patched autograd.grad.""" model = SimpleMLP().to(device=self.DEVICE, dtype=self.DTYPE) train_step = make_train_step(get_loss) tokens = torch.randint(0, 256, (2, 32), device=self.DEVICE) @@ -499,22 +498,6 @@ def test_backward_nodes_have_stack_trace(self): f"Backward nodes missing stack_trace: {bwd_nodes_missing_stack_trace}", ) - def test_patch_engine_restores_original(self): - """Verify that _patch_engine_run_backward restores the original function.""" - import torch.autograd - import torch.autograd.graph - - orig_fn = torch.autograd.graph._engine_run_backward - - with _patch_engine_run_backward(): - # Inside the context, it should be patched - self.assertIsNot(torch.autograd.graph._engine_run_backward, orig_fn) - self.assertIsNot(torch.autograd._engine_run_backward, orig_fn) - - # After the context, it should be restored - self.assertIs(torch.autograd.graph._engine_run_backward, orig_fn) - self.assertIs(torch.autograd._engine_run_backward, orig_fn) - @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestTraceModels(unittest.TestCase):