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

Merge branch 'main' into nightly

Browse source
This commit is contained in:
Daniel Han 2025-06-28 17:44:21 -07:00
commit 8943c1b5dd
8 changed files with 1633 additions and 8 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)

2
unsloth/models/_utils.py
View file

@ -280,7 +280,7 @@ pass
from transformers import __version__ as transformers_version
from transformers import PretrainedConfig
model_architectures = ["llama", "mistral", "gemma", "gemma2", "qwen2", "granite", "qwen3", "qwen3_moe"]
model_architectures = ["llama", "mistral", "gemma", "gemma2", "qwen2", "granite", "qwen3", "qwen3_moe", "falcon_h1"]
for model_name in model_architectures:
config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"

703
unsloth/models/falcon_h1.py Normal file
View file

@ -0,0 +1,703 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .llama import *
import os
from ._utils import __version__
from unsloth_zoo.utils import Version, _get_dtype
from .llama import (
LlamaRotaryEmbedding,
LlamaLinearScalingRotaryEmbedding,
_LlamaModel_fast_forward_inference,
)
try:
from transformers.models.falcon_h1.modeling_falcon_h1 import (
FalconH1Attention,
FalconH1DecoderLayer,
FalconH1Model,
FalconH1ForCausalLM,
FalconHybridMambaAttentionDynamicCache,
)
except:
from transformers import __version__ as transformers_version
transformers_version = Version(transformers_version)
if not transformers_version >= Version("4.53.0"): #TODO: Update when transformers is updated
raise ImportError(
f"Unsloth: Your transformers version of {transformers_version} does not support FalconH1.\n"\
f"The minimum required version is 4.53.0.\n"\
f'Try `pip install --upgrade "transformers>=4.53.0"`\n'\
f"to obtain the latest transformers build, then restart this session."\
)
pass
from transformers.modeling_attn_mask_utils import (
_prepare_4d_causal_attention_mask_for_sdpa,
)
# For Pytorch 2.1.1
try:
from transformers.models.falcon_h1.modeling_falcon_h1 import (
FalconH1Attention,
)
except:
# if we are on a old version of transformers technically it should fail in the try except above
# but if somehow we make it here, we need to raise an error since FalconH1Attention is not available
# or renamed
raise ImportError("Unsloth: Could not import FalconH1Attention from transformers.models.falcon_h1.modeling_falcon_h1.")
pass
def FalconH1Attention_fast_forward(
self,
hidden_states: torch.Tensor,
causal_mask: Optional[BlockDiagonalCausalMask] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: bool = False,
use_cache: bool = False,
padding_mask: Optional[torch.LongTensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
*args, **kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
# Clear inference
if hasattr(self, "paged_attention"):
del self.paged_attention_K
del self.paged_attention_V
del self.paged_attention
del self.temp_QA
del self.temp_KV
del self.RH_Q
del self.attention
pass
bsz, q_len, _ = hidden_states.size()
n_heads = self.config.num_attention_heads
n_groups = self.num_key_value_groups
n_kv_heads = self.config.num_key_value_heads
head_dim = self.head_dim
assert(n_kv_heads * n_groups == n_heads)
Q, K, V = self.apply_qkv(self, hidden_states)
Q = Q.view(bsz, q_len, n_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
K = K.view(bsz, q_len, n_kv_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
# Falcon H1 multiplies key states by a multiplier
K = K * self.config.key_multiplier
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
kv_seq_len = K.shape[-2]
if past_key_value is not None:
kv_seq_len += past_key_value[0].shape[-2]
if position_embeddings:
cos, sin = position_embeddings
else:
# Extend RoPE dynamically to fit in VRA
rotary_emb = self.rotary_emb
rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
if position_ids is None:
# Useful for LongRoPE
cos, sin = rotary_emb.get_cached(kv_seq_len)
else:
cos, sin = rotary_emb(V, seq_len = kv_seq_len)
Q, K = fast_rope_embedding(Q, K, cos, sin)
if past_key_value is not None:
K = torch.cat([past_key_value[0], K], dim = 2)
V = torch.cat([past_key_value[1], V], dim = 2)
pass
past_key_value = (K, V) if use_cache else None
# Attention module
if (not HAS_FLASH_ATTENTION and attention_mask is None):
# Xformers memory efficient attention
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
V = V.transpose(1, 2)
K_M = V_M = bsz * kv_seq_len
Q_M = bsz * q_len
# Group query attention
K = K .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
V = V .view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
if hidden_states.requires_grad:
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
else:
# Xformers does support the forward pass though
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
pass
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
A = A.view(bsz, q_len, n_heads, head_dim)
elif HAS_FLASH_ATTENTION and attention_mask is None:
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
V = V.transpose(1, 2)
sw = kv_seq_len
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
else:
# Grouped query attention
# if n_groups != 1:
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
# pass
# Must be contiguous or else results are False!
# https://github.com/pytorch/pytorch/issues/112577
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
# Needs (batch_size, n_heads, seq_len, head_dim)
# is_casual and attention_mask must not be both set!
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
# Go back to (batch_size, seq_len, n_heads, head_dim)
A = A.transpose(1, 2).contiguous()
pass
attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
attn_output = self.apply_o(self, attn_output)
attn_weights = None
return attn_output, attn_weights, past_key_value
pass
torch_matmul = torch.matmul
def FalconH1Attention_fast_forward_inference(
self,
hidden_states: torch.Tensor,
past_key_value: Optional[Tuple[torch.Tensor]],
position_ids,
do_prefill = False,
attention_mask = None,
):
"""
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
Fast inference using KV cache.
QK^T can be computed in 4 chunks
[Q, q] @ [K, k].T where q, k are the new tokens.
[QK^T, Qk^T]
[qK^T, qk^T]
Since the attention mask wipes Qk^T, we just get
[QK^T, 0]
[qK^T, qk^T]
Since softmax is row-wise, we get
softmax([QK^T, 0])
softmax([qK^T, qk^T])
We then multiply by [V]
[v]
softmax([QK^T, 0]) [softmax(QK^T)V] *
softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
But notice * [softmax(QK^T)V] is just the last attention.
We just need to compute the last final row.
This means we can pass in a row of Q, but we need to
remember K and V, which are called the KV cache.
"""
Xn = hidden_states
bsz, _, hd = hidden_states.size()
K1, V1 = past_key_value
dtype = Xn.dtype
n_heads = self.config.num_attention_heads
n_groups = self.num_key_value_groups
n_kv_heads = self.config.num_key_value_heads
head_dim = self.head_dim
# assert(n_kv_heads * n_groups == n_heads)
hidden_size = self.config.hidden_size
attention_size = n_heads*head_dim
seq_len = K1.shape[-2]
kv_seq_len = seq_len + 1
# Prefill phase
# if not hasattr(self, "paged_attention"):
device = hidden_states.device
if do_prefill:
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device)
self.paged_attention_K = self.paged_attention[:,0]
self.paged_attention_V = self.paged_attention[:,1]
self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)
self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device)
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device)
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
# Mistral Nemo 12b has weird dimensions
if attention_size != hidden_size:
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
else:
self.temp_O = self.temp_QA[1][:,:,:hidden_size]
pass
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device)
self.scalar = 1.0 / math_sqrt(self.head_dim)
self.half_head_dim = head_dim // 2
elif kv_seq_len >= self.paged_attention.shape[0]:
self.paged_attention.resize_((self.paged_attention.shape[0]+KV_CACHE_INCREMENT, 2, bsz, n_kv_heads, head_dim))
self.paged_attention_K = self.paged_attention[:,0]
self.paged_attention_V = self.paged_attention[:,1]
self.attention.resize_((bsz, n_heads, 1, self.attention.shape[-1]+KV_CACHE_INCREMENT))
pass
Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0])
Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0])
Kn = Kn * self.config.key_multiplier
Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1])
Qn = Qn.view(bsz, 1, n_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
Kn = Kn.view(bsz, 1, n_kv_heads, head_dim)#.transpose(1, 2) # we will transpose after normalisation
Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
Qn = Qn.transpose(1, 2)
Kn = Kn.transpose(1, 2)
# cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
# Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
# Need to do it prior 2 steps before hitting full on short KV cache
# or else error
self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
cos, sin = self.rotary_emb.get_cached(kv_seq_len)
cos = cos[position_ids].unsqueeze(1)
sin = sin[position_ids].unsqueeze(1)
h = self.half_head_dim
RH_Q = self.RH_Q
RH_Q[:,:,:,:h] = Qn[:,:,:,h:]
RH_Q[:,:,:,h:] = Qn[:,:,:,:h]
RH_Q[:,:,:,:h].neg_() # torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h])
Qn *= cos
Qn.addcmul_(RH_Q, sin)
RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
RH_K[:,:,:,:h] = Kn[:,:,:,h:]
RH_K[:,:,:,h:] = Kn[:,:,:,:h]
RH_K[:,:,:,:h].neg_() #torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h])
Kn *= cos
Kn.addcmul_(RH_K, sin)
# New KV cache
# Kn = torch.cat([K1, Kn], dim = 2)
# Vn = torch.cat([V1, Vn], dim = 2)
self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3)
self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)
Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3)
Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3)
# Handle sliding windows
sliding_window = getattr(self.config, "sliding_window", None)
if sliding_window is not None and kv_seq_len > sliding_window:
# From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193
slicing_tokens = 1 - sliding_window
Knn = Kn[:, :, slicing_tokens:, :]#.contiguous()
Vnn = Vn[:, :, slicing_tokens:, :]#.contiguous()
else:
Knn, Vnn = Kn, Vn
pass
# Grouped query attention
_, _, cached_len, _ = Knn.shape
if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1:
Knn = Knn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
Vnn = Vnn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
pass
# else:
# Knn, Vnn = Knn, Vnn
# pass
# Attention
if bsz == 1:
Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963
# It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows
A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len])
# if attention_mask is not None: A += attention_mask # Must add attention_mask for batched
A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
A = torch_matmul(A, Vnn, out = Qn)
else:
if SDPA_HAS_GQA:
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True)
else:
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False)
pass
A = A.transpose(1, 2)
A = A.reshape(bsz, 1, attention_size)
A = fast_linear_forward(self.o_proj, A, out = self.temp_O)
return A, (Kn, Vn)
pass
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon_h1/modeling_falcon_h1.py
def FalconH1DecoderLayer_fast_forward(
self,
hidden_states: torch.Tensor,
causal_mask = None,
attention_mask: Optional[torch.Tensor] = None,
mamba_attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
cache_position: Optional[torch.LongTensor] = None,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: Optional[bool] = False,
use_cache: Optional[bool] = False,
padding_mask: Optional[torch.LongTensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
*args, **kwargs,
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
"""
Args:
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
returned tensors for more detail.
use_cache (`bool`, *optional*):
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
(see `past_key_values`).
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
"""
if use_cache and hasattr(self, "_flag_for_generation"):
residual = hidden_states
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
attention_hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states = hidden_states,
causal_mask = causal_mask,
attention_mask = attention_mask,
position_ids = position_ids,
past_key_value = past_key_value,
output_attentions = output_attentions,
use_cache = use_cache,
padding_mask = padding_mask,
position_embeddings = position_embeddings,
)
attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
mamba_hidden_states = self.mamba(
hidden_states=hidden_states,
cache_params=past_key_value,
cache_position=cache_position,
attention_mask=mamba_attention_mask,
)
mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier
hidden_states = mamba_hidden_states + attention_hidden_states
hidden_states += residual
# Fully Connected
residual = hidden_states
hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states)
hidden_states = fast_swiglu_inference(self.mlp, hidden_states)
hidden_states += residual
else:
residual = hidden_states
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)
mamba_hidden_states = self.mamba(
hidden_states=hidden_states,
cache_params=past_key_value,
cache_position=cache_position,
attention_mask=mamba_attention_mask,
)
mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier
attention_hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states = hidden_states,
causal_mask = causal_mask,
attention_mask = attention_mask,
position_ids = position_ids,
past_key_value = past_key_value,
output_attentions = output_attentions,
use_cache = use_cache,
padding_mask = padding_mask,
position_embeddings = position_embeddings,
)
attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
hidden_states = mamba_hidden_states + attention_hidden_states
# residual connection after attention + Mamba
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = fast_rms_layernorm(self.pre_ff_layernorm, hidden_states)
hidden_states = self.feed_forward(hidden_states)
hidden_states = residual + hidden_states
pass
outputs = (hidden_states,)
if output_attentions: outputs += (self_attn_weights,)
if use_cache: outputs += (present_key_value,)
return outputs
pass
def _FalconH1_fast_forward_inference(attention_fast_forward_inference=FalconH1Attention_fast_forward_inference, mlp_fast_forward_inference=fast_swiglu_inference):
# This makes the attention and MLP customisable.
# Now for models like qwen3 or cohere which use custom attention operations, we can use this function
def FalconH1Model_fast_forward_inference_custom(
self,
input_ids,
past_key_values,
position_ids,
cache_position = None,
attention_mask = None,
mamba_attention_mask = None,
):
input_ids = input_ids[:,:self.max_seq_length]
bsz, q_len = input_ids.shape
hd = self.config.hidden_size
mlp_size = self.config.intermediate_size
gate_multiplier, down_multiplier = self.config.mlp_multipliers
X = self.model.embed_tokens(input_ids)
X = X * self.config.embedding_multiplier
X = X.to(_get_dtype(self.config.torch_dtype))
bsz, q_len, hd = X.shape
assert(q_len == 1)
# Get saved buffers to reduce memory movement
residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
XX, XX2 = _XX[0], _XX[1]
variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
seq_len = past_key_values[0][0].shape[-2]
if bsz != 1:
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
attention_mask,
(bsz, q_len),
X,
seq_len,
sliding_window = getattr(self.config, "sliding_window", None),
)
else:
attention_mask = None
pass
next_decoder_cache = []
for idx, decoder_layer in enumerate(self.model.layers):
residual.copy_(X) # residual = X
X = fast_rms_layernorm_inference(
decoder_layer.input_layernorm,
X,
XX = XX,
XX2 = XX2,
variance = variance,
)
attention_hidden_states, present_key_value = attention_fast_forward_inference(
decoder_layer.self_attn,
hidden_states = X * decoder_layer.attention_in_multiplier,
past_key_value = past_key_values[idx],
position_ids = position_ids,
attention_mask = attention_mask,
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
)
attention_hidden_states = attention_hidden_states * decoder_layer.attention_out_multiplier
mamba_hidden_states = decoder_layer.mamba(
hidden_states=X,
cache_params=present_key_value,
cache_position=cache_position,
attention_mask=mamba_attention_mask,
)
mamba_hidden_states = mamba_hidden_states * decoder_layer.ssm_out_multiplier
X = mamba_hidden_states + attention_hidden_states
X += residual
residual.copy_(X) # residual = X
X = fast_rms_layernorm_inference(
decoder_layer.pre_ff_layernorm,
X,
XX = XX,
XX2 = XX2,
variance = variance,
)
X = mlp_fast_forward_inference(
decoder_layer.feed_forward,
X,
temp_gate = temp_gate,
temp_up = temp_up,
gate_multiplier = gate_multiplier,
down_multiplier = down_multiplier
)
X += residual
next_decoder_cache.append(present_key_value)
pass
X = fast_rms_layernorm_inference(
self.model.final_layernorm,
X,
XX = XX,
XX2 = XX2,
variance = variance,
)
return BaseModelOutputWithPast(
last_hidden_state = X,
past_key_values = next_decoder_cache,
hidden_states = [],
attentions = [],
)
pass
return FalconH1Model_fast_forward_inference_custom
#Separate prepare_inputs_for_generation for Hybrid FalconH1
def _fast_prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
attention_mask=None,
inputs_embeds=None,
cache_position=None,
position_ids=None,
use_cache=True,
**kwargs,):
# Overwitten -- has a unique cache type, `FalconHybridMambaAttentionDynamicCache`
empty_past_kv = past_key_values is None
# If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
# Exception 1: when passing input_embeds, input_ids may be missing entries
# Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
# Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
# (we can't check exception 3 while compiling)
if not empty_past_kv:
if (
inputs_embeds is not None # Exception 1
or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3
):
input_ids = input_ids[:, -cache_position.shape[0] :]
elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2)
input_ids = input_ids[:, cache_position]
else:
past_key_values = FalconHybridMambaAttentionDynamicCache(
self.config,
input_ids.shape[0],
self.dtype,
devices=[
self.model.layers[i].mamba.conv1d.weight.device for i in range(self.config.num_hidden_layers)
],
)
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
if not empty_past_kv:
position_ids = position_ids[:, -input_ids.shape[1] :]
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
if inputs_embeds is not None and empty_past_kv:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases
model_inputs.update(
{
"position_ids": position_ids,
"past_key_values": past_key_values,
"use_cache": use_cache,
"attention_mask": attention_mask,
"logits_to_keep": self.config.num_logits_to_keep,
"cache_position": cache_position,
}
)
return model_inputs
pass
def fix_prepare_inputs_for_generation(module):
# Fix prepare_inputs_for_generation
if hasattr(module, "prepare_inputs_for_generation"):
module.prepare_inputs_for_generation = _fast_prepare_inputs_for_generation
pass
pass
class FastFalconH1Model(FastLlamaModel):
@staticmethod
def pre_patch():
init_name, function = patch_linear_scaling(
model_name = "FalconH1",
rope_module = LlamaRotaryEmbedding,
scaled_rope_module = LlamaLinearScalingRotaryEmbedding,
attention_module = FalconH1Attention,
)
if init_name is not None:
exec(function, globals())
FalconH1Attention.__init__ = eval(init_name)
pass
FalconH1Attention .forward = FalconH1Attention_fast_forward
FalconH1DecoderLayer .forward = FalconH1DecoderLayer_fast_forward
FalconH1Model .forward = LlamaModel_fast_forward
FalconH1ForCausalLM .forward = CausalLM_fast_forward(_FalconH1_fast_forward_inference(FalconH1Attention_fast_forward_inference))
PeftModelForCausalLM.forward = PeftModel_fast_forward
fix_prepare_inputs_for_generation(FalconH1ForCausalLM)
# Solves https://github.com/unslothai/unsloth/issues/168
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
# https://github.com/huggingface/transformers/pull/27931
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
import transformers.models.falcon_h1.modeling_falcon_h1
transformers.models.falcon_h1.modeling_falcon_h1.FalconH1RotaryEmbedding = LlamaRotaryEmbedding
return
pass
@staticmethod
def from_pretrained( #TODO: Change after release
model_name = "Qwen/FalconH1-7B",
max_seq_length = 4096,
dtype = None,
load_in_4bit = True,
token = None,
device_map = "sequential",
rope_scaling = None,
fix_tokenizer = True,
model_patcher = None,
tokenizer_name = None,
trust_remote_code = False,
**kwargs,
):
return FastLlamaModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
token = token,
device_map = device_map,
rope_scaling = rope_scaling,
fix_tokenizer = fix_tokenizer,
model_patcher = FastFalconH1Model,
tokenizer_name = tokenizer_name,
trust_remote_code = trust_remote_code,
**kwargs,
)
pass
pass

42
unsloth/models/llama.py
View file

@ -340,7 +340,7 @@ pass
torch_nn_functional_silu = torch.nn.functional.silu
def fast_swiglu_inference(self, X, temp_gate = None, temp_up = None):
def fast_swiglu_inference(self, X, temp_gate = None, temp_up = None, gate_multiplier = None, down_multiplier = None):
# gate = self.gate_proj(X)
# up = self.up_proj(X)
bsz, _, hd = X.shape
@ -348,12 +348,21 @@ def fast_swiglu_inference(self, X, temp_gate = None, temp_up = None):
# temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
gate = fast_linear_forward(self.gate_proj, X, out = temp_gate)
if gate_multiplier is not None:
gate *= gate_multiplier
up = fast_linear_forward(self. up_proj, X, out = temp_up)
gate = torch_nn_functional_silu(gate, inplace = True)
gate *= up
# X = self.down_proj(gate)
down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:hd])
if down_multiplier is not None:
down *= down_multiplier
return down
pass
@ -716,6 +725,7 @@ def LlamaModel_fast_forward(
IS_GEMMA2 = self.config.model_type.startswith("gemma2")
IS_COHERE = self.config.model_type.startswith("cohere")
IS_GRANITE = self.config.model_type.startswith("granite")
IS_FALCON_H1 = self.config.model_type.startswith("falcon_h1")
train_embed_tokens = self.embed_tokens.weight.requires_grad
@ -786,8 +796,8 @@ def LlamaModel_fast_forward(
pass
hidden_states = inputs_embeds
if IS_GRANITE: #granite has embedding multiplier
hidden_states = self.embedding_multiplier * hidden_states
if IS_GRANITE or IS_FALCON_H1: #granite has embedding multiplier
hidden_states = self.config.embedding_multiplier * hidden_states
if past_key_values is None and self.training:
use_cache = False
@ -942,11 +952,16 @@ def LlamaModel_fast_forward(
# Final layernorm
if use_cache:
hidden_states = \
(fast_rms_layernorm_inference_gemma if IS_GEMMA else fast_rms_layernorm_inference)\
(self.norm, hidden_states)
if IS_FALCON_H1:
hidden_states = fast_rms_layernorm_inference(self.final_layernorm, hidden_states)
else:
hidden_states = \
(fast_rms_layernorm_inference_gemma if IS_GEMMA else fast_rms_layernorm_inference)\
(self.norm, hidden_states)
elif IS_COHERE:
hidden_states = self.norm(hidden_states)
elif IS_FALCON_H1:
hidden_states = fast_rms_layernorm(self.final_layernorm, hidden_states, gemma = IS_GEMMA)
else:
hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA)
pass
@ -1155,6 +1170,10 @@ def CausalLM_fast_forward(fast_forward_inference):
if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None:
n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None)
if self.config.model_type == "falcon_h1":
hidden_states = hidden_states * self.config.lm_head_multiplier
loss = fused_linear_cross_entropy(
hidden_states = hidden_states,
lm_weight = lm_head,
@ -1188,6 +1207,8 @@ def CausalLM_fast_forward(fast_forward_inference):
# granite: https://github.com/huggingface/transformers/blob/4d1d0f29a493098e6bc6b904b82e29cb331827f5/src/transformers/models/granite/modeling_granite.py#L1103
# cohere: https://github.com/huggingface/transformers/blob/4d1d0f29a493098e6bc6b904b82e29cb331827f5/src/transformers/models/cohere/modeling_cohere.py#L1176
logit_scaling = 1 / getattr(self.config, "logits_scaling", 1)
elif self.config.model_type == "falcon_h1":
logit_scaling = self.config.lm_head_multiplier
if labels is not None:
shift_logits = logits
@ -1814,6 +1835,8 @@ class FastLlamaModel:
# Check if RoPE Scaling is even allowed
model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__]
IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1")
has_rope_scaling = False
try:
with open(inspect.getfile(model_function), "r") as file:
@ -1856,12 +1879,16 @@ class FastLlamaModel:
bnb_config = None
if load_in_4bit:
llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy()
if IS_FALCON_H1:
# we cannot quantize out_proj layer due to mamba kernels: https://github.com/tiiuae/Falcon-H1/issues/13#issuecomment-2918671274
llm_int8_skip_modules.append("out_proj")
bnb_config = BitsAndBytesConfig(
load_in_4bit = True,
bnb_4bit_use_double_quant = True,
bnb_4bit_quant_type = "nf4",
bnb_4bit_compute_dtype = dtype,
llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy(),
llm_int8_skip_modules = llm_int8_skip_modules,
)
pass
@ -2607,6 +2634,7 @@ class FastLlamaModel:
elif model_type == "cohere": apply_lora_mlp = apply_lora_mlp_swiglu
elif model_type == "granite": apply_lora_mlp = apply_lora_mlp_swiglu
elif model_type == "qwen3": apply_lora_mlp = apply_lora_mlp_swiglu
elif model_type == "falcon_h1": apply_lora_mlp = apply_lora_mlp_swiglu
elif model_type == "qwen3moe": apply_lora_mlp = apply_lora_mlp_swiglu
else:
raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!")

20
unsloth/models/loader.py
View file

@ -56,12 +56,17 @@ SUPPORTS_LLAMA32 = transformers_version > Version("4.45.0")
SUPPORTS_GRANITE = transformers_version >= Version("4.46.0")
SUPPORTS_QWEN3 = transformers_version >= Version("4.50.3")
SUPPORTS_QWEN3_MOE = transformers_version >= Version("4.50.3")
SUPPORTS_FALCON_H1 = transformers_version >= Version("4.53.0")
SUPPORTS_GEMMA3N = transformers_version >= Version("4.53.0")
if SUPPORTS_GEMMA:
from .gemma import FastGemmaModel
if SUPPORTS_GEMMA2:
from .gemma2 import FastGemma2Model
pass
if SUPPORTS_FALCON_H1:
from .falcon_h1 import FastFalconH1Model
pass
import torch
from ._utils import (
patch_compiling_bitsandbytes,
@ -129,6 +134,8 @@ class FastLanguageModel(FastLlamaModel):
pass
if token is None: token = get_token()
if isinstance(dtype, str) and dtype in ["float16", "bfloat16"]:
dtype = getattr(torch, dtype)
assert (dtype is None or dtype == torch.float16 or dtype == torch.bfloat16)
if use_gradient_checkpointing == "unsloth":
@ -313,6 +320,15 @@ class FastLanguageModel(FastLlamaModel):
f"to obtain the latest transformers build, then restart this session."\
)
dispatch_model = FastQwen3Model if model_type == "qwen3" else FastQwen3MoeModel
elif model_type == "falcon_h1":
dispatch_model = FastFalconH1Model
if not SUPPORTS_FALCON_H1:
raise ImportError(
f"Unsloth: Your transformers version of {transformers_version} does not support FalconH1.\n"\
f"The minimum required version is 4.50.3.\n"\
f'Try `pip install --upgrade "transformers>=4.50.3"`\n'\
f"to obtain the latest transformers build, then restart this session."\
)
# Temporary disable optimized Cohere until errors match
# elif model_type == "cohere":
# dispatch_model = FastCohereModel
@ -542,6 +558,10 @@ class FastModel(FastBaseModel):
elif "csm-1b" in lowered_model_name:
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" # Sesame fails
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.float16;if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
elif 'granite-4' in lowered_model_name:
# granite-4 rms norms are stored as 16 bit, but we upcast
os.environ["UNSLOTH_UPCAST_LAYERNORM"] = "1"
os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
elif "olmo-2" in lowered_model_name and transformers_version < Version("4.50.0.dev0"):
raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)
elif "gemma-3n" in lowered_model_name: