update (modelscope#1689)

kronoszhang · Aug 13, 2024 · 5297d8e · 5297d8e
1 parent bb208d2
commit 5297d8e
Show file tree

Hide file tree

Showing 9 changed files with 42 additions and 10 deletions.
diff --git a/docs/source/LLM/ORPO算法最佳实践.md b/docs/source/LLM/ORPO算法最佳实践.md
@@ -46,8 +46,11 @@ swift内置了处理方法将`answer_zh`作为`response`,将`answer_en`作为`re
 # Experimental environment: A100
 # DDP + MP
 # Memory usage: 4*24G
+nproc_per_node=2
+
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-NPROC_PER_NODE=2 \
+NPROC_PER_NODE=$nproc_per_node \
+MASTER_PORT=29500 \
 swift rlhf \
     --rlhf_type orpo \
     --model_type  llama3-8b-instruct \

diff --git a/docs/source/Multi-Modal/人类偏好对齐训练文档.md b/docs/source/Multi-Modal/人类偏好对齐训练文档.md
@@ -91,8 +91,11 @@ swift rlhf \
     --save_total_limit  2
 
 # DDP + MP
+nproc_per_node=2
+
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-NPROC_PER_NODE=2 \
+NPROC_PER_NODE=$nproc_per_node \
+MASTER_PORT=29500 \
 swift rlhf \
     --rlhf_type dpo \
     --model_type llava1_6-mistral-7b-instruct \

diff --git a/docs/source_en/LLM/Human-Preference-Alignment-Training-Documentation.md b/docs/source_en/LLM/Human-Preference-Alignment-Training-Documentation.md
@@ -90,8 +90,11 @@ swift rlhf \
 
 # DDP + MP
 # Memory usage: 4*24G
+nproc_per_node=2
+
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-NPROC_PER_NODE=2 \
+NPROC_PER_NODE=$nproc_per_node \
+MASTER_PORT=29500 \
 swift rlhf \
     --rlhf_type dpo \
     --model_type  llama3-8b-instruct \

diff --git a/docs/source_en/LLM/ORPO.md b/docs/source_en/LLM/ORPO.md
@@ -42,8 +42,11 @@ Swift has built-in methods for processing this dataset, using `answer_zh` as `re
 # Experimental environment: A100
 # DDP + MP
 # Memory usage: 4*24G
+nproc_per_node=2
+
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-NPROC_PER_NODE=2 \
+NPROC_PER_NODE=$nproc_per_node \
+MASTER_PORT=29500 \
 swift rlhf \
     --rlhf_type orpo \
     --model_type  llama3-8b-instruct \

diff --git a/docs/source_en/Multi-Modal/human-preference-alignment-training-documentation.md b/docs/source_en/Multi-Modal/human-preference-alignment-training-documentation.md
@@ -90,8 +90,11 @@ swift rlhf \
     --save_total_limit  2
 
 # DDP + MP
+nproc_per_node=2
+
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
-NPROC_PER_NODE=2 \
+NPROC_PER_NODE=$nproc_per_node \
+MASTER_PORT=29500 \
 swift rlhf \
     --rlhf_type dpo \
     --model_type llava1_6-mistral-7b-instruct \

diff --git a/swift/trainers/cpo_trainer.py b/swift/trainers/cpo_trainer.py
@@ -11,7 +11,6 @@
 from .utils import build_tokenized_answer, patch_trl, sort_by_max_length
 
 logger = get_logger()
-patch_trl()
 
 
 class CPOTrainer(PushToMsHubMixin, SwiftMixin, HFCPOTrainer):
@@ -21,6 +20,7 @@ def __init__(self, *args, template: Template, test_oom_error=False, **kwargs):
         kwargs.pop('gamma', None)
         self.streaming = kwargs.pop('streaming')
         is_vision = kwargs.pop('is_vision')
+        patch_trl(is_vision)
         self.keys = []  # keys appears in tokenize_row
         self.column_names = list(next(iter(kwargs.get('train_dataset'))).keys())
         self.need_filter: bool = False

diff --git a/swift/trainers/dpo_trainer.py b/swift/trainers/dpo_trainer.py
@@ -12,7 +12,6 @@
 from .utils import build_tokenized_answer, patch_trl, sort_by_max_length
 
 logger = get_logger()
-patch_trl()
 
 
 class DPOTrainer(PushToMsHubMixin, SwiftMixin, HFDPOTrainer):
@@ -22,6 +21,7 @@ def __init__(self, *args, template: Template, sft_beta=0., test_oom_error=False,
         self.sft_beta = sft_beta
         self.streaming = kwargs.pop('streaming')
         is_vision = kwargs.pop('is_vision')
+        patch_trl(is_vision)
         self.keys = []  # keys appears in tokenize_row
         self.column_names = list(next(iter(kwargs.get('train_dataset'))).keys())
         self.need_filter: bool = False

diff --git a/swift/trainers/orpo_trainer.py b/swift/trainers/orpo_trainer.py
@@ -11,7 +11,6 @@
 from .utils import build_tokenized_answer, patch_trl, sort_by_max_length
 
 logger = get_logger()
-patch_trl()
 
 
 class ORPOTrainer(PushToMsHubMixin, SwiftMixin, HFORPOTrainer):
@@ -20,6 +19,7 @@ def __init__(self, *args, template: Template, test_oom_error=False, **kwargs):
         self.template = template
         self.streaming = kwargs.pop('streaming')
         is_vision = kwargs.pop('is_vision')
+        patch_trl(is_vision)
         self.keys = []
         self.column_names = list(next(iter(kwargs.get('train_dataset'))).keys())
         self.need_filter: bool = False

diff --git a/swift/trainers/utils.py b/swift/trainers/utils.py
@@ -3,6 +3,7 @@
 
 import heapq
 import inspect
+from functools import partial
 from types import FunctionType, MethodType
 from typing import Dict, List, Optional, Union
 
@@ -120,7 +121,7 @@ def sort_by_max_length(dataset: HfDataset, num_dataset: int, is_encoder_decoder:
     return dataset.select(idx)
 
 
-def patch_trl():
+def patch_trl(is_vision_model: bool = False):
     from .callback import DefaultFlowCallbackNew, PrinterCallbackNew, ProgressCallbackNew
     from transformers import trainer
 
@@ -129,7 +130,10 @@ def patch_trl():
     trainer.PrinterCallback = PrinterCallbackNew
 
     # fix encoder-decoder error
-    patch_datacollator()
+    if is_vision_model:
+        patch_datacollator()
+        patch_dataset_map()
+
     patch_itds_map()
 
 
@@ -235,3 +239,16 @@ def new_map(self, *args, **kwargs):
         IterableDataset.map = new_map
         IterableDataset._old_map = old_map
         # model.forward = MethodType(_patch_ids_map(map_func), IterableDataset)
+
+
+def patch_dataset_map():
+    original_map = HfDataset.map
+    if not hasattr(HfDataset, '_old_map'):
+
+        def patched_map(self, function, **kwargs):
+            if 'writer_batch_size' not in kwargs:
+                kwargs['writer_batch_size'] = 10
+            return original_map(self, function, **kwargs)
+
+        HfDataset.map = patched_map
+        HfDataset._old_map = original_map