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

add instructions for installing on blackwell (#2812)

Browse source
This commit is contained in:
jeromeku 2025-06-27 04:54:55 -07:00 • committed by GitHub
commit 267ca3b209
4 changed files with 874 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

87
blackwell/README.md Normal file
View file

@ -0,0 +1,87 @@
## Unsloth Blackwell Compatibility
### 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` - safest is to use the nightly build: `uv pip install -U vllm --torch-backend=cu128 --extra-index-url https://wheels.vllm.ai/nightly`
- `xformers` - 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
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 --extra-index-url https://wheels.vllm.ai/nightly
```
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
```
4) Download and build `xformers`
```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) Update `triton`
```bash
uv pip install -U triton>=3.3.1
```
`triton>=3.3.1` is required for `Blackwell` support.
6) `transformers`
`transformers >= 4.53.0` breaks `unsloth` inference. Specifically, `transformers` with `gradient_checkpointing` enabled will automatically [switch off caching](https://github.com/huggingface/transformers/blob/67ddc82fbc7e52c6f42a395b4a6d278c55b77a39/src/transformers/modeling_layers.py#L52-L59).
When using `unsloth` `FastLanguageModel` to `generate` directly after training with `use_cache=True`, this will result in mismatch between expected and actual outputs [here](https://github.com/unslothai/unsloth/blob/bfa6a3678e2fb8097c5ece41d095a8051f099db3/unsloth/models/llama.py#L939).
Temporary solution is to switch off `gradient_checkpointing` (e.g., `model.disable_gradient_checkpointing()`) before generation if using `4.53.0` or stick with `4.52.4` for now:
```bash
uv pip install -U transformers==4.52.4
```
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 Normal file
View file

@ -0,0 +1,180 @@
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 Normal file
View file

@ -0,0 +1,178 @@
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))

429
blackwell/test_qwen3_grpo.py Normal file
View file

@ -0,0 +1,429 @@
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,
)
model.disable_gradient_checkpointing()
output = (
model.fast_generate(
[text],
sampling_params=sampling_params,
lora_request=None,
)[0]
.outputs[0]
.text
)
print(output)