ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Blackwell support

Browse source
This commit is contained in:
Daniel Han 2025-09-15 01:39:03 -07:00
commit 7e4788b646
5 changed files with 24 additions and 1003 deletions
Download patch file Download diff file Expand all files Collapse all files

32
README.md
View file

@ -38,14 +38,16 @@ Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-st
- See detailed documentation for Unsloth [here](https://docs.unsloth.ai/)
## ⚡ Quickstart
### Linux or WSL:
```
### Linux or WSL
```bash
pip install unsloth
```
### Windows
For Windows, ```pip install unsloth``` works only if you have Pytorch installed. For more info, read our [Windows Guide](https://docs.unsloth.ai/get-started/installing-+-updating/windows-installation).
For Windows, `pip install unsloth` works only if you have Pytorch installed. For more info, read our [Windows Guide](https://docs.unsloth.ai/get-started/installing-+-updating/windows-installation).
### Docker
Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://docs.unsloth.ai/get-started/install-and-update/docker).
### Blackwell
For RTX 50x, B200, 6000 GPUs, simply do `pip install unsloth`. Read our [Blackwell Guide](https://docs.unsloth.ai/basics/training-llms-with-blackwell-rtx-50-series-and-unsloth) for more details.
## 🦥 Unsloth.ai News
- 📣 **Memory-efficient RL** We're introducing even better RL. Our new kernels & algos allows faster RL with 50% less VRAM & 10× more context. [Read blog](https://docs.unsloth.ai/new/memory-efficient-rl)
@ -145,8 +147,14 @@ For **advanced installation instructions** or if you see weird errors during ins
1. Install `torch` and `triton`. Go to https://pytorch.org to install it. For example `pip install torch torchvision torchaudio triton`
2. Confirm if CUDA is installed correctly. Try `nvcc`. If that fails, you need to install `cudatoolkit` or CUDA drivers.
3. Install `xformers` manually. You can try installing `vllm` and seeing if `vllm` succeeds. Check if `xformers` succeeded with `python -m xformers.info` Go to https://github.com/facebookresearch/xformers. Another option is to install `flash-attn` for Ampere GPUs.
4. Double check that your versions of Python, CUDA, CUDNN, `torch`, `triton`, and `xformers` are compatible with one another. The [PyTorch Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix) may be useful.
3. Install `xformers` manually via:
```bash
pip install ninja
pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
```
Check if `xformers` succeeded with `python -m xformers.info` Go to https://github.com/facebookresearch/xformers. Another option is to install `flash-attn` for Ampere GPUs and ignore `xformers`
5. You can try installing `vllm` and seeing if `vllm` succeeds.
6. Double check that your versions of Python, CUDA, CUDNN, `torch`, `triton`, and `xformers` are compatible with one another. The [PyTorch Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix) may be useful.
5. Finally, install `bitsandbytes` and check it with `python -m bitsandbytes`
### Conda Installation (Optional)
@ -216,18 +224,26 @@ Or, run the below manually in a Python REPL:
try: import torch
except: raise ImportError('Install torch via `pip install torch`')
from packaging.version import Version as V
v = V(torch.__version__)
import re
v = V(re.match(r"[0-9\.]{3,}", torch.__version__).group(0))
cuda = str(torch.version.cuda)
is_ampere = torch.cuda.get_device_capability()[0] >= 8
if cuda != "12.1" and cuda != "11.8" and cuda != "12.4": raise RuntimeError(f"CUDA = {cuda} not supported!")
USE_ABI = torch._C._GLIBCXX_USE_CXX11_ABI
if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!")
if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
elif v < V('2.3.0'): x = 'cu{}{}-torch220'
elif v < V('2.4.0'): x = 'cu{}{}-torch230'
elif v < V('2.5.0'): x = 'cu{}{}-torch240'
elif v < V('2.6.0'): x = 'cu{}{}-torch250'
elif v < V('2.5.1'): x = 'cu{}{}-torch250'
elif v <= V('2.5.1'): x = 'cu{}{}-torch251'
elif v < V('2.7.0'): x = 'cu{}{}-torch260'
elif v < V('2.7.9'): x = 'cu{}{}-torch270'
elif v < V('2.8.0'): x = 'cu{}{}-torch271'
elif v < V('2.8.9'): x = 'cu{}{}-torch280'
else: raise RuntimeError(f"Torch = {v} too new!")
if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!")
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
```

207
blackwell/README.md
View file

@ -1,207 +0,0 @@
# Unsloth Blackwell Compatibility
For RTX 5060, RTX 5070, RTX 5080, RTX 5090 GPUs and also B200, B40, GB100, GB102, GB20* and GPUs listed in https://en.wikipedia.org/wiki/Blackwell_(microarchitecture)
## Overview
`Blackwell` (`sm100+`) requires all dependent libraries to be compiled with `cuda 12.8`.
The core libs for running unsloth which have dependencies on `CUDA` version are:
- `bitsandbytes` - already has wheels built with `CUDA 12.8` so `pip install` should work out of the box
- `triton` - requires `triton>=3.3.1`
- `torch` - requires installing with `pip install torch --extra-index-url https://download.pytorch.org/whl/cu128`
- `vllm` - vLLM 0.10.0 supports Blackwell now, but use CUDA 12.8: `uv pip install -U vllm --torch-backend=cu128`
- `xformers` - (Optional) as of 6/26, `xformers` wheels are not yet built with `sm100+` enabled as support was only recently [added](https://github.com/facebookresearch/xformers/commit/d9b3b6e2b38ca485c89507ef8ac1fbef2723cdfa) so will require a source build (see below).
## Installation
### Using uv
The installation order is important, since we want the overwrite bundled dependencies with specific versions (namely, `xformers` and `triton`).
1) I prefer to use `uv` over `pip` as it's faster and better for resolving dependencies, especially for libraries which depend on `torch` but for which a specific `CUDA` version is required per this scenario.
Install `uv`
```bash
curl -LsSf https://astral.sh/uv/install.sh | sh && source $HOME/.local/bin/env
```
Create a project dir and venv:
```bash
mkdir 'unsloth-blackwell' && cd 'unsloth-blackwell'
uv venv .venv --python=3.12 --seed
source .venv/bin/activate
```
2) Install `vllm`
```bash
uv pip install -U vllm --torch-backend=cu128
```
Note that we have to specify `cu128`, otherwise `vllm` will install `torch==2.7.0` but with `cu126`.
3) Install `unsloth` dependencies
```bash
uv pip install unsloth unsloth_zoo bitsandbytes
```
If you notice weird resolving issues due to Xformers, you can also install Unsloth from source without Xformers:
```bash
uv pip install -qqq \
"unsloth_zoo[base] @ git+https://github.com/unslothai/unsloth-zoo" \
"unsloth[base] @ git+https://github.com/unslothai/unsloth"
```
4) Download and build `xformers` (Optional)
Xformers is optional, but it is definitely faster and uses less memory. We'll use PyTorch's native SDPA if you do not want Xformers. Building Xformers from source might be slow, so beware!
```bash
# First uninstall xformers installed by previous libraries
uv pip uninstall xformers
# Clone and build
git clone --depth=1 https://github.com/facebookresearch/xformers --recursive
cd xformers
export TORCH_CUDA_ARCH_LIST="12.0"
python setup.py install
```
Note that we have to explicitly set `TORCH_CUDA_ARCH_LIST=12.0`.
5) `transformers`
Install any transformers version, but best to get the latest.
```bash
uv pip install -U transformers
```
### Using conda or mamba
1) Install `conda/mamba`
```bash
curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
```
Run the installation script
```bash
bash Miniforge3-$(uname)-$(uname -m).sh
```
Create a conda or mamba environment
```bash
conda create --name unsloth-blackwell python==3.12 -y
```
Activate newly created environment
```bash
conda activate unsloth-blackwell
```
2) Install `vllm`
Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:`
```bash
pip install -U vllm --extra-index-url https://download.pytorch.org/whl/cu128
```
Note that we have to specify `cu128`, otherwise `vllm` will install `torch==2.7.0` but with `cu126`.
3) Install `unsloth` dependencies
Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:`
```bash
pip install unsloth unsloth_zoo bitsandbytes
```
4) Download and build `xformers` (Optional)
Xformers is optional, but it is definitely faster and uses less memory. We'll use PyTorch's native SDPA if you do not want Xformers. Building Xformers from source might be slow, so beware!
You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:`
```bash
# First uninstall xformers installed by previous libraries
pip uninstall xformers
# Clone and build
git clone --depth=1 https://github.com/facebookresearch/xformers --recursive
cd xformers
export TORCH_CUDA_ARCH_LIST="12.0"
python setup.py install
```
Note that we have to explicitly set `TORCH_CUDA_ARCH_LIST=12.0`.
5) Update `triton`
Make sure you are inside the activated conda/mamba environment. You should see the name of your environment as a prefix to your terminal shell like this your `(unsloth-blackwell)user@machine:`
```bash
pip install -U triton>=3.3.1
```
`triton>=3.3.1` is required for `Blackwell` support.
6) `Transformers`
Install any transformers version, but best to get the latest.
```bash
uv pip install -U transformers
```
If you are using mamba as your package just replace conda with mamba for all commands shown above.
## WSL-Specific Notes
If you're using WSL (Windows Subsystem for Linux) and encounter issues during xformers compilation (reminder Xformers is optional, but faster for training) follow these additional steps:
1. **Increase WSL Memory Limit**
Create or edit the WSL configuration file:
```bash
# Create or edit .wslconfig in your Windows user directory
# (typically C:\Users\YourUsername\.wslconfig)
# Add these lines to the file
[wsl2]
memory=16GB # Minimum 16GB recommended for xformers compilation
processors=4 # Adjust based on your CPU cores
swap=2GB
localhostForwarding=true
```
After making these changes, restart WSL:
```powershell
wsl --shutdown
```
2. **Install xformers**
Use the following command to install xformers with optimized compilation for WSL:
```bash
# Set CUDA architecture for Blackwell GPUs
export TORCH_CUDA_ARCH_LIST="12.0"
# Install xformers from source with optimized build flags
pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
```
The `--no-build-isolation` flag helps avoid potential build issues in WSL environments.
## Post Installation notes:
After installation, your environment should look similar to `blackwell.requirements.txt`.
Note, might need to downgrade `numpy<=2.2` after all the installs.
## Test
Both `test_llama32_sft.py` and `test_qwen3_grpo.py` should run without issue if correct install. If not, check diff between your installed env and `blackwell.requirements.txt`.

180
blackwell/blackwell.requirements.txt
View file

@ -1,180 +0,0 @@
Using Python 3.12.11 environment at: unsloth-bw/.venv
accelerate==1.8.1
aiohappyeyeballs==2.6.1
aiohttp==3.12.13
aiosignal==1.3.2
airportsdata==20250622
annotated-types==0.7.0
anyio==4.9.0
astor==0.8.1
asttokens==3.0.0
attrs==25.3.0
bitsandbytes==0.46.0
blake3==1.0.5
cachetools==6.1.0
certifi==2025.6.15
charset-normalizer==3.4.2
click==8.2.1
cloudpickle==3.1.1
comm==0.2.2
compressed-tensors==0.10.1
cupy-cuda12x==13.4.1
cut-cross-entropy==25.1.1
datasets==3.6.0
debugpy==1.8.14
decorator==5.2.1
depyf==0.18.0
diffusers==0.34.0
dill==0.3.8
diskcache==5.6.3
distro==1.9.0
dnspython==2.7.0
docstring-parser==0.16
einops==0.8.1
email-validator==2.2.0
executing==2.2.0
fastapi==0.115.14
fastapi-cli==0.0.7
fastrlock==0.8.3
filelock==3.18.0
frozenlist==1.7.0
fsspec==2025.5.1
gguf==0.17.1
h11==0.16.0
hf-transfer==0.1.9
hf-xet==1.1.5
httpcore==1.0.9
httptools==0.6.4
httpx==0.28.1
huggingface-hub==0.33.1
idna==3.10
importlib-metadata==8.7.0
interegular==0.3.3
ipykernel==6.29.5
ipython==9.3.0
ipython-pygments-lexers==1.1.1
jedi==0.19.2
jinja2==3.1.6
jiter==0.10.0
jsonschema==4.24.0
jsonschema-specifications==2025.4.1
jupyter-client==8.6.3
jupyter-core==5.8.1
lark==1.2.2
llguidance==0.7.30
llvmlite==0.44.0
lm-format-enforcer==0.10.11
markdown-it-py==3.0.0
markupsafe==3.0.2
matplotlib-inline==0.1.7
mdurl==0.1.2
mistral-common==1.6.2
mpmath==1.3.0
msgpack==1.1.1
msgspec==0.19.0
multidict==6.5.1
multiprocess==0.70.16
nest-asyncio==1.6.0
networkx==3.5
ninja==1.11.1.4
numba==0.61.2
numpy==2.2.0
nvidia-cublas-cu12==12.8.3.14
nvidia-cuda-cupti-cu12==12.8.57
nvidia-cuda-nvrtc-cu12==12.8.61
nvidia-cuda-runtime-cu12==12.8.57
nvidia-cudnn-cu12==9.7.1.26
nvidia-cufft-cu12==11.3.3.41
nvidia-cufile-cu12==1.13.0.11
nvidia-curand-cu12==10.3.9.55
nvidia-cusolver-cu12==11.7.2.55
nvidia-cusparse-cu12==12.5.7.53
nvidia-cusparselt-cu12==0.6.3
nvidia-nccl-cu12==2.26.2
nvidia-nvjitlink-cu12==12.8.61
nvidia-nvtx-cu12==12.8.55
openai==1.92.2
opencv-python-headless==4.11.0.86
outlines==0.1.11
outlines-core==0.1.26
packaging==25.0
pandas==2.3.0
parso==0.8.4
partial-json-parser==0.2.1.1.post6
peft==0.15.2
pexpect==4.9.0
pillow==11.2.1
pip==25.1.1
platformdirs==4.3.8
prometheus-client==0.22.1
prometheus-fastapi-instrumentator==7.1.0
prompt-toolkit==3.0.51
propcache==0.3.2
protobuf==3.20.3
psutil==7.0.0
ptyprocess==0.7.0
pure-eval==0.2.3
py-cpuinfo==9.0.0
pyarrow==20.0.0
pybase64==1.4.1
pycountry==24.6.1
pydantic==2.11.7
pydantic-core==2.33.2
pygments==2.19.2
python-dateutil==2.9.0.post0
python-dotenv==1.1.1
python-json-logger==3.3.0
python-multipart==0.0.20
pytz==2025.2
pyyaml==6.0.2
pyzmq==27.0.0
ray==2.47.1
referencing==0.36.2
regex==2024.11.6
requests==2.32.4
rich==14.0.0
rich-toolkit==0.14.7
rpds-py==0.25.1
safetensors==0.5.3
scipy==1.16.0
sentencepiece==0.2.0
setuptools==80.9.0
shellingham==1.5.4
shtab==1.7.2
six==1.17.0
sniffio==1.3.1
stack-data==0.6.3
starlette==0.46.2
sympy==1.14.0
tiktoken==0.9.0
tokenizers==0.21.2
torch==2.7.0+cu128
torchaudio==2.7.0+cu128
torchvision==0.22.0+cu128
tornado==6.5.1
tqdm==4.67.1
traitlets==5.14.3
transformers==4.52.4
triton==3.3.1
trl==0.19.0
typeguard==4.4.4
typer==0.16.0
typing-extensions==4.14.0
typing-inspection==0.4.1
tyro==0.9.24
tzdata==2025.2
unsloth==2025.6.8
unsloth-zoo==2025.6.5
urllib3==2.5.0
uvicorn==0.34.3
uvloop==0.21.0
vllm==0.9.2.dev280+g04e1642e3
watchfiles==1.1.0
wcwidth==0.2.13
websockets==15.0.1
wheel==0.45.1
xformers==0.0.32+ff490c3.d20250626
xgrammar==0.1.19
xxhash==3.5.0
yarl==1.20.1
zipp==3.23.0

178
blackwell/test_llama32_sft.py
View file

@ -1,178 +0,0 @@
from unsloth import FastLanguageModel
from transformers import (
AutoModelForCausalLM,
DataCollatorForSeq2Seq,
AutoTokenizer,
)
from trl import SFTConfig, SFTTrainer
from unsloth.chat_templates import (
get_chat_template,
standardize_sharegpt,
train_on_responses_only,
)
from datasets import load_dataset
from peft import AutoPeftModelForCausalLM
import torch
max_seq_length = 2048
dtype = None
load_in_4bit = True
fourbit_models = [
"unsloth/Meta-Llama-3.1-8B-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
"unsloth/Meta-Llama-3.1-70B-bnb-4bit",
"unsloth/Meta-Llama-3.1-405B-bnb-4bit",
"unsloth/Mistral-Small-Instruct-2409",
"unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
"unsloth/Phi-3.5-mini-instruct",
"unsloth/Phi-3-medium-4k-instruct",
"unsloth/gemma-2-9b-bnb-4bit",
"unsloth/gemma-2-27b-bnb-4bit",
"unsloth/Llama-3.2-1B-bnb-4bit",
"unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
"unsloth/Llama-3.2-3B-bnb-4bit",
"unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
"unsloth/Llama-3.3-70B-Instruct-bnb-4bit",
]
model, tokenizer = FastLanguageModel.from_pretrained(
model_name="unsloth/Llama-3.2-1B-Instruct",
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit,
)
model: AutoModelForCausalLM = FastLanguageModel.get_peft_model(
model,
r=16,
target_modules=[
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
],
lora_alpha=16,
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=3407,
use_rslora=False,
loftq_config=None,
)
tokenizer = get_chat_template(tokenizer, chat_template="llama-3.1")
def formatting_prompts_func(examples):
convos = examples["conversations"]
texts = [
tokenizer.apply_chat_template(
convo, tokenize=False, add_generation_prompt=False
)
for convo in convos
]
return {"text": texts}
dataset = load_dataset("mlabonne/FineTome-100k", split="train")
dataset = standardize_sharegpt(dataset)
dataset = dataset.map(formatting_prompts_func, batched=True)
dataset[5]["conversations"]
dataset[5]["text"]
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=max_seq_length,
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
dataset_num_proc=2,
packing=False,
args=SFTConfig(
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
warmup_steps=5,
max_steps=10,
learning_rate=2e-4,
logging_steps=1,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=3407,
output_dir="outputs",
report_to="none",
),
)
trainer = train_on_responses_only(
trainer,
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
)
tokenizer.decode(trainer.train_dataset[5]["input_ids"])
space = tokenizer(" ", add_special_tokens=False).input_ids[0]
tokenizer.decode(
[space if x == -100 else x for x in trainer.train_dataset[5]["labels"]]
)
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(
torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3
)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")
trainer_stats = trainer.train()
used_memory = round(
torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3
)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory / max_memory * 100, 3)
lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
f"{round(trainer_stats.metrics['train_runtime'] / 60, 2)} minutes used for training."
)
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(
f"Peak reserved memory for training % of max memory = {lora_percentage} %."
)
tokenizer = get_chat_template(tokenizer, chat_template="llama-3.1")
FastLanguageModel.for_inference(model)
messages = [
{
"role": "user",
"content": "Continue the fibonnaci sequence: 1, 1, 2, 3, 5, 8,",
},
]
inputs = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt",
).to("cuda")
model.gradient_checkpointing_disable() # This is required if using transformers >= 4.53.0 and `use_cache=True`
outputs = model.generate(
input_ids=inputs,
max_new_tokens=64,
use_cache=True,
temperature=1.5,
min_p=0.1,
)
print(tokenizer.batch_decode(outputs))

430
blackwell/test_qwen3_grpo.py
View file

@ -1,430 +0,0 @@
from unsloth import FastLanguageModel
import torch
max_seq_length = 2048
lora_rank = 32
model, tokenizer = FastLanguageModel.from_pretrained(
model_name="unsloth/Qwen3-0.6B-Base",
max_seq_length=max_seq_length,
load_in_4bit=False,
fast_inference=True,
max_lora_rank=lora_rank,
gpu_memory_utilization=0.7,
)
model = FastLanguageModel.get_peft_model(
model,
r=lora_rank,
target_modules=[
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
],
lora_alpha=lora_rank * 2,
use_gradient_checkpointing="unsloth",
random_state=3407,
)
reasoning_start = "<start_working_out>"
reasoning_end = "<end_working_out>"
solution_start = "<SOLUTION>"
solution_end = "</SOLUTION>"
system_prompt = f"""You are given a problem.
Think about the problem and provide your working out.
Place it between {reasoning_start} and {reasoning_end}.
Then, provide your solution between {solution_start}{solution_end}"""
system_prompt
chat_template = (
"{% if messages[0]['role'] == 'system' %}"
"{{ messages[0]['content'] + eos_token }}"
"{% set loop_messages = messages[1:] %}"
"{% else %}"
"{{ '{system_prompt}' + eos_token }}"
"{% set loop_messages = messages %}"
"{% endif %}"
"{% for message in loop_messages %}"
"{% if message['role'] == 'user' %}"
"{{ message['content'] }}"
"{% elif message['role'] == 'assistant' %}"
"{{ message['content'] + eos_token }}"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}{{ '{reasoning_start}' }}"
"{% endif %}"
)
chat_template = chat_template.replace(
"'{system_prompt}'", f"'{system_prompt}'"
).replace("'{reasoning_start}'", f"'{reasoning_start}'")
tokenizer.chat_template = chat_template
tokenizer.apply_chat_template(
[
{"role": "user", "content": "What is 1+1?"},
{
"role": "assistant",
"content": f"{reasoning_start}I think it's 2.{reasoning_end}{solution_start}2{solution_end}",
},
{"role": "user", "content": "What is 2+2?"},
],
tokenize=False,
add_generation_prompt=True,
)
from datasets import load_dataset
import pandas as pd
import numpy as np
dataset = load_dataset("unsloth/OpenMathReasoning-mini", split="cot")
dataset = dataset.to_pandas()[["expected_answer", "problem", "generated_solution"]]
is_number = pd.to_numeric(
pd.Series(dataset["expected_answer"]), errors="coerce"
).notnull()
dataset = dataset.iloc[np.where(is_number)[0]]
dataset
def format_dataset(x):
expected_answer = x["expected_answer"]
problem = x["problem"]
thoughts = x["generated_solution"]
thoughts = thoughts.replace("<think>", "").replace("</think>", "")
thoughts = thoughts.strip()
final_prompt = (
reasoning_start
+ thoughts
+ reasoning_end
+ solution_start
+ expected_answer
+ solution_end
)
return [
{"role": "system", "content": system_prompt},
{"role": "user", "content": problem},
{"role": "assistant", "content": final_prompt},
]
dataset["Messages"] = dataset.apply(format_dataset, axis=1)
tokenizer.apply_chat_template(dataset["Messages"][0], tokenize=False)
dataset["N"] = dataset["Messages"].apply(
lambda x: len(tokenizer.apply_chat_template(x))
)
dataset = dataset.loc[dataset["N"] <= max_seq_length / 2].copy()
dataset.shape
from datasets import Dataset
dataset["text"] = tokenizer.apply_chat_template(
dataset["Messages"].values.tolist(), tokenize=False
)
dataset = Dataset.from_pandas(dataset)
dataset
from trl import SFTTrainer, SFTConfig
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
args=SFTConfig(
dataset_text_field="text",
per_device_train_batch_size=1,
gradient_accumulation_steps=1,
warmup_steps=5,
num_train_epochs=2,
learning_rate=2e-4,
logging_steps=5,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=3407,
report_to="none",
),
)
trainer.train()
text = tokenizer.apply_chat_template(
dataset[0]["Messages"][:2],
tokenize=False,
add_generation_prompt=True,
)
from transformers import TextStreamer
_ = model.generate(
**tokenizer(text, return_tensors="pt").to("cuda"),
temperature=0,
max_new_tokens=1024,
streamer=TextStreamer(tokenizer, skip_prompt=False),
)
del dataset
torch.cuda.empty_cache()
import gc
gc.collect()
from datasets import load_dataset
dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train")
dataset
dataset[0]["prompt"]
dataset[0]["solution"]
def extract_hash_answer(text):
return text
extract_hash_answer(dataset[0]["solution"])
dataset = dataset.map(
lambda x: {
"prompt": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": x["prompt"]},
],
"answer": extract_hash_answer(x["solution"]),
}
)
dataset[0]
import re
solution_end_regex = (
r"</SOLUTION>[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?"
)
match_format = re.compile(
rf"{reasoning_end}.*?"
rf"{solution_start}(.+?){solution_end_regex}"
rf"[\s]{{0,}}$",
flags=re.MULTILINE | re.DOTALL,
)
match_format
match_format.findall(
f"Let me think!<end_working_out><SOLUTION>\n2\n</SOLUTION>",
)
match_format.findall(
f"<start_working_out>Let me think!<end_working_out><SOLUTION> 2 </SOLUTION>\n\n",
)
def match_format_exactly(completions, **kwargs):
scores = []
for completion in completions:
score = 0
response = completion[0]["content"]
if match_format.search(response) is not None:
score += 3.0
scores.append(score)
return scores
def match_format_approximately(completions, **kwargs):
scores = []
for completion in completions:
score = 0
response = completion[0]["content"]
score += 0.5 if response.count(reasoning_end) == 1 else -1.0
score += 0.5 if response.count(solution_start) == 1 else -1.0
score += 0.5 if response.count(solution_end) == 1 else -1.0
scores.append(score)
return scores
def check_answer(prompts, completions, answer, **kwargs):
question = prompts[0][-1]["content"]
responses = [completion[0]["content"] for completion in completions]
extracted_responses = [
guess.group(1) if (guess := match_format.search(r)) is not None else None
for r in responses
]
scores = []
for guess, true_answer in zip(extracted_responses, answer):
score = 0
if guess is None:
scores.append(-2.0)
continue
if guess == true_answer:
score += 5.0
elif guess.strip() == true_answer.strip():
score += 3.5
else:
try:
ratio = float(guess) / float(true_answer)
if ratio >= 0.9 and ratio <= 1.1:
score += 2.0
elif ratio >= 0.8 and ratio <= 1.2:
score += 1.5
else:
score -= 2.5
except:
score -= 4.5
scores.append(score)
return scores
match_numbers = re.compile(
solution_start + r".*?[\s]{0,}([-]?[\d\.\,]{1,})", flags=re.MULTILINE | re.DOTALL
)
print(match_numbers.findall("<SOLUTION> 0.34 </SOLUTION>"))
print(match_numbers.findall("<SOLUTION> 123,456 </SOLUTION>"))
print(match_numbers.findall("<SOLUTION> -0.234 </SOLUTION>"))
print(match_numbers.findall("<SOLUTION>17</SOLUTION>"))
global PRINTED_TIMES
PRINTED_TIMES = 0
global PRINT_EVERY_STEPS
PRINT_EVERY_STEPS = 5
def check_numbers(prompts, completions, answer, **kwargs):
question = prompts[0][-1]["content"]
responses = [completion[0]["content"] for completion in completions]
extracted_responses = [
guess.group(1) if (guess := match_numbers.search(r)) is not None else None
for r in responses
]
scores = []
global PRINTED_TIMES
global PRINT_EVERY_STEPS
if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
print(
"*" * 20 + f"Question:\n{question}",
f"\nAnswer:\n{answer[0]}",
f"\nResponse:\n{responses[0]}",
f"\nExtracted:\n{extracted_responses[0]}",
)
PRINTED_TIMES += 1
for guess, true_answer in zip(extracted_responses, answer):
if guess is None:
scores.append(-2.5)
continue
try:
true_answer = float(true_answer.strip())
guess = float(guess.strip().replace(",", ""))
scores.append(3.5 if guess == true_answer else -1.5)
except:
scores.append(0)
continue
return scores
tokenized = dataset.map(
lambda x: {
"tokens": tokenizer.apply_chat_template(
x["prompt"], add_generation_prompt=True, tokenize=True
)
},
batched=True,
)
print(tokenizer.decode(tokenized[0]["tokens"]))
tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])})
import numpy as np
maximum_length = int(np.quantile(tokenized["L"], 0.9))
print("Max Length = ", maximum_length)
dataset = dataset.select(np.where(np.array(tokenized["L"]) <= maximum_length)[0])
del tokenized
max_prompt_length = maximum_length + 1
max_completion_length = max_seq_length - max_prompt_length
from vllm import SamplingParams
vllm_sampling_params = SamplingParams(
min_p=0.1,
top_p=1.0,
top_k=-1,
seed=3407,
stop=[tokenizer.eos_token],
include_stop_str_in_output=True,
)
from trl import GRPOConfig, GRPOTrainer
training_args = GRPOConfig(
vllm_sampling_params=vllm_sampling_params,
temperature=1.0,
learning_rate=5e-6,
weight_decay=0.01,
warmup_ratio=0.1,
lr_scheduler_type="linear",
optim="adamw_8bit",
logging_steps=1,
per_device_train_batch_size=1,
gradient_accumulation_steps=1,
num_generations=4,
max_prompt_length=max_prompt_length,
max_completion_length=max_completion_length,
max_steps=10,
save_steps=100,
report_to="none",
output_dir="outputs",
)
trainer = GRPOTrainer(
model=model,
processing_class=tokenizer,
reward_funcs=[
match_format_exactly,
match_format_approximately,
check_answer,
check_numbers,
],
args=training_args,
train_dataset=dataset,
)
trainer.train()
text = "What is the sqrt of 101?"
from vllm import SamplingParams
sampling_params = SamplingParams(
temperature=1.0,
top_k=50,
max_tokens=1024,
)
output = (
model.fast_generate(
[text],
sampling_params=sampling_params,
lora_request=None,
)[0]
.outputs[0]
.text
)
print(output)