From 267ca3b209042bc76e54d7074e824fd63fa8b605 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Fri, 27 Jun 2025 04:54:55 -0700 Subject: [PATCH 1/4] add instructions for installing on blackwell (#2812) --- blackwell/README.md | 87 ++++++ blackwell/blackwell.requirements.txt | 180 +++++++++++ blackwell/test_llama32_sft.py | 178 +++++++++++ blackwell/test_qwen3_grpo.py | 429 +++++++++++++++++++++++++++ 4 files changed, 874 insertions(+) create mode 100644 blackwell/README.md create mode 100644 blackwell/blackwell.requirements.txt create mode 100644 blackwell/test_llama32_sft.py create mode 100644 blackwell/test_qwen3_grpo.py diff --git a/blackwell/README.md b/blackwell/README.md new file mode 100644 index 0000000000..fdf5d2bc90 --- /dev/null +++ b/blackwell/README.md @@ -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`. diff --git a/blackwell/blackwell.requirements.txt b/blackwell/blackwell.requirements.txt new file mode 100644 index 0000000000..124e7eb119 --- /dev/null +++ b/blackwell/blackwell.requirements.txt @@ -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 diff --git a/blackwell/test_llama32_sft.py b/blackwell/test_llama32_sft.py new file mode 100644 index 0000000000..6448343182 --- /dev/null +++ b/blackwell/test_llama32_sft.py @@ -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)) + diff --git a/blackwell/test_qwen3_grpo.py b/blackwell/test_qwen3_grpo.py new file mode 100644 index 0000000000..bf56ad338c --- /dev/null +++ b/blackwell/test_qwen3_grpo.py @@ -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 = "" +reasoning_end = "" +solution_start = "" +solution_end = "" + +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("", "").replace("", "") + + 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"[\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!\n2\n", +) + +match_format.findall( + f"Let me think! 2 \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(" 0.34 ")) +print(match_numbers.findall(" 123,456 ")) +print(match_numbers.findall(" -0.234 ")) +print(match_numbers.findall("17")) + +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) From 96941573e66ceeccc47bcfca9d49c083a69ab95b Mon Sep 17 00:00:00 2001 From: Dhia Eddine Rhaiem <163106757+dhiaEddineRhaiem@users.noreply.github.com> Date: Sat, 28 Jun 2025 15:55:13 +0400 Subject: [PATCH 2/4] Add falcon h1 (#2650) * add falcon h1 * feat: add Falcon-H1 into unsloth * address comments * fix * Update unsloth/models/llama.py Co-authored-by: Daniel Han * Update unsloth/models/llama.py Co-authored-by: Daniel Han * fixes * fix comments --------- Co-authored-by: Daniel Han Co-authored-by: Younes B <49240599+younesbelkada@users.noreply.github.com> Co-authored-by: Younes Belkada Co-authored-by: ilyasch2 --- unsloth/models/_utils.py | 2 +- unsloth/models/falcon_h1.py | 699 ++++++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 42 ++- unsloth/models/loader.py | 14 + 4 files changed, 749 insertions(+), 8 deletions(-) create mode 100644 unsloth/models/falcon_h1.py diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7f2c2db105..a31bf7b604 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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}" diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py new file mode 100644 index 0000000000..df4622d5c0 --- /dev/null +++ b/unsloth/models/falcon_h1.py @@ -0,0 +1,699 @@ +# 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: + transformers_version = Version(transformers_version) + if not transformers_version >= Version("4.50.3"): #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.50.3.\n"\ + f'Try `pip install --upgrade "transformers>=4.50.3"`\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: + FalconH1Attention = FalconH1Attention +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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c83de8f944..bc46ba177b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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!") diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 34bf05cd8d..7800ec2b0a 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -26,6 +26,7 @@ from .mistral import FastMistralModel from .qwen2 import FastQwen2Model from .qwen3 import FastQwen3Model from .qwen3_moe import FastQwen3MoeModel +from .falcon_h1 import FastFalconH1Model from .cohere import FastCohereModel from transformers import AutoConfig from transformers import __version__ as transformers_version @@ -56,7 +57,9 @@ 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: @@ -129,6 +132,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 +318,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 From ebe935ee7417dc1783e850574bf3e87a16ed567e Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sat, 28 Jun 2025 07:27:56 -0500 Subject: [PATCH 3/4] granite force layernorm upcast (#2799) --- unsloth/models/loader.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 7800ec2b0a..f7e08ed10d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -556,6 +556,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: From be38f500d0adbb42f16c619bfb0b0f8e18b8e923 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sat, 28 Jun 2025 19:41:19 -0500 Subject: [PATCH 4/4] import undefined transformers_version for falcon model (#2822) * import undefined transformers_version for falcon model fixed falcon transformers version check and added error handling for FalconH1Attention bad import * Also, conditionally load module from falcon_h1 depending on if the transformers version supports is --- unsloth/models/falcon_h1.py | 12 ++++++++---- unsloth/models/loader.py | 4 +++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index df4622d5c0..8978e4db0d 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -30,12 +30,13 @@ try: FalconHybridMambaAttentionDynamicCache, ) except: + from transformers import __version__ as transformers_version transformers_version = Version(transformers_version) - if not transformers_version >= Version("4.50.3"): #TODO: Update when transformers is updated + 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.50.3.\n"\ - f'Try `pip install --upgrade "transformers>=4.50.3"`\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 @@ -48,7 +49,10 @@ try: FalconH1Attention, ) except: - FalconH1Attention = FalconH1Attention + # 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 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index f7e08ed10d..70dd896514 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -26,7 +26,6 @@ from .mistral import FastMistralModel from .qwen2 import FastQwen2Model from .qwen3 import FastQwen3Model from .qwen3_moe import FastQwen3MoeModel -from .falcon_h1 import FastFalconH1Model from .cohere import FastCohereModel from transformers import AutoConfig from transformers import __version__ as transformers_version @@ -65,6 +64,9 @@ if SUPPORTS_GEMMA: 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,