Fix deepspeed (#2778)

Jintao-Huang · web-flow · commit e471074f6427 · 2024-12-26T20:21:36.000+08:00
diff --git a/docs/source/Instruction/命令行参数.md b/docs/source/Instruction/命令行参数.md
@@ -87,7 +87,7 @@
 
 - 🔥output_dir: 默认为`output/<model_name>`
 - 🔥gradient_checkpointing: 是否使用gradient_checkpointing，默认为True
-- 🔥deepspeed: 默认为None. 可以设置为'zero2', 'zero3', 'zero2_offload', 'zero3_offload'来使用ms-swift内置的deepspeed配置文件
+- 🔥deepspeed: 默认为None. 可以设置为'zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'来使用ms-swift内置的deepspeed配置文件
 - 🔥per_device_train_batch_size: 默认值1
 - 🔥per_device_eval_batch_size: 默认值1
 - weight_decay: weight衰减系数，默认值0.1
diff --git a/docs/source_en/Instruction/Command-line-parameters.md b/docs/source_en/Instruction/Command-line-parameters.md
@@ -88,7 +88,7 @@ This parameter list inherits from transformers `Seq2SeqTrainingArguments`, with
 
 - 🔥output_dir: Default is `output/<model_name>`.
 - 🔥gradient_checkpointing: Whether to use gradient checkpointing, default is True.
-- 🔥deepspeed: Default is None. Can be set to 'zero2', 'zero3', 'zero2_offload', 'zero3_offload' to use the built-in deepspeed configuration files from ms-swift.
+- 🔥deepspeed: Default is None. Can be set to 'zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload' to use the built-in deepspeed configuration files from ms-swift.
 - 🔥per_device_train_batch_size: Default is 1.
 - 🔥per_device_eval_batch_size: Default is 1.
 - weight_decay: Weight decay coefficient, default value is 0.1.
diff --git a/swift/llm/argument/train_args.py b/swift/llm/argument/train_args.py
@@ -164,7 +164,10 @@ def _init_deepspeed(self):
                                  f'local_world_size: {self.local_world_size}.')
 
             ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'ds_config'))
-            deepspeed_mapping = {name: f'{name}.json' for name in ['zero2', 'zero3', 'zero2_offload', 'zero3_offload']}
+            deepspeed_mapping = {
+                name: f'{name}.json'
+                for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload']
+            }
             for ds_name, ds_config in deepspeed_mapping.items():
                 if self.deepspeed == ds_name:
                     self.deepspeed = os.path.join(ds_config_folder, ds_config)
diff --git a/swift/llm/argument/tuner_args.py b/swift/llm/argument/tuner_args.py
@@ -222,4 +222,4 @@ def _init_multimodal_full(self):
         if self.freeze_parameters:
             logger.info(f'freeze_parameters: {self.freeze_parameters}')
         if self.trainable_parameters:
-            logger.info(f'trainable_parameters: {self.trainable_parameters}')
+            logger.info(f'additional trainable_parameters: {self.trainable_parameters}')
diff --git a/swift/llm/ds_config/zero0.json b/swift/llm/ds_config/zero0.json
@@ -0,0 +1,31 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 0,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/swift/llm/ds_config/zero1.json b/swift/llm/ds_config/zero1.json
@@ -0,0 +1,35 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 1,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
diff --git a/swift/llm/ds_config/zero2.json b/swift/llm/ds_config/zero2.json
@@ -12,24 +12,6 @@
         "enabled": "auto"
     },
 
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",
-            "betas": "auto",
-            "eps": "auto",
-            "weight_decay": "auto"
-        }
-    },
-
-    "scheduler": {
-        "type": "WarmupCosineLR",
-        "params": {
-            "total_num_steps": "auto",
-            "warmup_num_steps": "auto"
-        }
-    },
-
     "zero_optimization": {
         "stage": 2,
         "offload_optimizer": {
diff --git a/swift/llm/ds_config/zero2_offload.json b/swift/llm/ds_config/zero2_offload.json
@@ -12,24 +12,6 @@
         "enabled": "auto"
     },
 
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",
-            "betas": "auto",
-            "eps": "auto",
-            "weight_decay": "auto"
-        }
-    },
-
-    "scheduler": {
-        "type": "WarmupCosineLR",
-        "params": {
-            "total_num_steps": "auto",
-            "warmup_num_steps": "auto"
-        }
-    },
-
     "zero_optimization": {
         "stage": 2,
         "offload_optimizer": {
diff --git a/swift/llm/ds_config/zero3.json b/swift/llm/ds_config/zero3.json
@@ -12,24 +12,6 @@
         "enabled": "auto"
     },
 
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",
-            "betas": "auto",
-            "eps": "auto",
-            "weight_decay": "auto"
-        }
-    },
-
-    "scheduler": {
-        "type": "WarmupCosineLR",
-        "params": {
-            "total_num_steps": "auto",
-            "warmup_num_steps": "auto"
-        }
-    },
-
     "zero_optimization": {
         "stage": 3,
         "offload_optimizer": {
diff --git a/swift/llm/ds_config/zero3_offload.json b/swift/llm/ds_config/zero3_offload.json
@@ -12,24 +12,6 @@
         "enabled": "auto"
     },
 
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",
-            "betas": "auto",
-            "eps": "auto",
-            "weight_decay": "auto"
-        }
-    },
-
-    "scheduler": {
-        "type": "WarmupCosineLR",
-        "params": {
-            "total_num_steps": "auto",
-            "warmup_num_steps": "auto"
-        }
-    },
-
     "zero_optimization": {
         "stage": 3,
         "offload_optimizer": {
diff --git a/swift/trainers/trainers.py b/swift/trainers/trainers.py
@@ -149,8 +149,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             labels = inputs['labels']
             # fix https://github.com/huggingface/transformers/issues/34263
             if num_items_in_batch is not None:
-                if getattr(self.args, 'average_tokens_across_devices', False):
-                    outputs.loss *= self.accelerator.num_processes
                 outputs.loss = outputs.loss * (labels[:, 1:] != -100).sum() / num_items_in_batch
 
             if isinstance(outputs, dict) and 'loss' not in outputs:
diff --git a/swift/ui/llm_train/advanced.py b/swift/ui/llm_train/advanced.py
@@ -159,6 +159,6 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                         scale=20,
                         allow_custom_value=True,
                         value=None,
-                        choices=['zero2', 'zero3', 'zero2_offload', 'zero3_offload'])
+                        choices=['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'])
                 with gr.Row():
                     gr.Textbox(elem_id='more_params', lines=4, scale=20)
diff --git a/swift/utils/utils.py b/swift/utils/utils.py
@@ -37,7 +37,7 @@ def check_json_format(obj: Any, token_safe: bool = True) -> Any:
     elif isinstance(obj, Mapping):
         res = {}
         for k, v in obj.items():
-            if token_safe and isinstance(k, str) and '_token' in k:
+            if token_safe and isinstance(k, str) and '_token' in k and isinstance(v, str):
                 res[k] = None
             else:
                 res[k] = check_json_format(v, token_safe)
