Enable trl GRPO trainer

Makefile

@@ -159,11 +159,15 @@ slow_tests_video_llava_example: test_installs
 slow_tests_fsdp: test_installs
 	python -m pytest tests/test_fsdp_examples.py -v -s --token $(TOKEN)
 
-slow_tests_trl: test_installs
+slow_tests_trl_ddpo: test_installs
 	python -m pip install trl==0.9.6
 	python -m pip install peft==0.12.0
 	python -m pytest tests/test_trl.py -v -s -k "test_calculate_loss"
 
+slow_tests_trl_grpo: test_installs
+	python -m pip install -r examples/trl/requirements_grpo.txt
+	python -m pytest tests/test_trl.py -v -s -k "GaudiGRPOTrainerTester"
+
 slow_tests_object_segmentation: test_installs
 	python -m pytest tests/test_object_segmentation.py
 

examples/trl/README.md

@@ -4,10 +4,74 @@
 ## Requirements
 
 First, you should install the requirements:
+
+- For **GRPO example**:
+```bash
+$ pip install -U -r requirements_grpo.txt
+```
+
+- For **all other examples**:
 ```bash
 $ pip install -U -r requirements.txt
 ```
 
+## GRPO Training
+
+Installing DeepSpeed
+
+```sh
+pip install git+https://github.com/HabanaAI/[email protected]
+```
+
+Running single card training
+
+```sh
+PT_HPU_MAX_COMPOUND_OP_SIZE=10 PT_HPU_LAZY_MODE=1 python3 grpo.py \
+    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
+    --dataset_name AI-MO/NuminaMath-TIR \
+    --per_device_train_batch_size 8 \
+    --per_device_eval_batch_size 8 \
+    --do_train \
+    --do_eval \
+    --use_habana \
+    --use_lazy_mode \
+    --bf16 True \
+    --gradient_accumulation_steps=16 \
+    --max_prompt_length 512 \
+    --num_generations 4 \
+    --max_completion_length 64 \
+    --use_peft True \
+    --lora_target_modules q_proj k_proj \
+    --num_train_epochs 1 \
+    --save_strategy="epoch"
+```
+
+
+Runnig multi-card training
+
+```sh
+PT_HPU_MAX_COMPOUND_OP_SIZE=10 PT_HPU_LAZY_MODE=1 python3 ../gaudi_spawn.py --world_size 8 --use_deepspeed grpo.py \
+    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
+    --dataset_name AI-MO/NuminaMath-TIR \
+    --per_device_train_batch_size 8 \
+    --per_device_eval_batch_size 8 \
+    --do_train \
+    --do_eval \
+    --use_habana \
+    --use_lazy_mode \
+    --bf16 True \
+    --gradient_accumulation_steps=16 \
+    --gradient_checkpointing \
+    --max_prompt_length 512 \
+    --num_generations 4 \
+    --max_completion_length 64 \
+    --use_peft True \
+    --lora_target_modules q_proj k_proj \
+    --max_steps=500 \
+    --logging_steps=10 \
+    --save_steps=100
+```
+
 ## Supervised Finetuning
 
 1. The following example is for the supervised Lora finetune with Qwen2 model for conversational format dataset.

examples/trl/grpo.py

@@ -0,0 +1,210 @@
+import contextlib
+import io
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+import transformers
+from datasets import load_dataset
+from math_verify import LatexExtractionConfig, parse, verify
+from peft import LoraConfig
+from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser
+from transformers.integrations.deepspeed import (
+    is_deepspeed_available,
+)
+from transformers.trainer_utils import is_main_process
+
+from optimum.habana import GaudiConfig
+from optimum.habana.trl import GaudiGRPOConfig, GaudiGRPOTrainer
+from optimum.habana.utils import set_seed
+
+
+logger = logging.getLogger(__name__)
+SYSTEM_PROMPT = (
+    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
+    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
+    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
+    "<think> reasoning process here </think><answer> answer here </answer>"
+)
+
+
+def make_conversation(example):
+    return {
+        "prompt": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+        ],
+    }
+
+
+ideal_length = 50
+
+
+def reward_len(completions, **kwargs):
+    return [-abs(ideal_length - len(completion)) for completion in completions]  # penalize response when len!=50
+
+
+def format_reward(completions, **kwargs):
+    # Checks if the reasoning process is enclosed within <think> and </think> tags,
+    # while the final answer is enclosed within <answer> and </answer> tags.
+    pattern = r"^<think>.*?</think>\s*<answer>.*?</answer>$"
+    completion_contents = [completion[0]["content"] for completion in completions]
+    matches = [re.match(pattern, content) for content in completion_contents]
+    return [1.0 if match else 0.0 for match in matches]
+
+
+def accuracy_reward(completions, **kwargs):
+    # Checks if the completion is the same as the ground truth.
+    solutions = kwargs["solution"]
+    completion_contents = [completion[0]["content"] for completion in completions]
+    rewards = []
+    for content, solution in zip(completion_contents, solutions):
+        gold_parsed = parse(solution, extraction_mode="first_match", extraction_config=[LatexExtractionConfig()])
+        answer_parsed = parse(content, extraction_mode="first_match", extraction_config=[LatexExtractionConfig()])
+        if len(gold_parsed) != 0:
+            try:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+                    rewards.append(float(verify(answer_parsed, gold_parsed)))
+            except Exception:
+                rewards.append(0.0)
+        else:
+            rewards.append(1.0)
+    return rewards
+
+
+@dataclass
+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=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."})
+    dataset_test_split: str = field(default="test[:5%]", metadata={"help": "Dataset split to use for evaluation."})
+    reward_model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Reward model id of a pretrained model hosted inside a model repo on huggingface.co or "
+            "local path to a directory containing model weights saved using `PreTrainedModel.save_pretrained`."
+        },
+    )
+
+    use_flash_attention: Optional[bool] = field(
+        default=True, metadata={"help": "Whether to use Habana flash attention for fine-tuning."}
+    )
+    flash_attention_recompute: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to enable recompute in Habana flash attention for fine-tuning."}
+    )
+    flash_attention_causal_mask: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to enable causal mask in Habana flash attention for fine-tuning."}
+    )
+
+    # LoraConfig
+    lora_alpha: Optional[float] = field(default=32, metadata={"help": "the lora alpha parameter"})
+    lora_dropout: Optional[float] = field(default=0.1, metadata={"help": "the lora dropout parameter"})
+    lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"})
+    lora_target_modules: List[str] = field(
+        default_factory=lambda: None,
+        metadata={"help": "Target modules for the LoRA method."},
+    )
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser((GaudiGRPOConfig, ScriptArguments))
+    (training_args, script_args) = parser.parse_args_into_dataclasses()
+
+    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, "
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.bf16}"
+    )
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    if is_main_process(training_args.local_rank):
+        transformers.utils.logging.set_verbosity_info()
+        transformers.utils.logging.enable_default_handler()
+        transformers.utils.logging.enable_explicit_format()
+    logger.info(f"Training/evaluation parameters {training_args}")
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    use_deepspeed = training_args.world_size > 1
+
+    if script_args.use_peft:
+        peft_config = LoraConfig(
+            r=script_args.lora_r,
+            lora_alpha=script_args.lora_alpha,
+            lora_dropout=script_args.lora_dropout,
+            target_modules=script_args.lora_target_modules,
+            task_type="CAUSAL_LM",
+        )
+    else:
+        peft_config = None
+
+    tokenizer = AutoTokenizer.from_pretrained(script_args.model_name_or_path, trust_remote_code=True)
+    if training_args.chat_template is not None:
+        tokenizer.chat_template = training_args.chat_template
+
+    train_dataset, test_dataset = load_dataset(
+        script_args.dataset_name,
+        data_dir=None if script_args.subset == "None" else script_args.subset,
+        num_proc=script_args.num_workers if not script_args.streaming else None,
+        split=[script_args.dataset_train_split, script_args.dataset_test_split],
+    )
+
+    train_dataset = train_dataset.map(make_conversation)
+    test_dataset = test_dataset.map(make_conversation)
+    train_dataset = train_dataset.remove_columns(["messages", "problem"])
+
+    low_cpu_mem_usage = True
+    if is_deepspeed_available() and use_deepspeed:
+        from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+
+        if is_deepspeed_zero3_enabled():
+            low_cpu_mem_usage = False
+
+    model = AutoModelForCausalLM.from_pretrained(
+        script_args.model_name_or_path,
+        low_cpu_mem_usage=low_cpu_mem_usage,
+        torch_dtype=torch.bfloat16,
+    )
+
+    model.config.use_cache = False
+    if not script_args.use_flash_attention and (
+        script_args.flash_attention_recompute or script_args.flash_attention_recompute
+    ):
+        assert "Need to enable use_flash_attention"
+    model.generation_config.use_flash_attention = script_args.use_flash_attention
+    model.generation_config.flash_attention_recompute = script_args.flash_attention_recompute
+    model.generation_config.flash_attention_causal_mask = script_args.flash_attention_causal_mask
+
+    reward_funcs = [format_reward, accuracy_reward]
+    if script_args.reward_model_name_or_path:
+        reward_funcs = AutoModelForSequenceClassification.from_pretrained(
+            script_args.reward_model_name_or_path,
+            trust_remote_code=True,
+        )
+
+    if getattr(tokenizer, "pad_token", None) is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    gaudi_config = GaudiConfig()
+    gaudi_config.use_fused_adam = True
+    gaudi_config.use_fused_clip_norm = True
+
+    trainer = GaudiGRPOTrainer(
+        model=model,
+        reward_funcs=reward_funcs,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        processing_class=tokenizer,
+        gaudi_config=gaudi_config,
+        peft_config=peft_config,
+    )
+
+    trainer.train()
+
+    print("Done!")

examples/trl/requirements_grpo.txt

@@ -0,0 +1,8 @@
+trl == 0.17.0
+peft == 0.12.0
+datasets
+tyro
+evaluate
+scikit-learn == 1.5.2
+accelerate
+math_verify

optimum/habana/trl/__init__.py

@@ -1,10 +1,21 @@
+import importlib.metadata
+
+from packaging import version
+
 from .models.modeling_base import adapt_PreTrainedModelWrapper_to_gaudi
 from .models.modeling_sd_base import GaudiDefaultDDPOStableDiffusionPipeline
 from .trainer.ddpo_trainer import GaudiDDPOTrainer
 from .trainer.dpo_config import GaudiDPOConfig
 from .trainer.dpo_trainer import GaudiDPOTrainer
-from .trainer.ppo_config import GaudiPPOConfig
-from .trainer.ppo_trainer import GaudiPPOTrainer
+
+
+trl_version = importlib.metadata.version("trl")
+if version.parse(trl_version) < version.parse("0.17.0"):
+    from .trainer.ppo_config import GaudiPPOConfig
+    from .trainer.ppo_trainer import GaudiPPOTrainer
+else:
+    from .trainer.grpo_config import GaudiGRPOConfig
+    from .trainer.grpo_trainer import GaudiGRPOTrainer
 from .trainer.reward_trainer import GaudiRewardTrainer, RewardDataCollatorWithPadding
 from .trainer.sft_config import GaudiSFTConfig
 from .trainer.sft_trainer import GaudiSFTTrainer

optimum/habana/trl/trainer/__init__.py

@@ -16,13 +16,22 @@
 
 # There is a circular import in the PPOTrainer if we let isort sort these
 # isort: on
+import importlib.metadata
+from packaging import version
 
 from .sft_trainer import GaudiSFTTrainer
 from .dpo_trainer import GaudiDPOTrainer
-from .ppo_config import GaudiPPOConfig
-from .ppo_trainer import GaudiPPOTrainer
+
 from .reward_trainer import GaudiRewardTrainer, RewardDataCollatorWithPadding
 
 from .ddpo_trainer import GaudiDDPOTrainer
 from .dpo_config import GaudiDPOConfig
 from .sft_config import GaudiSFTConfig
+
+trl_version = importlib.metadata.version("trl")
+if version.parse(trl_version) < version.parse("0.17.0"):
+    from .ppo_config import GaudiPPOConfig
+    from .ppo_trainer import GaudiPPOTrainer
+else:
+    from .grpo_trainer import GaudiGRPOTrainer
+    from .grpo_config import GaudiGRPOConfig