chore: import upstream snapshot with attribution

swift/callbacks/__init__.py

@@ -0,0 +1,2 @@
+from .base import TrainerCallback
+from .mapping import callbacks_map

swift/callbacks/activation_cpu_offload.py

@@ -0,0 +1,607 @@
+"""Functionality for CPU offloading of tensors saved for backward pass."""
+import functools
+import torch
+from torch.distributed.fsdp import FSDPModule as FSDP2
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from transformers.trainer_callback import TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+from transformers.utils import is_torch_npu_available
+from typing import Any, Optional
+
+from swift.utils import get_logger
+from .base import TrainerCallback
+
+logger = get_logger()
+
+is_cuda_available = torch.cuda.is_available()
+is_npu_available = is_torch_npu_available()
+
+
+def _get_unique_tensor_key(tensor):
+    key = (tensor.untyped_storage().data_ptr() + tensor.storage_offset(), tensor.dtype)
+    return key
+
+
+def get_device_name() -> str:
+    """Function that gets the torch.device based on the current machine.
+    This currently only supports CPU, CUDA, NPU.
+    Returns:
+        device
+    """
+    if is_cuda_available:
+        device = 'cuda'
+    elif is_npu_available:
+        device = 'npu'
+    else:
+        device = 'cpu'
+    return device
+
+
+class FSDPParameterFilter:
+
+    def __init__(self):
+        self.model_parameters_storage = set()
+
+    def __call__(self, tensor):
+        return tensor.untyped_storage().data_ptr() not in self.model_parameters_storage
+
+    def update_model_parameters(self, model):
+        new_storage = set()
+        for p in model.parameters():
+            new_storage.add(p.data.untyped_storage().data_ptr())
+        self.model_parameters_storage = new_storage
+
+
+def get_torch_device() -> Any:
+    """Return the corresponding torch attribute based on the device type string.
+    Returns:
+        module: The corresponding torch device namespace, or torch.cuda if not found.
+    """
+    device_name = get_device_name()
+    try:
+        return getattr(torch, device_name)
+    except AttributeError:
+        logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.")
+        return torch.cuda
+
+
+class CpuOffloadHookWithOffloadHandler:
+    """Context-manager that offloads/recovers tensors through an offload hander.
+
+    The hook just offloads/recovers the tensor object to the handler through `tensor_push`
+    and `tensor_pop` interface. How the offload-handler manages the offloading, recovering
+    or prefetching timing is transparent to this hook.
+    """
+
+    def __init__(
+        self,
+        offload_handler: 'OffloadHandler',
+        handler_extra_kwargs: Optional[dict[str, Any]] = None,
+    ) -> None:
+        if handler_extra_kwargs is None:
+            handler_extra_kwargs = {}
+        self.offload_handler: OffloadHandler = offload_handler
+        self.handler_extra_kwargs: dict[str, Any] = handler_extra_kwargs
+        self.inside_context = False
+
+    def __enter__(self):
+        self.inside_context = True
+        torch._C._autograd._push_saved_tensors_default_hooks(self.on_save_for_backward, self.on_get_saved_tensor)
+
+    def __exit__(self, *args: Any):
+        self.inside_context = False
+        torch._C._autograd._pop_saved_tensors_default_hooks()
+
+    def on_save_for_backward(self, tensor: torch.Tensor) -> Any:
+        retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs)
+        return retrieve_identifier
+
+    def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor:
+        tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs)
+        return tensor
+
+
+class OffloadHandler:
+    """A base class for CPU offload-handler."""
+
+    def __init__(self) -> None:
+        pass
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
+        """Tensor push."""
+        raise NotImplementedError(
+            '`tensor_push is not implented in OffloadHandler class. Inherit this class and implement your '
+            'custom tensor_push.')
+
+    def tensor_pop(self, tensor_tag: Any, **kwargs):
+        """Tensor pop."""
+        raise NotImplementedError(
+            '`tensor_pop is not implented in OffloadHandler class. Inherit this class and implement your '
+            'custom tensor_pop.')
+
+
+class GroupCommitFunction(torch.autograd.Function):
+    """this is a dummy op with output identical to input.
+    However, it is necessary for marking a timepoint for offload handler to
+    accomplish all synchronizations. Implementing it as a function is necessary
+    because we need to actions in both forward and backward.
+    """
+
+    @staticmethod
+    def forward(ctx, tensor, cpu_offload_handler):
+        # pylint: disable=missing-function-docstring
+        cpu_offload_handler.on_group_commit_forward()
+        ctx.cpu_offload_handler = cpu_offload_handler
+        # return the identical tensor
+        return tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # pylint: disable=missing-function-docstring
+        cpu_offload_handler = ctx.cpu_offload_handler
+        cpu_offload_handler.on_group_commit_backward()
+        return grad_output, None
+
+
+group_prefetch_offload_commit = GroupCommitFunction.apply
+
+
+class SynchronizedGroupOffloadHandler(OffloadHandler):
+    """Offload Handler that offloads/reloads in a synchronized way.
+    The device-to-host and host-to-device copying happen in the same stream
+    as the computation kernels, thus the copying will block computation.
+    """
+
+    def __init__(self, num_offload_group, tensor_need_offloading_checker=(lambda _: True)) -> None:
+        super().__init__()
+
+        self.num_offload_group = num_offload_group
+        self.tensor_need_offloading_checker = tensor_need_offloading_checker
+
+        self.groupid_reset()
+
+    def groupid_reset(self):
+        """Groupid reset."""
+        # Data structures to label saved tensors and book-keep their cpu copies.
+        # Currently, on push, create a new cpu tensor and copies; on pop, copies
+        # the tensor back to gpu and deletes the cpu tensor.
+        # These will increment whenever `group_commit()` is invoked
+        self.current_group, self.tensor_count_current_group = (0, 0)
+        self.torch_tensor_count = 0
+        self.tensor_tag_to_state = {}
+
+    def on_group_commit_forward(self):
+        """On group commit forward."""
+        # finishing up with updating current group and tensor count
+        self.current_group += 1  # increment
+        self.tensor_count_current_group = 0  # reset
+
+    def on_group_commit_backward(self):
+        """On group commit backward."""
+        self.current_group -= 1
+        assert self.current_group >= 0
+
+    @staticmethod
+    def offload(src_tensor, pin_memory=True):
+        """Offload."""
+        # NPU doesn't fully support async H2D/D2H with pinned memory; use sync copy.
+        if is_npu_available:
+            cpu_backup = torch.empty(
+                src_tensor.size(),
+                dtype=src_tensor.dtype,
+                layout=src_tensor.layout,
+                device='cpu',
+                pin_memory=False,
+            )
+            cpu_backup.copy_(src_tensor, non_blocking=False)
+        else:
+            cpu_backup = torch.empty(
+                src_tensor.size(),
+                dtype=src_tensor.dtype,
+                layout=src_tensor.layout,
+                device='cpu',
+                pin_memory=pin_memory,
+            )
+            cpu_backup.copy_(src_tensor, non_blocking=True)
+        state = (src_tensor.device, cpu_backup)
+        return state
+
+    @staticmethod
+    def reload(state, non_blocking=None):
+        """Reload."""
+        dev, cpu_backup = state
+        if non_blocking is None:
+            non_blocking = cpu_backup.is_pinned()
+        return cpu_backup.to(dev, non_blocking=non_blocking)
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs):
+        """Tensor push."""
+        # obtain a unique tensor tag
+        tensor_tag = (self.current_group, self.tensor_count_current_group)
+        self.tensor_count_current_group += 1
+        assert tensor_tag not in self.tensor_tag_to_state
+        if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(tensor):
+            state = SynchronizedGroupOffloadHandler.offload(tensor)
+            self.tensor_tag_to_state[tensor_tag] = state
+        else:
+            # will be offloaded together after group commit
+            self.tensor_tag_to_state[tensor_tag] = tensor
+
+        return tensor_tag
+
+    def tensor_pop(self, tensor_tag, **kwargs):
+        """Tensor pop."""
+        assert tensor_tag in self.tensor_tag_to_state
+        state = self.tensor_tag_to_state.pop(tensor_tag)
+        if isinstance(state, tuple):
+            tensor = SynchronizedGroupOffloadHandler.reload(state)
+        else:
+            tensor = state
+        return tensor
+
+
+class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler):
+    """Compared to synchronize, this uses more memory because of the buffer but
+    achieves better performance due to the overlapping. D2h and h2d copying are
+    completely hidden behind computation if computation time of a layer is longer
+    than host-device communication time. Bulk offloading with delay and bulk reloading
+    with prefetch are implemented."""
+
+    def __init__(
+            self,
+            num_offload_group,  # must be <= actual number of groups (number of commits)
+            num_model_group,
+            tensor_need_offloading_checker=(lambda t: True),
+    ) -> None:
+        super().__init__(
+            num_offload_group=num_offload_group,
+            tensor_need_offloading_checker=tensor_need_offloading_checker,
+        )
+        # Number of layers in the model
+        self.num_layers = num_model_group
+        # Data Structure to maintain reference to activation tensors
+        self.tensor_tag_to_buf = {}
+        # Tracking the number of layers offloaded
+        self.offloaded_group_count = 0
+        # Core data structure that decides the window for offloading
+        self.layer_window_map = {}
+        self.group_offload_mapping = {}
+
+        # Logic to make offloading load balance across computation
+        # for optimal CPU/GPU interconnect usage
+        constant = 0
+        for i in range(self.num_offload_group):
+            self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1
+            if i < (self.num_layers % self.num_offload_group):
+                self.layer_window_map[i] += i + 1
+                constant = i + 1
+            else:
+                self.layer_window_map[i] += constant
+
+        # allocate streams and events for synchronization
+        self.d2h_stream = get_torch_device().Stream()
+        self.h2d_stream = get_torch_device().Stream()
+
+    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:
+        torch_stray_tensor = isinstance(
+            tensor,
+            torch._subclasses.fake_tensor.FakeTensor | torch._subclasses.functional_tensor.FunctionalTensor,
+        )
+        need_offload = not torch_stray_tensor
+        need_offload = need_offload and self.tensor_need_offloading_checker(tensor)
+
+        if need_offload:
+            # obtain a unique tensor tag
+            tensor_tag = (self.current_group, self.tensor_count_current_group)
+            self.tensor_count_current_group += 1
+
+            assert tensor_tag not in self.tensor_tag_to_state
+            self.tensor_tag_to_state[tensor_tag] = tensor
+
+            if self.current_group < self.num_offload_group:
+                self.tensor_tag_to_buf[tensor_tag] = tensor
+        else:
+            tensor_tag = tensor
+        return tensor_tag
+
+    def tensor_pop(self, tensor_tag, **kwargs):
+        """Tensor pop."""
+        if isinstance(tensor_tag, torch.Tensor):
+            return tensor_tag
+        assert tensor_tag in self.tensor_tag_to_state
+        tensor = self.tensor_tag_to_state.pop(tensor_tag)
+        self.tensor_tag_to_buf.pop(tensor_tag, None)
+
+        # the tensor should have been copied back in on_group_commit_backward()
+        # which invokes bulk_reload_group.
+        assert not isinstance(tensor, tuple)
+        return tensor
+
+    def bulk_offload_group(self, group_to_offload):
+        """Bulk offload group."""
+        offload_mapping = {}
+        offload_size = 0
+        with get_torch_device().stream(self.d2h_stream):
+            for tensor_tag, state in self.tensor_tag_to_state.items():
+                group_id, _ = tensor_tag
+                if group_id == group_to_offload:
+                    assert not isinstance(state, tuple)
+                    key = _get_unique_tensor_key(state)
+                    if key not in offload_mapping:
+                        offload_mapping[key] = state
+                    # if offload, return the reference to cpu copy
+                    self.tensor_tag_to_state[tensor_tag] = (key, state.shape)
+            for key, tensor in offload_mapping.items():
+                state = SynchronizedGroupOffloadHandler.offload(tensor)
+                offload_size += tensor.numel() * tensor.element_size()
+                offload_mapping[key] = state
+
+            self.group_offload_mapping[group_to_offload] = offload_mapping
+
+    def synchronize_on_group_commit_forward(self, current_group):
+        """Synchronize on group commit forward."""
+
+        # For the first group, kickstart the offload after we have
+        # the first compute completion
+        if current_group == 0:
+            self.d2h_stream.wait_stream(get_torch_device().current_stream())
+            self.bulk_offload_group(current_group)
+
+        # Window map data structure helps us synchronize based on number
+        # of layers offloaded
+        if self.layer_window_map[self.offloaded_group_count] == current_group:
+            # Stream synchronization both ways
+            self.d2h_stream.wait_stream(get_torch_device().current_stream())
+            get_torch_device().current_stream().wait_stream(self.d2h_stream)
+
+            # Time to free the activation memory after usage
+            for tensor_tag, _ in self.tensor_tag_to_buf.items():
+                if tensor_tag[0] == self.offloaded_group_count:
+                    self.tensor_tag_to_buf[tensor_tag] = None
+
+            # Time to offload the next group
+            if self.offloaded_group_count < (self.num_offload_group - 1):
+                self.bulk_offload_group(self.offloaded_group_count + 1)
+
+            # Increment the offload group count to keep track
+            self.offloaded_group_count += 1
+
+    def on_group_commit_forward(self):
+        """This function will cause host device synchronization"""
+        # handle synchronization events
+        self.synchronize_on_group_commit_forward(self.current_group)
+
+        super().on_group_commit_forward()
+
+    @torch.no_grad
+    def bulk_reload_group(self, group_to_reload):
+        """Bulk reload group."""
+        assert group_to_reload < self.num_offload_group
+
+        with get_torch_device().stream(self.h2d_stream):
+            # move back tensors
+            offload_mapping = self.group_offload_mapping.pop(group_to_reload)
+            assert offload_mapping is not None
+            for key, state in offload_mapping.items():
+                offload_mapping[key] = SynchronizedGroupOffloadHandler.reload(state)
+            for tensor_label, state in self.tensor_tag_to_state.items():
+                group_id, _ = tensor_label
+                if group_id == group_to_reload and not isinstance(state, torch.Tensor):
+                    assert isinstance(state, tuple), f'{group_id} {state}'
+                    key, shape = state
+                    recovered_tensor = offload_mapping[key].view(shape)
+                    self.tensor_tag_to_state[tensor_label] = recovered_tensor
+
+    def on_group_commit_backward(self):
+        # first decrement the current group.
+        # after last commit in forward, the group will +1; in backward it -1.
+        # Finally it should be decremented to 0.
+        self.current_group -= 1
+        assert self.current_group >= 0
+
+        # Layer window data structure helps us to reload at right times
+        if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group:
+            # Stream synchronization both ways
+            self.h2d_stream.wait_stream(get_torch_device().current_stream())
+            get_torch_device().current_stream().wait_stream(self.h2d_stream)
+
+            # Time to reload the next group
+            self.bulk_reload_group(self.offloaded_group_count - 1)
+
+            # Decrease the offloading group counter
+            self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0
+
+        # Last group computation needs to wait till all the reloads complete
+        if self.current_group == 0:
+            get_torch_device().current_stream().wait_stream(self.h2d_stream)
+            self.offloaded_group_count = 0
+
+
+def get_activation_offload_context(num_layers: int = 1,
+                                   model_layers: int = 1,
+                                   tensor_need_offloading_checker=(lambda t: True)):
+    cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler(
+        num_offload_group=num_layers,
+        num_model_group=model_layers,
+        tensor_need_offloading_checker=tensor_need_offloading_checker,
+    )
+
+    def group_prefetch_offload_commit_async(tensor):
+        return group_prefetch_offload_commit(tensor, cpu_offload_handler)
+
+    return (
+        CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler),
+        group_prefetch_offload_commit_async,
+    )
+
+
+class ActivationHandler:
+
+    def __init__(self, offload_ctx, sync_func, tensor_filter, enable_ckpt):
+        self._offload_ctx = offload_ctx
+        self._sync_func = sync_func
+        self._enable_ckpt = enable_ckpt
+        self._tensor_filter = tensor_filter
+        if enable_ckpt:
+            self.checkpoint_fn = functools.partial(
+                torch.utils.checkpoint.checkpoint,
+                use_reentrant=True,
+            )
+
+    def pre_forward(self, module):
+        if module.training:
+            self._offload_ctx.__enter__()
+            self._tensor_filter.update_model_parameters(module)
+
+    def post_forward(self, module):
+        if module.training:
+            self._offload_ctx.__exit__(None, None, None)
+
+    def _pack_kwargs(self, *args, **kwargs):
+        kwarg_keys = []
+        flat_args = list(args)
+        for k, v in kwargs.items():
+            kwarg_keys.append(k)
+            flat_args.append(v)
+
+        return tuple(flat_args), tuple(kwarg_keys)
+
+    def _unpack_kwargs(self, flat_args, kwarg_keys):
+        assert len(kwarg_keys) <= len(flat_args), f'too many keys {len(kwarg_keys)} vs. {len(flat_args)}'
+        if len(kwarg_keys) == 0:
+            return flat_args, {}
+        args = flat_args[:-len(kwarg_keys)]
+        kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys):], strict=True))
+        return args, kwargs
+
+    def _ckpt_forward(self, forward_method, *args, **kwargs):
+        flat_args, kwarg_keys = self._pack_kwargs(*args, **kwargs)
+
+        def my_function(*inputs):
+            # unpack back into args and kwargs
+            unpacked_args, unpacked_kwargs = self._unpack_kwargs(inputs, kwarg_keys)
+            # run original module
+            return forward_method(*unpacked_args, **unpacked_kwargs)
+
+        return self.checkpoint_fn(
+            my_function,
+            *flat_args,
+        )
+
+    def forward(self, module, forward_method, *args, **kwargs):
+        if not module.training:
+            return forward_method(*args, **kwargs)
+        if not self._enable_ckpt:
+            ret = forward_method(*args, **kwargs)
+        else:
+            ret = self._ckpt_forward(forward_method, *args, **kwargs)
+        binded_tensor = ret
+        if isinstance(ret, tuple):
+            binded_tensor = ret[0]
+        binded_tensor = self._sync_func(binded_tensor)
+        final_ret = binded_tensor
+        if isinstance(ret, tuple):
+            final_ret = (final_ret, ) + ret[1:]
+        return final_ret
+
+    def wrap_module_forward_method(self, module):
+        orig_method = module.forward
+        handler = self
+
+        @functools.wraps(orig_method)
+        def wrapped_method(model_self, *args, **kwargs):
+            handler.pre_forward(model_self)
+            out = handler.forward(model_self, orig_method, *args, **kwargs)
+            handler.post_forward(model_self)
+            return out
+
+        module.forward = wrapped_method.__get__(module, type(module))
+
+
+def enable_activation_offloading(model, strategy, enable_ckpt=False):
+    """
+    Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation
+    groups asynchronously. This means that the offloading of the i-th activation group and the computation of the i+1-th
+    activation group happen at the same time, and there are at most two activation groups in GPU memory.
+
+    Args:
+        model: the model to enable activation offloading
+        strategy: the training strategy of the model, such as "fsdp"
+        enable_ckpt: whether activation checkpointing(also called gradient checkpointing) has been enabled for the model
+
+    Note:
+        For best efficiency, activation offloading is usually combined with activation checkpointing. However, this
+        implementation of activation offloading is conflicted with the implementation of activation checkpointing in
+        some training strategies. This function resolves this conflict, and therefore requires the "strategy" and
+        "enable_ckpt" arguments.
+
+    Returns:
+
+    """
+
+    assert strategy == 'fsdp' or strategy == 'fsdp2', 'activation offloading only supports fsdp strategy'
+    layers = []
+
+    def get_layers(module):
+        for name, child in module.named_children():
+            if not isinstance(child, FSDP | FSDP2):
+                get_layers(child)
+            else:
+                wrapped_module = child
+                if isinstance(child, FSDP):
+                    wrapped_module = child._fsdp_wrapped_module
+                # In some cases, torch.nn.Embedding is wrapped with FSDP alone. However, the activation
+                # size of torch.nn.Embedding is small, so it's not necessary to offload it.
+                if not isinstance(wrapped_module, torch.nn.Embedding):
+                    layers.append(child)
+
+    get_layers(model)
+    if len(layers) < 3:
+        logger.warning(f'Find only {len(layers)} fsdp layers, not necessary to enable async activation offloading')
+        return
+
+    tensor_filter = FSDPParameterFilter()
+    context, sync_func = get_activation_offload_context(len(layers) - 1, len(layers), tensor_filter)
+    if enable_ckpt:
+        # The implementation of activation checkpointing in transformers library is incompatible with
+        # activation offloading,
+        # so it will be disabled, but this implementation supports another version of activation checkpointing, so that
+        # these two features can be enabled at the same time.
+        for module in model.modules():
+            if hasattr(module, 'gradient_checkpointing_disable'):
+                module.gradient_checkpointing_disable()
+
+    handler = ActivationHandler(context, sync_func, tensor_filter, enable_ckpt)
+    for layer in layers:
+        module = layer
+        if isinstance(layer, FSDP):
+            module = module._fsdp_wrapped_module
+        handler.wrap_module_forward_method(module)
+
+
+class ActivationCpuOffloadCallBack(TrainerCallback):
+
+    def __init__(self, args: TrainingArguments, trainer):
+        super().__init__(args, trainer)
+
+    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        """
+        Event called at the beginning of training.
+        """
+        model = kwargs['model']
+
+        # Check if model is wrapped with FSDP
+        if isinstance(model, FSDP) or isinstance(model, FSDP2):
+            if args is not None and hasattr(args, 'fsdp_config'):
+                fsdp_config = args.fsdp_config
+                # Check if fsdp_config is a dictionary and has activation_cpu_offload enabled
+                if isinstance(fsdp_config, dict) and fsdp_config.get('activation_cpu_offload', False):
+                    # Get FSDP version from fsdp_config
+                    strategy = fsdp_config.get('version', None)
+                    if strategy is not None:
+                        fsdp_version = 'fsdp' if strategy == 1 else 'fsdp2'
+                        # Get activation checkpointing setting from fsdp_config
+                        enable_ckpt = fsdp_config.get('activation_checkpointing', False)
+                        if enable_ckpt and hasattr(model, 'enable_input_require_grads'):
+                            model.enable_input_require_grads()
+                        enable_activation_offloading(model, strategy=fsdp_version, enable_ckpt=enable_ckpt)

swift/callbacks/adalora.py

@@ -0,0 +1,31 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import types
+from typing import TYPE_CHECKING
+
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class AdaloraCallback(TrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.global_step = 0
+        self.args = args
+
+    # offload original_modules to cpu, to save memory
+    def on_train_begin(self, _args, state, control, **kwargs):
+        model = kwargs['model']
+        model.peft_config['default'].total_step = state.max_steps
+
+        def zero_grad(_self, *args, **kwargs):
+            _self.update_and_allocate(self.global_step + 1)
+            _self._zero_grad(*args, **kwargs)
+
+        model._zero_grad = model.zero_grad
+        model.zero_grad = types.MethodType(zero_grad, model)
+
+    def on_step_end(self, _args, state, control, **kwargs):
+        self.global_step = state.global_step

swift/callbacks/base.py

@@ -0,0 +1,13 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from transformers import TrainerCallback as HfTrainerCallback
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class TrainerCallback(HfTrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        self.args = args
+        self.trainer = trainer

swift/callbacks/deepspeed_elastic.py

@@ -0,0 +1,45 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+
+from swift.utils import ShutdownManager, get_device
+from .base import TrainerCallback
+
+
+class DeepspeedElasticCallback(TrainerCallback):
+    """Compatibility marker for enabling DeepSpeed elastic setup during argument initialization."""
+
+
+class GracefulExitCallback(TrainerCallback):
+
+    def __init__(self, args=None, trainer=None):
+        if args is not None and trainer is not None:
+            super().__init__(args, trainer)
+        shutdown_manager = ShutdownManager()
+        shutdown_manager.register()
+        self.shutdown_manager = shutdown_manager
+        self._pending_stop = False
+
+    def on_step_end(self, args, state, control, **kwargs):
+        device_type = get_device()
+
+        local_req = 1 if self.shutdown_manager.should_shutdown() else 0
+        if dist.is_available() and dist.is_initialized():
+
+            t = torch.tensor([local_req], dtype=torch.uint8, device=device_type)
+            # all_reduce with MAX: if any rank has 1 -> result 1 everywhere
+            dist.all_reduce(t, op=dist.ReduceOp.MAX)
+            any_req = bool(int(t.item()))
+        else:
+            any_req = bool(local_req)
+
+        if any_req:
+            control.should_save = True
+            self._pending_stop = True
+        return control
+
+    def on_save(self, args, state, control, **kwargs):
+        if self._pending_stop:
+            control.should_training_stop = True
+            self._pending_stop = False
+        return control

swift/callbacks/early_stop.py

@@ -0,0 +1,34 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+from transformers import TrainerControl, TrainerState
+from typing import TYPE_CHECKING
+
+from swift.utils import get_logger
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+logger = get_logger()
+
+
+class EarlyStopCallback(TrainerCallback):
+    """An early stop implementation"""
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.best_metric = None
+        self.interval = 0
+        self.total_interval = args.early_stop_interval
+
+    def on_save(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        operator = np.greater if args.greater_is_better else np.less
+        if self.best_metric is None or operator(state.best_metric, self.best_metric):
+            self.best_metric = state.best_metric
+            self.interval = 0
+        else:
+            self.interval += 1
+
+        if self.interval >= self.total_interval:
+            logger.info(f'Training stop because of eval metric is stable at step {state.global_step}')
+            control.should_training_stop = True

swift/callbacks/lisa.py

@@ -0,0 +1,57 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import torch
+from typing import TYPE_CHECKING
+
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class LISACallback(TrainerCallback):
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        assert args.tuner_type == 'full', 'LISA only supports full parameter training.'
+        super().__init__(args, trainer)
+        self.n_layers = args.lisa_activated_layers
+        self.step_interval = args.lisa_step_interval
+        self.model = self.trainer.model
+        layers_name = None
+        layers = None
+        for name, module in self.model.named_modules():
+            if isinstance(module, torch.nn.ModuleList):
+                layers_name = name
+                layers = module
+                break
+        assert layers_name is not None
+        self.layers_attribute = layers_name
+        self.total_layers = len(layers)
+
+        # Freeze all layers upon initialization
+        self.freeze_all_layers()
+        self.active_layers_indices = []
+        self.switch_active_layers()
+
+    def freeze_all_layers(self):
+        layers = self.model.get_submodule(self.layers_attribute)
+        for layer in layers:
+            for param in layer.parameters():
+                param.requires_grad = False
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        # Check if it's time to switch active layers, including at step 0
+        if state.global_step % self.step_interval == 0 or state.global_step == 1:
+            self.switch_active_layers()
+
+    def switch_active_layers(self):
+        # First, disable gradients for all layers
+        self.freeze_all_layers()
+
+        # Randomly select n_layers to activate
+        layers = self.model.get_submodule(self.layers_attribute)
+        self.active_layers_indices = np.random.choice(range(self.total_layers), self.n_layers, replace=False)
+        # Enable gradients only for the selected layers
+        for idx in self.active_layers_indices:
+            for param in layers[idx].parameters():
+                param.requires_grad = True

swift/callbacks/mapping.py

@@ -0,0 +1,17 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .activation_cpu_offload import ActivationCpuOffloadCallBack
+from .adalora import AdaloraCallback
+from .deepspeed_elastic import DeepspeedElasticCallback, GracefulExitCallback
+from .early_stop import EarlyStopCallback
+from .lisa import LISACallback
+from .perf_log import PerfMetricsLogCallback
+
+callbacks_map = {
+    'activation_cpu_offload': ActivationCpuOffloadCallBack,
+    'adalora': AdaloraCallback,
+    'deepspeed_elastic': DeepspeedElasticCallback,
+    'early_stop': EarlyStopCallback,
+    'graceful_exit': GracefulExitCallback,
+    'lisa': LISACallback,
+    'perf_log': PerfMetricsLogCallback
+}

swift/callbacks/perf_log.py

@@ -0,0 +1,137 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import time
+import torch
+from transformers import TrainerControl, TrainerState
+from typing import TYPE_CHECKING
+
+from swift.utils import (empty_cache, get_current_device, get_device_count, get_dist_setting, get_env_args, get_logger,
+                         synchronize)
+from .base import TrainerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+logger = get_logger()
+
+device_flops_map = {
+    'GB200': 2.5e15,
+    'B200': 2.25e15,
+    'MI300X': 1336e12,
+    'H100': 312e12,
+    'H800': 312e12,
+    'H200': 989e12,
+    'A100': 312e12,
+    'A800': 312e12,
+    'L40S': 362.05e12,
+    'L40': 181.05e12,
+    'A40': 149.7e12,
+    'L20': 119.5e12,
+    'H20': 148e12,
+    '910B': 354e12,
+    'Ascend910': 354e12,
+    'RTX 3070 Ti': 21.75e12
+}
+
+
+class PerfMetricsLogCallback(TrainerCallback):
+    """An callback for perf metrics (MFU etc) log implementation"""
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        super().__init__(args, trainer)
+        self.max_tflops = None
+        self.elapsed = 0.0
+        self.step_start_time = None
+
+    def on_init_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+
+        # Top priority. Specify by ENV
+        tflops = get_env_args('DEVICE_TFLOPS', float, None)
+        # `state.total_flos` is summed across all ranks (cluster-global) by HF Trainer,
+        # so the theoretical denominator must use the TOTAL number of GPUs in use across the entire cluster.
+        _, _, world_size, local_world_size = get_dist_setting()
+        local_n_gpu = get_device_count()
+        gpus_per_process = max(1, local_n_gpu // max(local_world_size, 1))
+        device_count = max(world_size * gpus_per_process, 1)
+        if tflops is not None:
+            logger.info(f"Specify theoretical max TFLOPS through ENV 'DEVICE_TFLOPS'. [{tflops} TFLOPS]")
+        else:
+            # Run a estimating test.
+            dtype = kwargs.get('model').dtype
+            device = torch.device(get_current_device())
+            logger.info(f'Estimating device TFLOPS baseline. Device: [{device}] dtype: [{dtype}]')
+            tflops = self._estimate_device_tflops_by_dtype(device, dtype)
+            logger.info(f'Estimate test finished. [{tflops} TFLOPS] Device count: [{device_count}]')
+
+        self.max_tflops = tflops * device_count
+
+    def on_step_begin(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        self.step_start_time = time.time()
+
+    def on_step_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs):
+        self.elapsed += time.time() - self.step_start_time
+
+    def on_log(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, logs=None, **kwargs):
+        total_flos = getattr(state, 'total_flos', 0)
+        actual_flops = total_flos / self.elapsed
+        theoretical_max_flops = self.max_tflops * 1e12
+        mfu = actual_flops / theoretical_max_flops
+        logger.debug(f'Total_flos[{total_flos}] elapsed_time[{self.elapsed}]sec Average MFU[{mfu}]')
+        logs['MFU'] = round(mfu, 6)
+
+    @staticmethod
+    def _estimate_device_tflops_by_dtype(device: torch.device, dtype: torch.dtype, repeats: int = 60, dim: int = 8192):
+        # Set matrix dimension
+        shape = (dim, dim)
+        backend = device.type
+        if backend == 'npu':
+            import torch_npu
+
+        # Initialize matrices
+        a = torch.randn(*shape, device=device, dtype=dtype)
+        b = torch.randn(*shape, device=device, dtype=dtype)
+
+        # Warm-up
+        for _ in range(5):
+            c = torch.matmul(a, b)
+        synchronize(device)
+
+        # Run benchmark test
+        start = time.time()
+        for _ in range(repeats):
+            c = torch.matmul(a, b)
+        synchronize(device)
+        end = time.time()
+        total_time = end - start
+        avg_time = total_time / repeats
+
+        # Adjust repeat count and retest if test duration is too short
+        if total_time < 3:
+            repeats = int(6 / avg_time)
+            start = time.time()
+            for _ in range(repeats):
+                c = torch.matmul(a, b)
+            synchronize(device)
+            end = time.time()
+            total_time = end - start
+            avg_time = total_time / repeats
+
+        del a, b, c
+        empty_cache()
+
+        tflops = (2 * dim**3 / avg_time) / 1e12
+        logger.info(f'[Device {device}] Total time: {total_time:.4f}s, dtype: {dtype}, Perf: {tflops:.4f} TFLOPS')
+        return tflops
+
+    @staticmethod
+    def _retrieve_flops_from_map(device):
+        """Retrieve theoretical FLOPS from Map.    """
+        # This function is never used.
+
+        device_name = device.get_device_name()
+        flops = None
+        for name, value in device_flops_map.items():
+            if name in device_name:
+                flops = value
+                break
+
+        return flops