checkpointing.py 7.38 KB
Newer Older
chenych's avatar
chenych committed
1
# Copyright 2025 HuggingFace Inc., Daniel Han-Chen & the Unsloth team and the LlamaFactory team.
chenych's avatar
chenych committed
2
#
luopl's avatar
luopl committed
3
# This code is inspired by the HuggingFace's Transformers and PEFT library,
chenych's avatar
chenych committed
4
5
# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/modeling_utils.py
# https://github.com/huggingface/peft/blob/v0.10.0/src/peft/utils/other.py
luopl's avatar
luopl committed
6
7
# and the Unsloth library.
# https://github.com/unslothai/unsloth/blob/July-2024/unsloth/models/_utils.py
chenych's avatar
chenych committed
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#
# 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.

import inspect
luopl's avatar
luopl committed
22
from functools import WRAPPER_ASSIGNMENTS, partial, wraps
chenych's avatar
chenych committed
23
from types import MethodType
chenych's avatar
chenych committed
24
from typing import TYPE_CHECKING, Any, Callable, Optional, Union
chenych's avatar
chenych committed
25
26
27

import torch

luopl's avatar
luopl committed
28
from ...extras import logging
chenych's avatar
chenych committed
29
30
31
32
33
34
35
36
37
from ...extras.constants import LAYERNORM_NAMES


if TYPE_CHECKING:
    from transformers import PreTrainedModel

    from ...hparams import ModelArguments


luopl's avatar
luopl committed
38
logger = logging.get_logger(__name__)
chenych's avatar
chenych committed
39
40


luopl's avatar
luopl committed
41
42
def get_unsloth_gradient_checkpointing_func() -> Callable:
    class UnslothGradientCheckpointing(torch.autograd.Function):
chenych's avatar
chenych committed
43
        r"""Saves VRAM by smartly offloading to RAM."""
luopl's avatar
luopl committed
44
45
46
47
48
49
50
51
52
53
54

        @staticmethod
        @torch.cuda.amp.custom_fwd
        def forward(
            ctx: "torch.autograd.Function",
            forward_function: "torch.Module",
            hidden_states: "torch.Tensor",
            *args: Union["torch.Tensor", Any],
        ) -> "torch.Tensor":
            saved_hidden_states = hidden_states.to("cpu", non_blocking=True)
            with torch.no_grad():
chenych's avatar
chenych committed
55
                outputs = forward_function(hidden_states, *args)
luopl's avatar
luopl committed
56
57
58
59

            ctx.save_for_backward(saved_hidden_states)
            ctx.forward_function = forward_function
            ctx.args = args
chenych's avatar
chenych committed
60
            return outputs
luopl's avatar
luopl committed
61
62
63
64
65
66
67
68

        @staticmethod
        @torch.cuda.amp.custom_bwd
        def backward(ctx: "torch.autograd.Function", grad_output: "torch.Tensor") -> "torch.Tensor":
            (hidden_states,) = ctx.saved_tensors
            hidden_states = hidden_states.to("cuda", non_blocking=True).detach()
            hidden_states.requires_grad_(True)
            with torch.enable_grad():
chenych's avatar
chenych committed
69
70
                outputs = ctx.forward_function(hidden_states, *ctx.args)
                output = outputs[0] if isinstance(outputs, tuple) else outputs
luopl's avatar
luopl committed
71
72
73
74
75
76
77
78

            torch.autograd.backward(output, grad_output)
            return (None, hidden_states.grad) + (None,) * len(ctx.args)

    return UnslothGradientCheckpointing.apply


def get_custom_gradient_checkpointing_func(gradient_checkpointing_func: Callable) -> Callable:
chenych's avatar
chenych committed
79
    r"""Only applies gradient checkpointing to trainable layers."""
luopl's avatar
luopl committed
80

luopl's avatar
luopl committed
81
    @wraps(gradient_checkpointing_func, assigned=WRAPPER_ASSIGNMENTS + ("__self__",))
luopl's avatar
luopl committed
82
    def custom_gradient_checkpointing_func(func: Callable, *args: Union["torch.Tensor", Any], **kwargs):
chenych's avatar
chenych committed
83
84
85
86
        if isinstance(func, partial):
            module: torch.nn.Module = func.func.__self__
        else:
            module: torch.nn.Module = func.__self__
luopl's avatar
luopl committed
87

chenych's avatar
chenych committed
88
        has_grad = False
luopl's avatar
luopl committed
89
        if any(param.requires_grad for param in module.parameters()):
chenych's avatar
chenych committed
90
            has_grad = True
luopl's avatar
luopl committed
91
92
93
            for arg in args:
                if torch.is_tensor(arg) and torch.is_floating_point(arg):
                    arg.requires_grad_(True)
chenych's avatar
chenych committed
94
                    break  # assume the first tensor is always the hidden states
luopl's avatar
luopl committed
95

chenych's avatar
chenych committed
96
97
98
99
        if has_grad:
            return gradient_checkpointing_func(func, *args, **kwargs)
        else:
            return func(*args, **kwargs)
luopl's avatar
luopl committed
100
101
102
103

    return custom_gradient_checkpointing_func


chenych's avatar
chenych committed
104
def _gradient_checkpointing_enable(
luopl's avatar
luopl committed
105
    self: "PreTrainedModel",
chenych's avatar
chenych committed
106
    gradient_checkpointing_kwargs: Optional[dict[str, Any]] = None,
luopl's avatar
luopl committed
107
    use_unsloth_gc: bool = False,
chenych's avatar
chenych committed
108
) -> None:
chenych's avatar
chenych committed
109
    r"""Activates gradient checkpointing for the current model.
chenych's avatar
chenych committed
110
111
112
113
114
115

    Modification of the original method to enable gradient checkpointing for block-wise optimizer.
    """
    from torch.utils.checkpoint import checkpoint

    if not self.supports_gradient_checkpointing:
luopl's avatar
luopl committed
116
        raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
chenych's avatar
chenych committed
117
118
119
120

    if gradient_checkpointing_kwargs is None:
        gradient_checkpointing_kwargs = {"use_reentrant": True}

luopl's avatar
luopl committed
121
122
123
124
    if use_unsloth_gc:
        gradient_checkpointing_func = get_unsloth_gradient_checkpointing_func()
    else:
        gradient_checkpointing_func = partial(checkpoint, **gradient_checkpointing_kwargs)
chenych's avatar
chenych committed
125

luopl's avatar
luopl committed
126
    gradient_checkpointing_func = get_custom_gradient_checkpointing_func(gradient_checkpointing_func)
chenych's avatar
chenych committed
127
128
129
    if "value" in inspect.signature(self._set_gradient_checkpointing).parameters:  # old GC format
        self.apply(partial(self._set_gradient_checkpointing, value=True))
        self.enable_input_require_grads()
luopl's avatar
luopl committed
130
        logger.warning_rank0_once("You are using the old GC format, some features (e.g. BAdam) will be invalid.")
chenych's avatar
chenych committed
131
    else:  # have already enabled input require gradients
luopl's avatar
luopl committed
132
        self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func)
chenych's avatar
chenych committed
133
134
135


def _fp32_forward_post_hook(
chenych's avatar
chenych committed
136
    module: "torch.nn.Module", args: tuple["torch.Tensor"], output: "torch.Tensor"
chenych's avatar
chenych committed
137
138
139
140
141
) -> "torch.Tensor":
    return output.to(torch.float32)


def prepare_model_for_training(model: "PreTrainedModel", model_args: "ModelArguments") -> None:
chenych's avatar
chenych committed
142
143
144
145
146
147
    r"""Prepare the model before training.

    Include:
    (1) cast the layernorm in fp32
    (2) make output embedding layer require grads
    (3) add the upcasting of the lm_head in fp32.
chenych's avatar
chenych committed
148
149
    """
    if model_args.upcast_layernorm:
luopl's avatar
luopl committed
150
        logger.info_rank0("Upcasting layernorm weights in float32.")
chenych's avatar
chenych committed
151
152
153
154
155
156
        for name, param in model.named_parameters():
            if param.ndim == 1 and any(ln_name in name for ln_name in LAYERNORM_NAMES):
                param.data = param.data.to(torch.float32)

    if not model_args.disable_gradient_checkpointing:
        if not getattr(model, "supports_gradient_checkpointing", False):
luopl's avatar
luopl committed
157
            logger.warning_rank0("Current model does not support gradient checkpointing.")
chenych's avatar
chenych committed
158
159
160
        else:
            # use_reentrant=False might increase VRAM usage (have not been empirically verified yet)
            # According to: https://github.com/huggingface/transformers/issues/28339
luopl's avatar
luopl committed
161
162
163
164
            gradient_checkpointing_enable = partial(
                _gradient_checkpointing_enable, use_unsloth_gc=model_args.use_unsloth_gc
            )
            model.gradient_checkpointing_enable = MethodType(gradient_checkpointing_enable, model)
luopl's avatar
luopl committed
165
166
167
            model.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": model_args.use_reentrant_gc}
            )
chenych's avatar
chenych committed
168
            setattr(model.config, "use_cache", False)  # turn off when gradient checkpointing is enabled
luopl's avatar
luopl committed
169
            logger.info_rank0("Gradient checkpointing enabled.")
chenych's avatar
chenych committed
170
171
172
173

    if model_args.upcast_lmhead_output:
        output_layer = model.get_output_embeddings()
        if isinstance(output_layer, torch.nn.Linear) and output_layer.weight.dtype != torch.float32:
luopl's avatar
luopl committed
174
            logger.info_rank0("Upcasting lm_head outputs in float32.")
chenych's avatar
chenych committed
175
            output_layer.register_forward_hook(_fp32_forward_post_hook)