deepspeed precision

Signed-off-by: ssbuild <[email protected]>

Author: ssbuild

config/reward_config/main.py

@@ -24,6 +24,8 @@
 if global_args['quantization_config'] is not None:
     global_args['quantization_config'].load_in_4bit = load_in_bit == 4
     global_args['quantization_config'].load_in_8bit = load_in_bit == 8
+    if load_in_bit == 0:
+        global_args["quantization_config"] = None
 
 if enable_lora:
     enable_ptv2 = False
@@ -49,15 +51,15 @@
     train_info_args['use_fast_tokenizer'] = True
 
 
-def get_deepspeed_config():
+def get_deepspeed_config(precision='fp16'):
     '''
         lora prompt finetuning 使用 deepspeed_offload.json
         普通finetuning 使用deepspeed.json
     '''
     # 是否开启deepspeed
     if not enable_deepspeed:
         return None
-
+    precision = str(precision).lower()
     # 选择 deepspeed 配置文件
     is_need_update_config = False
     if enable_lora:
@@ -79,5 +81,17 @@ def get_deepspeed_config():
             optimizer['params']['eps'] = train_info_args.get('adam_epsilon', 1e-8)
             # deepspeed_offload 优化器有效
             train_info_args['optimizer'] = optimizer['type']
+
+    if precision == 'bf16':
+        if 'fp16' in deepspeed_config:
+            deepspeed_config["fp16"]["enbale"] = False
+        if 'bf16' in deepspeed_config:
+            deepspeed_config["bf16"]["enbale"] = True
+        else:
+            deepspeed_config['bf16'] = {"enbale": True}
+    elif precision == 'fp16':
+        if 'bf16' in deepspeed_config:
+            deepspeed_config["bf16"]["enbale"] = False
+
     return deepspeed_config
 

config/rlhf_config/main.py

@@ -22,6 +22,8 @@
 if global_args['quantization_config'] is not None:
     global_args['quantization_config'].load_in_4bit = load_in_bit == 4
     global_args['quantization_config'].load_in_8bit = load_in_bit == 8
+    if load_in_bit == 0:
+        global_args["quantization_config"] = None
 
 if enable_lora:
     enable_ptv2 = False
@@ -47,15 +49,15 @@
     train_info_args['use_fast_tokenizer'] = True
 
 
-def get_deepspeed_config():
+def get_deepspeed_config(precision='fp16'):
     '''
         lora prompt finetuning 使用 deepspeed_offload.json
         普通finetuning 使用deepspeed.json
     '''
     # 是否开启deepspeed
     if not enable_deepspeed:
         return None
-
+    precision = str(precision).lower()
     # 选择 deepspeed 配置文件
     is_need_update_config = False
     if enable_lora:
@@ -78,5 +80,17 @@ def get_deepspeed_config():
 
             # deepspeed_offload 优化器有效
             train_info_args['optimizer'] = optimizer['type']
+
+    if precision == 'bf16':
+        if 'fp16' in deepspeed_config:
+            deepspeed_config["fp16"]["enbale"] = False
+        if 'bf16' in deepspeed_config:
+            deepspeed_config["bf16"]["enbale"] = True
+        else:
+            deepspeed_config['bf16'] = {"enbale": True}
+    elif precision == 'fp16':
+        if 'bf16' in deepspeed_config:
+            deepspeed_config["bf16"]["enbale"] = False
+
     return deepspeed_config
 

stage2_reward/train.py

@@ -34,9 +34,16 @@
 
     dataHelper.make_dataset_all()
 
+    is_bf16_supported = torch.cuda.is_bf16_supported()
+    # 精度 根据实际情况做调整
+    if is_bf16_supported:
+        precision = 'bf16'
+    else:
+        precision = '16'
 
-
-    deepspeed_config = get_deepspeed_config()
+    if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
+        precision = "32"
+    deepspeed_config = get_deepspeed_config(precision)
     strategy = 'ddp' if torch.cuda.device_count() > 1 else 'auto'
     if deepspeed_config is not None and len(deepspeed_config):
         strategy = DeepSpeedStrategy(config=deepspeed_config, )
@@ -56,15 +63,7 @@
     if deepspeed_config is not None and len(deepspeed_config):
         strategy = DeepSpeedStrategy(config=deepspeed_config, )
 
-    is_bf16_supported = torch.cuda.is_bf16_supported()
-    # 精度 根据实际情况做调整
-    if is_bf16_supported:
-        precision = 'bf16'
-    else:
-        precision = '16'
 
-    if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
-        precision = "32"
 
     trainer = Trainer(
         callbacks=[checkpoint_callback, LearningRateMonitor(logging_interval='step')],

stage3_rlhf/train.py

@@ -36,7 +36,16 @@
 
     dataHelper.make_dataset_all()
 
-    deepspeed_config = get_deepspeed_config()
+    is_bf16_supported = torch.cuda.is_bf16_supported()
+    # 精度 根据实际情况做调整
+    if is_bf16_supported:
+        precision = 'bf16'
+    else:
+        precision = '16'
+
+    if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
+        precision = "32"
+    deepspeed_config = get_deepspeed_config(precision)
     strategy = 'ddp' if torch.cuda.device_count() >= 1 else 'auto'
     if deepspeed_config is not None and len(deepspeed_config):
         strategy = DeepSpeedStrategy(config=deepspeed_config, )
@@ -53,15 +62,7 @@
         training_args=training_args,
         lora_args=lora_args, )
 
-    is_bf16_supported = torch.cuda.is_bf16_supported()
-    # 精度 根据实际情况做调整
-    if is_bf16_supported:
-        precision = 'bf16'
-    else:
-        precision = '16'
 
-    if global_args["quantization_config"] is not None and global_args["quantization_config"].load_in_8bit:
-        precision = "32"
     trainer = PPOTrainer(
         callbacks=[ checkpoint_callback],
         max_epochs=training_args.max_epochs,