fix (#5147)

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

Author: 2 people

swift/cli/sft.py

@@ -1,6 +1,9 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 import os
 
+if int(os.environ.get('UNSLOTH_PATCH_TRL', '0')) != 0:
+    import unsloth
+
 from swift.llm import sft_main
 
 if __name__ == '__main__':

swift/llm/model/register.py

@@ -139,31 +139,35 @@ def load_by_unsloth(args):
 
     @contextmanager
     def _patch_distributed_function():
-        from unsloth_zoo import utils
+        from unsloth_zoo import utils, compiler
 
         def distributed_function(n=1, function=None, *args, **kwargs):
             return function(*args, **kwargs)
 
         _origin_distributed_function = utils.distributed_function
         utils.distributed_function = distributed_function
+        compiler.distributed_function = distributed_function
         yield
         utils.distributed_function = _origin_distributed_function
+        compiler.distributed_function = _origin_distributed_function
 
     with _patch_distributed_function():
         if model_meta.is_multimodal:
             from unsloth import FastVisionModel as UnslothModel
+        elif model_info.is_moe_model:
+            from unsloth import FastModel 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,
-        max_seq_length=args.max_length,
-        full_finetuning=args.train_type == 'full',
-        load_in_4bit=args.quant_bits == 4,
-        load_in_8bit=args.quant_bits == 8,
-        device_map=args.device_map,
-    )
+        model, processor = UnslothModel.from_pretrained(
+            model_name=args.adapters and args.adapters[0] or args.model_dir,
+            dtype=args.torch_dtype,
+            max_seq_length=args.max_length,
+            full_finetuning=args.train_type == 'full',
+            load_in_4bit=args.quant_bits == 4,
+            load_in_8bit=args.quant_bits == 8,
+            device_map=args.device_map,
+        )
     if isinstance(model, PeftModel):
         base_model = model.model
     else:

swift/trainers/trainers.py

@@ -333,7 +333,7 @@ def _prepare_inputs(self, inputs):
             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['logits_to_keep'] = int(logits_to_keep.sum())
 
         inputs['compute_loss_func'] = compute_loss_func
         inputs['loss_kwargs'] = loss_kwargs
@@ -387,8 +387,10 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         if getattr(self.args, 'average_tokens_across_devices', False) and self.model_accepts_loss_kwargs:
             loss *= self.accelerator.num_processes
 
-        if outputs.logits is not None and labels is not None and not return_outputs:
+        if (outputs.logits is not None and labels is not None and not return_outputs
+                and self.args.tuner_backend != 'unsloth'):
             # Liger does not have logits
+            # Unsloth has a bug with output logits
             self._compute_acc(outputs, labels)
         return (loss, outputs) if return_outputs else loss
 

swift/tuners/base.py

@@ -526,7 +526,7 @@ def mark_trainable_callback(model, cfg):
 
     def save_pretrained(self,
                         save_directory: str,
-                        safe_serialization: bool = False,
+                        safe_serialization: bool = True,
                         adapter_name: Union[str, List[str]] = None,
                         **kwargs):
         """Save the adapters to a local directory.