Upgrade to TRL 0.17.0 version

examples/trl/README.md

@@ -21,7 +21,7 @@ $ pip install -U -r requirements.txt
         --subset '' \
         --output_dir ./model_qwen \
         --num_train_epochs 1 \
-        --per_device_train_batch_size 16 \
+        --per_device_train_batch_size 4 \
         --eval_strategy "no" \
         --save_strategy "no" \
         --learning_rate 3e-4 \
@@ -39,15 +39,15 @@ $ pip install -U -r requirements.txt
         --lora_alpha=16 \
         --lora_dropout=0.05 \
         --lora_target_modules "q_proj" "v_proj" "k_proj" "o_proj" \
-        --max_seq_length 512 \
+        --max_length 512 \
         --adam_epsilon 1e-08 \
         --use_flash_attention
     ```
 
-2. Supervised fine-tuning of the mistralai/Mixtral-8x7B-Instruct-v0.1 on 4 cards:
+2. Supervised fine-tuning of the mistralai/Mixtral-8x7B-Instruct-v0.1 on 8 cards:
 
     ```bash
-    PT_HPU_LAZY_MODE=1 DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED=1 PT_ENABLE_INT64_SUPPORT=1 python ../gaudi_spawn.py --world_size 4 --use_deepspeed sft.py \
+    PT_HPU_LAZY_MODE=1 DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED=1 python ../gaudi_spawn.py --world_size 8 --use_deepspeed sft.py \
         --model_name_or_path mistralai/Mixtral-8x7B-Instruct-v0.1 \
         --dataset_name "philschmid/dolly-15k-oai-style" \
         --subset 'data/' \
@@ -69,7 +69,7 @@ $ pip install -U -r requirements.txt
         --lora_target_modules "q_proj" "v_proj" \
         --bf16 \
         --remove_unused_columns=False \
-        --max_seq_length 512 \
+        --max_length 512 \
         --run_name="sft_mixtral" \
         --report_to=none \
         --use_habana \

examples/trl/dpo.py

@@ -1,4 +1,4 @@
-# copy from https://github.com/huggingface/trl/blob/v0.7.6/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py, enable it for Gaudi2
+# copy from https://github.com/huggingface/trl/blob/v0.17.0/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py, enable it for Gaudi2
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional
 
@@ -162,14 +162,18 @@ def return_prompt_and_responses(samples) -> Dict[str, str]:
         bf16=True,
         remove_unused_columns=False,
         run_name="dpo_llama2",
-        gradient_checkpointing_kwargs={"use_reentrant": script_args.gradient_checkpointing_use_reentrant},
+        gradient_checkpointing_kwargs=dict(use_reentrant=script_args.gradient_checkpointing_use_reentrant),
         use_habana=True,
         use_lazy_mode=True,
         use_hpu_graphs_for_training=not script_args.gradient_checkpointing and (not script_args.deepspeed),
         use_hpu_graphs_for_inference=not script_args.deepspeed,
         seed=script_args.seed,
         deepspeed=script_args.deepspeed,
         overwrite_output_dir=True,
+        beta=script_args.beta,
+        max_prompt_length=script_args.max_prompt_length,
+        max_length=script_args.max_length,
+        force_use_ref_model=True,
     )
 
     # Set seed before initializing model.
@@ -243,14 +247,10 @@ def return_prompt_and_responses(samples) -> Dict[str, str]:
         model_ref,
         gaudi_config=gaudi_config,
         args=training_args,
-        beta=script_args.beta,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
         peft_config=peft_config,
-        max_prompt_length=script_args.max_prompt_length,
-        max_length=script_args.max_length,
-        force_use_ref_model=True,
     )
 
     # 6. train