fix peft memory cost issue (#2133)

Author: lvyufeng

.gitignore

@@ -171,3 +171,5 @@ tests/diffusers/
 tests/transformers/
 tests/huggingface_transformers/
 .gradio/
+
+huanhuan.json

examples/transformers/peft/lora/Qwen2.5-7B-Instruct-Lora.ipynb

@@ -55,9 +55,8 @@
 from .func import vmap
 from .configs import set_pyboost
 
-from . import profiler, cuda, optim, amp, compiler, jit, version, __future__, overrides, \
-    return_types, linalg, fx, backends, testing, nn, fft, _jit_internal, utils
-
+from . import profiler, cuda, amp, compiler, jit, version, __future__, overrides, \
+    return_types, linalg, fx, backends, testing, nn, fft, _jit_internal, utils, optim
 from ._lowrank import svd_lowrank
 from .random import get_rng_state, initial_seed, manual_seed, seed, set_rng_state
 

mindnlp/core/__init__.py

@@ -3,6 +3,7 @@
 from mindspore.ops.auto_generate import gen_ops_prim
 from mindspore.ops.auto_generate import pyboost_inner_prim
 from mindspore._c_expression import _empty_instance
+from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloatStatusV2
 
 from mindnlp import core
 from mindnlp.core._C import default_generator
@@ -162,4 +163,33 @@ def reverse_v2(input, dims):
         dims = (dims,)
     return pyboost_inner_prim.reverse_v2_impl(input, dims)
 
-__all__.append('reverse_v2')
+adam_op = ops.Adam().set_device('Ascend')
+def raw_adam(param, exp_avg, exp_avg_sq, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad):
+    # var, m, v, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad
+    return adam_op(param, exp_avg, exp_avg_sq, beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad)
+
+__all__.append('raw_adam')
+
+depend_op = ops.Depend().set_device('Ascend')
+def depend(*args):
+    return depend_op(*args)
+
+__all__.append('depend')
+
+npu_get_float_status_op = NPUGetFloatStatusV2().set_device('Ascend')
+def npu_get_float_status_v2(status):
+    return npu_get_float_status_op(status)
+
+__all__.append('npu_get_float_status_v2')
+
+npu_clear_float_status_op = NPUClearFloatStatusV2().set_device('Ascend')
+def npu_clear_float_status_v2(status):
+    return npu_clear_float_status_op(status)
+
+__all__.append('npu_clear_float_status_v2')
+
+stop_gradient_op = ops.StopGradient().set_device('Ascend')
+def stop_gradient(*args):
+    return stop_gradient_op(*args)
+
+__all__.append('stop_gradient')

mindnlp/core/_prims/ascend.py

@@ -511,3 +511,10 @@ def tril_ext(input, diagonal):
     return core.Tensor.from_numpy(out)
 
 __all__.append('tril_ext')
+
+def randperm_ext(n, seed, offset, dtype):
+    out = np.random.permutation(n)
+    return core.Tensor.from_numpy(out)
+
+__all__.append('randperm_ext')
+

mindnlp/core/_prims/numpy.py

@@ -97,7 +97,11 @@ def tensor(data, *, dtype=None, device=None, requires_grad=False):
     if isinstance(data, Tensor):
         UserWarning("To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than core.tensor(sourceTensor).")
         out = Tensor(data)
-        out._device = data.device
+        if device is not None:
+            out._device = device
+        else:
+            out._device = data.device
+
         return out
 
     # if isinstance(data, list):
@@ -1306,7 +1310,7 @@ def is_floating_point(self):
     def is_leaf(self):
         if not self.requires_grad:
             return True
-        if self.requires_grad and self._user_created:
+        if self.requires_grad and hasattr(self, 'param_info'):
             return True
         return False
 
@@ -2397,9 +2401,7 @@ def zero_(self):
 
     # Tensor.detach
     def detach(self):
-        out = self.data
-        out._requires_grad = False
-        return out
+        return ops.stop_gradient(self)
 
     # Tensor.detach_
     def detach_(self):
@@ -2501,7 +2503,7 @@ def __mod__(self, other):
         return ops.fmod(self, other)
 
     def backward(self):
-        pass
+        return self
 
     def log_softmax(self, dim):
         return ops.log_softmax(self, dim)

mindnlp/core/_tensor.py

@@ -7,12 +7,32 @@
 from enum import Enum
 from typing import Any, cast, Dict, Iterable, List, Optional, overload, Tuple, Union
 
+import numpy as np
+import mindspore
 from mindnlp import core
-
+from mindnlp.core.configs import DEVICE_TARGET, SOC
 
 __all__ = ["OptState", "GradScaler"]
 
 
+def non_finite_check(inputs):
+    """all finite check"""
+    if DEVICE_TARGET == 'Ascend':
+        status = core.tensor(np.array([0] * 8), dtype=core.int32, device='npu')
+        status = core.depend(status, inputs)
+        found_inf = core.npu_get_float_status_v2(status)
+        status = core.depend(status, found_inf)
+        clear_status = core.npu_clear_float_status_v2(status)
+        found_inf = core.depend(found_inf, clear_status)
+        return found_inf.sum()
+
+    found_inf = core.all_finite(inputs)  # pylint: disable=invalid-unary-operand-type
+    # else:
+    #     outputs = _hypermap(_partial(_overflow), inputs)
+    #     flag_sum = ops.addn(outputs).reshape(())
+    #     status_finite = ops.less(flag_sum, 1)
+    return found_inf
+
 class _MultiDeviceReplicator:
     """Lazily serves copies of a tensor to requested devices.
 
@@ -276,11 +296,15 @@ def _unscale_grads_(
 
             for device, per_dtype_grads in per_device_and_dtype_grads.items():
                 for grads in per_dtype_grads.values():
-                    core._amp_foreach_non_finite_check_and_unscale_(
-                        grads,
-                        per_device_found_inf.get(device),
-                        per_device_inv_scale.get(device),
-                    )
+                    # core._amp_foreach_non_finite_check_and_unscale_(
+                    #     grads,
+                    #     per_device_found_inf.get(device),
+                    #     per_device_inv_scale.get(device),
+                    # )
+                    found_inf = per_device_found_inf.get(device)
+                    found_inf.copy_(non_finite_check(grads).to(found_inf.dtype))
+                    for grad in grads:
+                        grad *= per_device_inv_scale.get(device)
 
         return per_device_found_inf._per_device_tensors
 
@@ -518,14 +542,32 @@ def update(self, new_scale: Optional[Union[float, core.Tensor]] = None) -> None:
                 for i in range(1, len(found_infs)):
                     found_inf_combined += found_infs[i]
 
-            core._amp_update_scale_(
-                _scale,
-                _growth_tracker,
-                found_inf_combined,
-                self._growth_factor,
-                self._backoff_factor,
-                self._growth_interval,
-            )
+            # core._amp_update_scale_(
+            #     _scale,
+            #     _growth_tracker,
+            #     found_inf_combined,
+            #     self._growth_factor,
+            #     self._backoff_factor,
+            #     self._growth_interval,
+            # )
+            if found_inf_combined > 0:
+                _scale.copy_(_scale * self._backoff_factor)
+                _growth_tracker.copy_(_growth_tracker * 0)
+            else:
+                successful = self._growth_interval + 1
+                if successful == self._growth_interval:
+                    new_scale = _scale * self._growth_factor
+                    if core.isfinite(new_scale):
+                        _scale.copy_(new_scale)
+                    _growth_tracker.copy_(_growth_tracker * 0)
+                else:
+                    _growth_tracker.copy_(
+                        core.tensor(successful,
+                                    dtype=_growth_tracker.dtype,
+                                    device=_growth_tracker.device
+                        )
+                    )
+
 
         # To prepare for next iteration, clear the data collected from optimizers this iteration.
         self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)

mindnlp/core/amp/grad_scaler.py

@@ -13,6 +13,8 @@
 except:
     Function = None
 
+from mindnlp import core
+
 grad_ = GradOperation(False, True, False)
 grad_sens_ = GradOperation(False, True, True)
 grad_input_sens_ = GradOperation(True, True, True)
@@ -29,7 +31,7 @@ def fn_aux(*args):
         outputs = fn(*args)
         no_grad_outputs = ()
         for out in outputs[1:]:
-            no_grad_outputs += (stop_gradient(out),)
+            no_grad_outputs += (out.detach(),)
         return outputs[0], no_grad_outputs
 
     if has_aux:
@@ -39,19 +41,19 @@ def fn_aux(*args):
 
     def value_and_grad_f(*args, **kwargs):
         _pynative_executor.set_grad_flag(True)
-        _pynative_executor.new_graph(fn, *args, **kwargs)
+        _pynative_executor.new_graph(fn_, *args, **kwargs)
         values = fn_(*args, **kwargs)
-        _pynative_executor.end_graph(fn, values, *args, **kwargs)
+        _pynative_executor.end_graph(fn_, values, *args, **kwargs)
 
         run_args = args
         if kwargs:
             run_args = args + tuple(kwargs.values())
 
         grads = _pynative_executor.check_run(grad_, fn_, params_or_argnums, None, *run_args)
         grads = _pynative_executor.grad(fn_, grad_, params_or_argnums, None, *run_args)
-        grads = tuple(mindspore.Tensor(grad) for grad in grads)
         if attach_grads:
             for param, grad in zip(params_or_argnums, grads):
+                grad = core.tensor(grad, device=param.device)
                 if param.grad is None:
                     param.grad = grad
                 else:

mindnlp/core/autograd/function.py

@@ -20,6 +20,13 @@
 HalfTensor = core.FloatTensor
 BFloat16Tensor = core.BFloat16Tensor
 
+
+class DefaultGenerators:
+    def __getitem__(self, idx):
+        return core.default_generator
+
+default_generators = DefaultGenerators()
+
 def set_compile_mode(*args, **kwargs):
     pass
 

mindnlp/core/npu/__init__.py

@@ -0,0 +1,23 @@
+from typing import Union
+
+from mindnlp import core
+
+
+def get_rng_state(device: Union[int, str, core.device] = "npu") -> core.Tensor:
+    r"""Return the random number generator state of the specified GPU as a ByteTensor.
+
+    Args:
+        device (msadapter.device or int, optional): The device to return the RNG state of.
+            Default: ``'cuda'`` (i.e., ``msadapter.device('cuda')``, the current CUDA device).
+
+    .. warning::
+        This function eagerly initializes CUDA.
+    """
+    
+    if isinstance(device, str):
+        device = core.device(device)
+    elif isinstance(device, int):
+        device = core.device("npu", device)
+    idx = device.index
+    default_generator = core.npu.default_generators[idx]
+    return default_generator.get_state()