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:

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

Usage

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:
```
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.