Repository: huggingface/transformers
Pull request: [Trainer] Move NEFTune impl to standalone functions

# What does this PR do ?

This PR refactor NEFTune implementation from Trainer methods to standalone functions. Part of PRs to unbloat Trainer. 

Changed files:
filename: src/transformers/integrations/__init__.py
patch: @@ -130,6 +130,11 @@
         "replace_with_mxfp4_linear",
         "swizzle_mxfp4",
     ],
+    "neftune": [
+        "activate_neftune",
+        "deactivate_neftune",
+        "neftune_post_forward_hook",
+    ],
     "peft": ["PeftAdapterMixin"],
     "quanto": ["replace_with_quanto_layers"],
     "spqr": ["replace_with_spqr_linear"],
@@ -270,6 +275,7 @@
         replace_with_mxfp4_linear,
         swizzle_mxfp4,
     )
+    from .neftune import activate_neftune, deactivate_neftune, neftune_post_forward_hook
     from .peft import PeftAdapterMixin
     from .quanto import replace_with_quanto_layers
     from .spqr import replace_with_spqr_linear

filename: src/transformers/integrations/neftune.py
patch: @@ -0,0 +1,114 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+NEFTune: Noisy Embeddings for Fine-Tuning.
+
+Implementation based on https://github.com/neelsjain/NEFTune
+Paper: https://huggingface.co/papers/2310.05914
+"""
+
+import torch
+
+from ..trainer_utils import _is_peft_model
+
+
+def neftune_post_forward_hook(module, input, output):
+    """
+    Implements the NEFTune forward pass for the model using forward hooks. Note this works only for torch.nn.Embedding
+    layers. This method is slightly adapted from the original source code that can be found here:
+    https://github.com/neelsjain/NEFTune. Simply add it to your model as follows:
+    ```python
+    from transformers.integrations.neftune import neftune_post_forward_hook
+
+    model = ...
+    model.embed_tokens.neftune_noise_alpha = 0.1
+    model.embed_tokens.register_forward_hook(neftune_post_forward_hook)
+    ```
+    Args:
+        module (`torch.nn.Module`):
+            The embedding module where the hook is attached. Note that you need to set `module.neftune_noise_alpha` to
+            the desired noise alpha value.
+        input (`torch.Tensor`):
+            The input tensor to the model.
+        output (`torch.Tensor`):
+            The output tensor of the model (i.e. the embeddings).
+    """
+    if module.training:
+        dims = torch.tensor(output.size(1) * output.size(2))
+        mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)
+        output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
+    return output
+
+
+def activate_neftune(model, neftune_noise_alpha, accelerator=None):
+    """
+    Activates NEFTune (Noisy Embeddings for Fine-Tuning) on the model.
+
+    NEFTune adds noise to embedding vectors during training, which has been shown to improve
+    fine-tuning performance. See https://huggingface.co/papers/2310.05914 for details.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to activate NEFTune on.
+        neftune_noise_alpha (`float`):
+            The noise alpha value controlling the magnitude of the noise.
+        accelerator (`Accelerator`, *optional*):
+            The accelerator instance. If provided, the model will be unwrapped before
+            accessing embeddings. Required when using distributed training.
+
+    Returns:
+        `torch.utils.hooks.RemovableHandle`: The hook handle that can be used to deactivate NEFTune.
+    """
+    if accelerator is not None:
+        unwrapped_model = accelerator.unwrap_model(model)
+    else:
+        unwrapped_model = model
+
+    if _is_peft_model(unwrapped_model):
+        embeddings = unwrapped_model.base_model.model.get_input_embeddings()
+    else:
+        embeddings = unwrapped_model.get_input_embeddings()
+
+    embeddings.neftune_noise_alpha = neftune_noise_alpha
+    hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook)
+
+    return hook_handle
+
+
+def deactivate_neftune(model, hook_handle, accelerator=None):
+    """
+    Deactivates NEFTune on the model.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to deactivate NEFTune on.
+        hook_handle (`torch.utils.hooks.RemovableHandle`):
+            The hook handle returned by `activate_neftune`.
+        accelerator (`Accelerator`, *optional*):
+            The accelerator instance. If provided, the model will be unwrapped before
+            accessing embeddings.
+    """
+    if accelerator is not None:
+        unwrapped_model = accelerator.unwrap_model(model)
+    else:
+        unwrapped_model = model
+
+    if _is_peft_model(unwrapped_model):
+        embeddings = unwrapped_model.base_model.model.get_input_embeddings()
+    else:
+        embeddings = unwrapped_model.get_input_embeddings()
+
+    hook_handle.remove()
+    if hasattr(embeddings, "neftune_noise_alpha"):
+        del embeddings.neftune_noise_alpha

filename: src/transformers/trainer.py
patch: @@ -62,6 +62,7 @@
 from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend
 from .image_processing_utils import BaseImageProcessor
 from .integrations.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_available
+from .integrations.neftune import activate_neftune, deactivate_neftune
 from .integrations.peft import MIN_PEFT_VERSION
 from .integrations.tpu import tpu_spmd_dataloader
 from .modelcard import TrainingSummary
@@ -113,6 +114,7 @@
     SaveStrategy,
     TrainerMemoryTracker,
     TrainOutput,
+    _is_peft_model,
     check_target_module_exists,
     default_compute_objective,
     denumpify_detensorize,
@@ -121,7 +123,6 @@
     get_last_checkpoint,
     has_length,
     load_sharded_checkpoint,
-    neftune_post_forward_hook,
     number_of_arguments,
     rotate_checkpoints,
     seed_worker,
@@ -204,7 +205,7 @@
     from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat
 
 if is_peft_available():
-    from peft import PeftMixedModel, PeftModel
+    from peft import PeftModel
 
 if is_accelerate_available():
     from accelerate import Accelerator, skip_first_batches
@@ -225,13 +226,6 @@
         from accelerate.utils import DeepSpeedSchedulerWrapper
 
 
-def _is_peft_model(model):
-    if is_peft_available():
-        classes_to_check = (PeftModel, PeftMixedModel)
-        return isinstance(model, classes_to_check)
-    return False
-
-
 def _get_fsdp_ckpt_kwargs():
     if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters):
         return {"adapter_only": True}
@@ -763,42 +757,6 @@ def __init__(
             xs.set_global_mesh(xs.Mesh(np.array(range(num_devices)), (num_devices, 1), axis_names=("fsdp", "tensor")))
         self.is_fsdp_xla_v1_enabled = self.is_fsdp_xla_enabled and not self.is_fsdp_xla_v2_enabled
 
-    def _activate_neftune(self, model):
-        r"""
-        Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper:
-        https://huggingface.co/papers/2310.05914
-        """
-        unwrapped_model = self.accelerator.unwrap_model(model)
-
-        if _is_peft_model(unwrapped_model):
-            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
-        else:
-            embeddings = unwrapped_model.get_input_embeddings()
-
-        del unwrapped_model
-
-        embeddings.neftune_noise_alpha = self.neftune_noise_alpha
-        hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook)
-        self.neftune_hook_handle = hook_handle
-        return model
-
-    def _deactivate_neftune(self, model):
-        """
-        Deactivates the neftune method. Make sure to call `_activate_neftune` first.
-        """
-        if not hasattr(self, "neftune_hook_handle"):
-            raise ValueError("Neftune is not activated make sure to call `trainer._activate_neftune()` first")
-
-        unwrapped_model = self.accelerator.unwrap_model(model)
-
-        if _is_peft_model(unwrapped_model):
-            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
-        else:
-            embeddings = unwrapped_model.get_input_embeddings()
-
-        self.neftune_hook_handle.remove()
-        del embeddings.neftune_noise_alpha, unwrapped_model
-
     def add_callback(self, callback):
         """
         Add a callback to the current list of [`~transformers.TrainerCallback`].
@@ -2106,7 +2064,7 @@ def train(
 
         # Attach NEFTune hooks if necessary
         if self.neftune_noise_alpha is not None:
-            self.model = self._activate_neftune(self.model)
+            self.neftune_hook_handle = activate_neftune(self.model, self.neftune_noise_alpha, self.accelerator)
 
         # do_train is not a reliable argument, as it might not be set and .train() still called, so
         # the following is a workaround:
@@ -2701,7 +2659,7 @@ def _inner_training_loop(
         # After training we make sure to retrieve back the original forward pass method
         # for the embedding layer by removing the forward post hook.
         if self.neftune_noise_alpha is not None:
-            self._deactivate_neftune(self.model)
+            deactivate_neftune(self.model, self.neftune_hook_handle, self.accelerator)
 
         return TrainOutput(self.state.global_step, train_loss, metrics)
 

filename: src/transformers/trainer_utils.py
patch: @@ -38,6 +38,7 @@
     WEIGHTS_INDEX_NAME,
     ExplicitEnum,
     check_torch_load_is_safe,
+    is_peft_available,
     is_psutil_available,
     is_torch_available,
     is_torch_cuda_available,
@@ -60,6 +61,15 @@
     import torch
     from safetensors.torch import load_file as safe_load_file
 
+if is_peft_available():
+    from peft import PeftMixedModel, PeftModel
+
+
+def _is_peft_model(model):
+    if is_peft_available():
+        return isinstance(model, (PeftModel, PeftMixedModel))
+    return False
+
 
 def seed_worker(worker_id: int, num_workers: int, rank: int):
     """
@@ -126,32 +136,6 @@ def set_seed(seed: int, deterministic: bool = False):
         torch.xpu.manual_seed_all(seed)
 
 
-def neftune_post_forward_hook(module, input, output):
-    """
-    Implements the NEFTune forward pass for the model using forward hooks. Note this works only for torch.nn.Embedding
-    layers. This method is slightly adapted from the original source code that can be found here:
-    https://github.com/neelsjain/NEFTune Simply add it to your model as follows:
-    ```python
-    model = ...
-    model.embed_tokens.neftune_noise_alpha = 0.1
-    model.embed_tokens.register_forward_hook(neftune_post_forward_hook)
-    ```
-    Args:
-        module (`torch.nn.Module`):
-            The embedding module where the hook is attached. Note that you need to set `module.neftune_noise_alpha` to
-            the desired noise alpha value.
-        input (`torch.Tensor`):
-            The input tensor to the model.
-        output (`torch.Tensor`):
-            The output tensor of the model (i.e. the embeddings).
-    """
-    if module.training:
-        dims = torch.tensor(output.size(1) * output.size(2))
-        mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)
-        output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
-    return output
-
-
 class EvalPrediction:
     """
     Evaluation output (always contains labels), to be used to compute metrics.

filename: tests/trainer/test_trainer.py
patch: @@ -52,6 +52,7 @@
     set_seed,
 )
 from transformers.hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS
+from transformers.integrations.neftune import activate_neftune
 from transformers.testing_utils import (
     ENDPOINT_STAGING,
     TOKEN,
@@ -1689,7 +1690,7 @@ def test_neftune(self):
         )
         trainer = Trainer(tiny_gpt2, args, train_dataset=train_dataset)
 
-        trainer.model = trainer._activate_neftune(trainer.model)
+        activate_neftune(trainer.model, trainer.args.neftune_noise_alpha)
 
         dummy_input = torch.LongTensor([[1, 0, 1]]).to(torch_device)