1/2 Last step cleaning the LlamaModel codebase

src/transformers/modeling_utils.py

@@ -33,7 +33,7 @@
 from enum import Enum
 from functools import partial, wraps
 from threading import Thread
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union
 from zipfile import is_zipfile
 
 import torch
@@ -73,14 +73,9 @@
     verify_tp_plan,
 )
 from .loss.loss_utils import LOSS_MAPPING
+from .modeling_layers import GradientCheckpointingLayer
 from .pytorch_utils import (  # noqa: F401
-    Conv1D,
-    apply_chunking_to_forward,
-    find_pruneable_heads_and_indices,
     id_tensor_storage,
-    prune_conv1d_layer,
-    prune_layer,
-    prune_linear_layer,
 )
 from .quantizers import AutoHfQuantizer, HfQuantizer
 from .quantizers.quantizers_utils import get_module_from_name
@@ -2001,6 +1996,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, PushToHubMixin, PeftAdapterMi
     # In practice, it means that they support attention interface functions, fully pass the kwargs
     # through all modules up to the Attention layer, can slice logits with Tensor, and have a default TP plan
     _supports_attention_backend = False
+    _can_record_outputs = None
 
     @property
     def dummy_inputs(self) -> dict[str, torch.Tensor]:
@@ -2050,6 +2046,11 @@ def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
         self._keep_in_fp32_modules = copy.copy(self.__class__._keep_in_fp32_modules)
 
         self._no_split_modules = self._no_split_modules or []
+        _param_to_record = {}
+        for module in self.modules():
+            if hasattr(module, "return_hooks"):
+                _param_to_record.update({module.return_hooks[0]: (module, module.return_hooks[1])})
+        self._can_record_outputs: Dict[str, Tuple[nn.Module, int]] = _param_to_record
 
     def post_init(self):
         """