update

examples/tutorials/rlhf/gpt2_sst2_ppo/step_2_train_model_two_gpu.py

@@ -105,7 +105,7 @@ def main():
         args.sft_model_name,
         num_labels=1
     )
-    value_model.transformer = model.transformer
+    # value_model.transformer = model.transformer
 
     ref_model = AutoModelForCausalLM.from_pretrained(args.sft_model_name)
     reward_model = AutoModelForSequenceClassification.from_pretrained(
@@ -132,10 +132,10 @@ def main():
         remove_columns=valid_dataset.column_names,
     )
     train_dataset = train_dataset.filter(
-        function=lambda x: 8 < len(x["input_ids"]) < 1024
+        function=lambda x: 8 < len(x["input_ids"]) < 512
     )
     valid_dataset = valid_dataset.filter(
-        function=lambda x: 8 < len(x["input_ids"]) < 1024
+        function=lambda x: 8 < len(x["input_ids"]) < 512
     )
     train_dataset = train_dataset.map(
         lambda example: token_truncate(example, tokenizer),
@@ -164,7 +164,7 @@ def main():
         num_mini_batches=2,
         num_sample_generations=100,
         # total_episodes=100000,  # 最多训练多少个样本。
-        response_length=1024,
+        response_length=64,
         # stop_token=tokenizer.eos_token, # stop_token 和 stop_token_id 只设置一个。
         stop_token_id=tokenizer.eos_token_id,
         batch_size=16,