run model debugging with forward arg (#39905)

* run model debugging a lot simpler

* fixup

* Update src/transformers/utils/generic.py

* fixup

* mode syle?

* guard a bit

Author: ArthurZucker

src/transformers/model_debugging_utils.py

@@ -21,26 +21,23 @@
 from io import StringIO
 from typing import Optional
 
-from safetensors.torch import save_file
-
-from transformers.utils.import_utils import requires
-
-from .utils import is_torch_available
+from .utils.import_utils import is_torch_available, requires
 
 
 if is_torch_available():
     import torch
     import torch.distributed.tensor
+    from safetensors.torch import save_file
 
-
+    # Note to code inspectors: this toolbox is intended for people who add models to `transformers`.
+    _torch_distributed_available = torch.distributed.is_available()
+else:
+    _torch_distributed_available = False
 from .utils import logging
 
 
 logger = logging.get_logger(__name__)
 
-# Note to code inspectors: this toolbox is intended for people who add models to `transformers`.
-_torch_distributed_available = torch.distributed.is_available()
-
 
 def _is_rank_zero():
     """Return True if rank=0 or we aren't running distributed."""

src/transformers/utils/generic.py

@@ -52,6 +52,8 @@
     # required for @can_return_tuple decorator to work with torchdynamo
     import torch  # noqa: F401
 
+    from ..model_debugging_utils import model_addition_debugger_context
+
 
 class cached_property(property):
     """
@@ -1032,7 +1034,13 @@ def make_capture_wrapper(module, orig_forward, key, index):
             def wrapped_forward(*args, **kwargs):
                 if key == "hidden_states" and len(collected_outputs[key]) == 0:
                     collected_outputs[key] += (args[0],)
-                output = orig_forward(*args, **kwargs)
+                if kwargs.get("debug_io", False):
+                    with model_addition_debugger_context(
+                        module, kwargs.get("debug_io_dir", "~/model_debug"), kwargs.get("prune_layers")
+                    ):
+                        output = orig_forward(*args, **kwargs)
+                else:
+                    output = orig_forward(*args, **kwargs)
                 if not isinstance(output, tuple):
                     collected_outputs[key] += (output,)
                 elif output[index] is not None: