incorporating the review

Author: schoi-habana

examples/trl/README.md

@@ -41,7 +41,7 @@ PT_HPU_MAX_COMPOUND_OP_SIZE=10 PT_HPU_LAZY_MODE=1 python3 grpo.py \
     --num_generations 4 \
     --max_completion_length 64 \
     --use_peft True \
-    --lora_target_modules q_proj, k_proj \
+    --lora_target_modules q_proj k_proj \
     --num_train_epochs 1 \
     --save_strategy="epoch"
 ```
@@ -66,7 +66,7 @@ PT_HPU_MAX_COMPOUND_OP_SIZE=10 PT_HPU_LAZY_MODE=1 python3 ../gaudi_spawn.py --wo
     --num_generations 4 \
     --max_completion_length 64 \
     --use_peft True \
-    --lora_target_modules q_proj, k_proj \
+    --lora_target_modules q_proj k_proj \
     --max_steps=500 \
     --logging_steps=10 \
     --save_steps=100

examples/trl/grpo.py

@@ -79,7 +79,7 @@ class ScriptArguments:
     model_name_or_path: Optional[str] = field(default="Qwen/Qwen2-0.5B-Instruct", metadata={"help": "the model name"})
     dataset_name: Optional[str] = field(default=None, metadata={"help": "the dataset name"})
     use_peft: Optional[bool] = field(default=False, metadata={"help": "whether to use peft"})
-    num_workers: Optional[int] = field(default=4, metadata={"help": "the number of workers"})
+    num_workers: Optional[int] = field(default=1, metadata={"help": "the number of workers"})
     subset: Optional[str] = field(default=None, metadata={"help": "the subset to use"})
     streaming: Optional[bool] = field(default=False, metadata={"help": "whether to stream the dataset"})
     dataset_train_split: str = field(default="train[:5%]", metadata={"help": "Dataset split to use for training."})

examples/trl/requirements_grpo.txt

@@ -1,6 +1,6 @@
 trl == 0.17.0
 peft == 0.12.0
-datasets == 3.0.0
+datasets
 tyro
 evaluate
 scikit-learn == 1.5.2

optimum/habana/trl/trainer/grpo_trainer.py

@@ -49,11 +49,11 @@
     selective_log_softmax,
 )
 
-from optimum.habana.transformers import trainer as habana_trainer
-from optimum.habana.transformers.trainer import _get_input_update_settings
 from optimum.utils import logging
 
 from ... import GaudiConfig, GaudiTrainer
+from ...transformers import trainer as habana_trainer
+from ...transformers.trainer import _get_input_update_settings
 from .grpo_config import GaudiGRPOConfig