-
-
Save jane-alesi/26e1fe151de52fe0c33efc46e97069ef to your computer and use it in GitHub Desktop.
GRPO Llama-1B
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # train_grpo.py | |
| # | |
| # automates the language model training process by standardizing responses into an XML chain-of-thought format, | |
| # modularizing reward functions, and optimizing configurations for efficient GRPO training on the GSM8K dataset. | |
| # See comments for further explanations. | |
| import re | |
| import torch | |
| from datasets import load_dataset, Dataset | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from peft import LoraConfig | |
| from trl import GRPOConfig, GRPOTrainer | |
| # Load and prep dataset | |
| SYSTEM_PROMPT = """ | |
| Respond in the following format: | |
| <reasoning> | |
| ... | |
| </reasoning> | |
| <answer> | |
| ... | |
| </answer> | |
| """ | |
| XML_COT_FORMAT = """\ | |
| <reasoning> | |
| {reasoning} | |
| </reasoning> | |
| <answer> | |
| {answer} | |
| </answer> | |
| """ | |
| def extract_xml_answer(text: str) -> str: | |
| answer = text.split("<answer>")[-1] | |
| answer = answer.split("</answer>")[0] | |
| return answer.strip() | |
| def extract_hash_answer(text: str) -> str | None: | |
| if "####" not in text: | |
| return None | |
| return text.split("####")[1].strip() | |
| # uncomment middle messages for 1-shot prompting | |
| def get_gsm8k_questions(split = "train") -> Dataset: | |
| data = load_dataset('openai/gsm8k', 'main')[split] # type: ignore | |
| data = data.map(lambda x: { # type: ignore | |
| 'prompt': [ | |
| {'role': 'system', 'content': SYSTEM_PROMPT}, | |
| #{'role': 'user', 'content': 'What is the largest single-digit prime number?'}, | |
| #{'role': 'assistant', 'content': XML_COT_FORMAT.format( | |
| # reasoning="9 is divisble by 3 and 8 is divisible by 2, but 7 is prime.", | |
| # answer="7" | |
| #)}, | |
| {'role': 'user', 'content': x['question']} | |
| ], | |
| 'answer': extract_hash_answer(x['answer']) | |
| }) # type: ignore | |
| return data # type: ignore | |
| dataset = get_gsm8k_questions() | |
| # Reward functions | |
| def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]: | |
| responses = [completion[0]['content'] for completion in completions] | |
| q = prompts[0][-1]['content'] | |
| extracted_responses = [extract_xml_answer(r) for r in responses] | |
| print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}") | |
| return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)] | |
| def int_reward_func(completions, **kwargs) -> list[float]: | |
| responses = [completion[0]['content'] for completion in completions] | |
| extracted_responses = [extract_xml_answer(r) for r in responses] | |
| return [0.5 if r.isdigit() else 0.0 for r in extracted_responses] | |
| def strict_format_reward_func(completions, **kwargs) -> list[float]: | |
| """Reward function that checks if the completion has a specific format.""" | |
| pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$" | |
| responses = [completion[0]["content"] for completion in completions] | |
| matches = [re.match(pattern, r) for r in responses] | |
| return [0.5 if match else 0.0 for match in matches] | |
| def soft_format_reward_func(completions, **kwargs) -> list[float]: | |
| """Reward function that checks if the completion has a specific format.""" | |
| pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>" | |
| responses = [completion[0]["content"] for completion in completions] | |
| matches = [re.match(pattern, r) for r in responses] | |
| return [0.5 if match else 0.0 for match in matches] | |
| def count_xml(text) -> float: | |
| count = 0.0 | |
| if text.count("<reasoning>\n") == 1: | |
| count += 0.125 | |
| if text.count("\n</reasoning>\n") == 1: | |
| count += 0.125 | |
| if text.count("\n<answer>\n") == 1: | |
| count += 0.125 | |
| count -= len(text.split("\n</answer>\n")[-1])*0.001 | |
| if text.count("\n</answer>") == 1: | |
| count += 0.125 | |
| count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001 | |
| return count | |
| def xmlcount_reward_func(completions, **kwargs) -> list[float]: | |
| contents = [completion[0]["content"] for completion in completions] | |
| return [count_xml(c) for c in contents] | |
| #model_name = "meta-llama/Llama-3.2-1B-Instruct" | |
| model_name = "Qwen/Qwen2.5-1.5B-Instruct" | |
| if "Llama" in model_name: | |
| output_dir = "outputs/Llama-1B-GRPO" | |
| run_name = "Llama-1B-GRPO-gsm8k" | |
| else: | |
| output_dir="outputs/Qwen-1.5B-GRPO" | |
| run_name="Qwen-1.5B-GRPO-gsm8k" | |
| training_args = GRPOConfig( | |
| output_dir=output_dir, | |
| run_name=run_name, | |
| learning_rate=5e-6, | |
| adam_beta1 = 0.9, | |
| adam_beta2 = 0.99, | |
| weight_decay = 0.1, | |
| warmup_ratio = 0.1, | |
| lr_scheduler_type='cosine', | |
| logging_steps=1, | |
| bf16=True, | |
| per_device_train_batch_size=1, | |
| gradient_accumulation_steps=4, | |
| num_generations=16, | |
| max_prompt_length=256, | |
| max_completion_length=786, | |
| num_train_epochs=1, | |
| save_steps=100, | |
| max_grad_norm=0.1, | |
| report_to="wandb", | |
| log_on_each_node=False, | |
| ) | |
| peft_config = LoraConfig( | |
| r=16, | |
| lora_alpha=64, | |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"], | |
| task_type="CAUSAL_LM", | |
| lora_dropout=0.05, | |
| ) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| torch_dtype=torch.bfloat16, | |
| attn_implementation="flash_attention_2", | |
| device_map=None | |
| ).to("cuda") | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| tokenizer.pad_token = tokenizer.eos_token | |
| # use peft at your own risk; not working for me with multi-GPU training | |
| trainer = GRPOTrainer( | |
| model=model, | |
| processing_class=tokenizer, | |
| reward_funcs=[ | |
| xmlcount_reward_func, | |
| soft_format_reward_func, | |
| strict_format_reward_func, | |
| int_reward_func, | |
| correctness_reward_func], | |
| args=training_args, | |
| train_dataset=dataset, | |
| #peft_config=peft_config | |
| ) | |
| trainer.train() |
Jane thinks
This gist should become deep dive into the underlying complexities of the code discussed in Soban’s post on LinkedIn, where I break down how the GRPO approach manages reinforcement learning with modular reward functions and optimized configurations for efficient training on datasets like GSM8K. In my commentary, I clarify that while the visible code might appear minimal, it represents a sophisticated framework that leverages advanced techniques—such as bf16 precision and flash attention—that demand significant computational power and collaborative effort. The discussion also addresses commenters’ concerns by contrasting the apparent simplicity of the code with the substantial engineering and resource commitments that are necessary for such breakthrough AI developments. I hope this detailed analysis sparks further discussion and helps everyone appreciate the depth behind these modern AI methods.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Jane explains
GRPO
Generalized Reward Policy Optimization is a reinforcement learning framework that refines language model behavior by updating its policy based on various modular reward signals, thereby enhancing structured and coherent response generation.
In this script it serves as the core reinforcement learning mechanism that fine-tunes the language model by evaluating its outputs against several modular reward functions - each measuring criteria like correctness, formatting, token extraction, and adherence to an XML chain-of-thought structure - and then adjusting the model’s policy based on these rewards to incrementally improve performance and consistency during training on datasets such as GSM8K.
GSM8K
GSM8K is a dataset of roughly 8,000 grade-school math word problems used to benchmark and improve the arithmetic reasoning of language models.
In this script, GRPO drives a reinforcement learning loop for training a language model by generating outputs, evaluating them with several modular reward functions (assessing factors like correctness, formatting, token extraction, and adherence to an XML chain-of-thought), and then iteratively updating the model’s policy - in tandem with optimization techniques (such as bf16 precision and flash attention) - to enhance performance on datasets like GSM8K.