Support ddp of unsloth (#5141)

Co-authored-by: tastelikefeet <[email protected]>

Author: 2 people

swift/llm/model/register.py

@@ -2,6 +2,7 @@
 import os
 import platform
 import re
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import asdict, dataclass, field
 from functools import partial
@@ -133,10 +134,27 @@ def load_by_unsloth(args):
     os.environ['UNSLOTH_DISABLE_STATISTICS'] = '1'
     model_info = args.model_info
     model_meta = args.model_meta
-    if model_meta.is_multimodal:
-        from unsloth import FastVisionModel as UnslothModel
-    else:
-        from unsloth import FastLanguageModel as UnslothModel
+
+    os.environ['UNSLOTH_IS_PRESENT'] = '1'
+
+    @contextmanager
+    def _patch_distributed_function():
+        from unsloth_zoo import utils
+
+        def distributed_function(n=1, function=None, *args, **kwargs):
+            return function(*args, **kwargs)
+
+        _origin_distributed_function = utils.distributed_function
+        utils.distributed_function = distributed_function
+        yield
+        utils.distributed_function = _origin_distributed_function
+
+    with _patch_distributed_function():
+        if model_meta.is_multimodal:
+            from unsloth import FastVisionModel as UnslothModel
+        else:
+            from unsloth import FastLanguageModel as UnslothModel
+
     model, processor = UnslothModel.from_pretrained(
         model_name=args.adapters and args.adapters[0] or args.model_dir,
         dtype=args.torch_dtype,

swift/trainers/arguments.py

@@ -23,6 +23,7 @@ class TrainArgumentsMixin:
     per_device_train_batch_size: int = 1
     per_device_eval_batch_size: int = 1
     gradient_accumulation_steps: Optional[int] = None
+    tuner_backend: Optional[str] = None
 
     gradient_checkpointing: bool = True
     vit_gradient_checkpointing: Optional[bool] = None

swift/trainers/trainers.py

@@ -332,6 +332,8 @@ def _prepare_inputs(self, inputs):
             inputs['labels'], logits_to_keep = self.get_logits_to_keep(inputs['labels'])
             if logits_to_keep is not None:
                 inputs['logits_to_keep'] = logits_to_keep
+                if self.args.tuner_backend == 'unsloth':
+                    inputs['logits_to_keep'] = logits_to_keep.sum()
 
         inputs['compute_loss_func'] = compute_loss_func
         inputs['loss_kwargs'] = loss_kwargs