---
language:
- en
license: apache-2.0
base_model: meta-llama/Meta-Llama-3.1-8B-Instruct
tags:
- llama
- unsloth
- trl
- grpo
- reasoning
- math
- gsm8k
- rocm
- amd
datasets:
- openai/gsm8k
pipeline_tag: text-generation
---

# llama31-8b-grpo-gsm8k-run1

A reasoning-focused fine-tune of Llama 3.1 8B Instruct, trained using Group Relative Policy Optimization (GRPO) on the GSM8K math reasoning dataset. Trained entirely on AMD ROCm hardware via the AMD Academy lab environment.

---

## Model Details

| Parameter | Value |
|-----------|-------|
| Base model | meta-llama/Meta-Llama-3.1-8B-Instruct |
| Fine-tuning method | GRPO (reinforcement learning, no supervised labels) |
| Dataset | GSM8K (7,473 train examples) |
| Training steps | 250 |
| LoRA rank | 32 |
| Quantization | None (BF16) |
| Effective batch size | 12 (6 generations × 2 accumulation × 1 GPU) |
| Learning rate | 5e-6 |
| Hardware | AMD Instinct MI300X VF (191.688 GB HBM) |
| Framework | Unsloth 2025.3.19 + TRL + ROCm |
| Torch | 2.7.0 + ROCm |
| vLLM | 0.7.4 |

---

## What is GRPO?

GRPO (Group Relative Policy Optimization) trains the model using reinforcement learning without a separate reward model. For each prompt, the model generates a group of completions, scores each one using reward functions, and updates based on relative scores within the group — no human-labeled reasoning chains required.

This is the same algorithm used by DeepSeek to train R1-Zero, where the model spontaneously learned to allocate more thinking time to harder problems — the so-called "aha moment."

---

## Training Approach

Five reward signals were used simultaneously:

| Reward Function | Max Score | What it checks |
|----------------|-----------|----------------|
| `correctness_reward_func` | 2.0 | Extracted answer matches gold answer exactly |
| `int_reward_func` | 0.5 | Answer is a valid integer |
| `strict_format_reward_func` | 0.5 | Exact newline structure inside tags |
| `soft_format_reward_func` | 0.5 | Tags present anywhere in output |
| `xmlcount_reward_func` | 0.5 | Partial credit for tag placement; penalises text after `</answer>` |

**Maximum possible reward per completion: 3.5**

The `xmlcount_reward_func` applies a `-0.001` penalty per character appearing after `</answer>`, discouraging the model from rambling after giving its answer.

---

## Expected Output Format

The model is trained to respond in the following format:

```text
<reasoning>
Step 1: ...
Step 2: ...
</reasoning>
<answer>
42
</answer>
```

---

## Usage

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "sarapatel/llama31-8b-grpo-gsm8k-run1",
    torch_dtype="auto",
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained("sarapatel/llama31-8b-grpo-gsm8k-run1")

SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

prompt = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many in May. How many clips did she sell altogether?"}
]

inputs = tokenizer.apply_chat_template(
    prompt,
    return_tensors="pt",
    add_generation_prompt=True
).to(model.device)

outputs = model.generate(inputs, max_new_tokens=512, temperature=0.8, top_p=0.95)
print(tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True))
```

---

## ROCm Findings

This model was trained entirely on AMD ROCm hardware. Key findings during the training process:

**What worked:**
- Unsloth's vLLM fast inference loaded and ran correctly under ROCm
- BF16 full-precision training was stable throughout all 250 steps
- `gpu_memory_utilization = 0.6` was sufficient for the full run

**What did not work:**
- `load_in_4bit = True` fails immediately with:

  ```text
  RuntimeError: bitsandbytes quantization is currently not supported in rocm.
  ```

  `bitsandbytes` is a CUDA-only library. Its 4bit/8bit quantization kernels have no ROCm equivalent. If you are on AMD hardware, keep `load_in_4bit = False`.

- Increasing `num_generations` beyond 6 with full BF16 weights caused OOM kernel crashes around step 100/250 due to peak VRAM pressure during simultaneous generation.

---

## Limitations

- 250 training steps is a short run — reward curves had not fully stabilized by the end. Expect better reasoning quality with 1000+ steps.
- Not evaluated on the GSM8K test set — no benchmark numbers are available for this checkpoint.
- `bitsandbytes` 4bit quantization is unavailable on ROCm. Use BF16 for inference, or explore GPTQ via AutoGPTQ as an alternative quantization path on AMD hardware.

---

## Training Metrics

Training was tracked via Weights & Biases under project `grpo-gsm8k-amd`, run name `llama31-8b-noquant-250steps`.

---

## References

- DeepSeek R1 paper: https://arxiv.org/abs/2501.12948
- GSM8K paper: https://arxiv.org/abs/2110.14168
- Unsloth: https://github.com/unslothai/unsloth
- TRL (GRPOTrainer): https://github.com/huggingface/trl

---

## Author

Saransh Patel - BTech Computer Science

Research interests: ML Systems, LLM Inference Optimization, GPU Programming, AMD ROCm

GitHub: https://github.com/SARPAT  
Site: https://sarpat.github.io

# Uploaded  model

- **Developed by:** sarapatel
- **License:** apache-2.0
- **Finetuned from model :** unsloth/Meta-Llama-3.1-8B-Instruct

This llama model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library.

[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)