update requirement

requirements.txt

@@ -1,19 +1,17 @@
-accelerate==0.29.3
-transformers==4.39.3
-peft==0.10.0
-bitsandbytes==0.43.1
-loguru==0.7.2
-numpy==1.26.4
-pandas==2.2.1
-tqdm==4.66.4
-deepspeed==0.14.2
-sentencepiece==0.2.0
-transformers-stream-generator==0.0.5
-tiktoken==0.6.0
-einops==0.8.0
-httpx
-scipy
-torch==2.2.2
-datasets==2.16.1
+accelerate
+transformers==4.44.2
+peft
+bitsandbytes
+loguru
+numpy
+pandas
+tqdm
+deepspeed==0.13.1
+sentencepiece
+transformers-stream-generator
+tiktoken
+einops
+torch==2.3.0
+datasets
 trl
 wandb

rlhf/rlhf_train.py

@@ -158,13 +158,13 @@ def main():
             llm_int8_has_fp16_weight=False,
         )
         model_kwargs.update(quantization_config=quantization_config)
-
-    # 如果模型不支持AutoModelForSequenceClassification需要在对应config文件中添加映射
-    try:
-        reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1,
-                                                                          **model_kwargs)
-    except Exception as e:
-        assert False, "模型不支持AutoModelForSequenceClassification需要在对应config文件中添加映射"
+    if args.rlhf_type in ['PPO', 'RLOO']:
+        # 如果模型不支持AutoModelForSequenceClassification需要在对应config文件中添加映射
+        try:
+            reward_model = AutoModelForSequenceClassification.from_pretrained(config.reward_model_path, num_labels=1,
+                                                                              **model_kwargs)
+        except Exception as e:
+            assert False, "模型不支持AutoModelForSequenceClassification需要在对应config文件中添加映射"
 
     ref_policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path, **model_kwargs)
     policy = AutoModelForCausalLM.from_pretrained(config.sft_model_path, **model_kwargs)