Compare commits
3 commits
main
...
reduce-com
| Author | SHA1 | Date | |
|---|---|---|---|
|
c8a76c78ea | ||
|
|
acc19887a4 | ||
|
|
0050c73e05 |
61 changed files with 667 additions and 1681 deletions
|
|
@ -16,11 +16,8 @@ import os, importlib.util, platform
|
|||
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
# ── Windows console UTF-8 safety ─────────────────────────────────────────────
|
||||
# Legacy Windows consoles (cp1252) can't encode Unsloth's emoji/box-drawing
|
||||
# glyphs and crash with UnicodeEncodeError. Force stdout/stderr to UTF-8 only on
|
||||
# Windows and only when not already UTF-8; no-op elsewhere. errors="replace"
|
||||
# guarantees we never crash on an unencodable glyph.
|
||||
# Force stdout/stderr to UTF-8 on Windows: legacy cp1252 consoles crash on
|
||||
# Unsloth's emoji/box-drawing glyphs. errors="replace" avoids unencodable-glyph crashes.
|
||||
if platform.system() == "Windows":
|
||||
import sys as _sys
|
||||
for _name in ("stdout", "stderr"):
|
||||
|
|
@ -34,9 +31,8 @@ if platform.system() == "Windows":
|
|||
|
||||
|
||||
def _is_mlx_available():
|
||||
# Transitional import barrier: keep non-Apple-Silicon imports from touching
|
||||
# unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this
|
||||
# can collapse back to the centralized zoo runtime call below.
|
||||
# Transitional barrier: avoid importing unsloth_zoo on GPU hosts until
|
||||
# unsloth_zoo.mlx is import-safe there.
|
||||
if (
|
||||
os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1"
|
||||
or platform.system() != "Darwin"
|
||||
|
|
@ -62,9 +58,8 @@ if _IS_MLX:
|
|||
"Unsloth: MLX support requires `unsloth-zoo` with MLX modules. "
|
||||
"Reinstall with `pip install unsloth-zoo` or rerun install.sh."
|
||||
) from _e
|
||||
# An older unsloth-zoo satisfies `import unsloth_zoo` but lacks the
|
||||
# mlx.trainer / mlx.loader submodules. Surface a friendly install hint
|
||||
# instead of a raw ImportError on the submodule path.
|
||||
# Older unsloth-zoo imports fine but lacks mlx.trainer/mlx.loader; give a
|
||||
# friendly install hint instead of a raw submodule ImportError.
|
||||
try:
|
||||
from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig
|
||||
from unsloth_zoo.mlx.loader import FastMLXModel
|
||||
|
|
@ -75,8 +70,8 @@ if _IS_MLX:
|
|||
"`pip install -U unsloth-zoo` or rerun install.sh."
|
||||
) from _e
|
||||
|
||||
# Load raw_text helpers without executing dataprep/__init__.py, which
|
||||
# imports synthetic.py -> torch and would defeat the torch-free MLX path.
|
||||
# Load raw_text helpers directly: dataprep/__init__.py imports torch via
|
||||
# synthetic.py, which would break the torch-free MLX path.
|
||||
from pathlib import Path as _Path
|
||||
|
||||
_raw_text_path = _Path(__file__).resolve().parent / "dataprep" / "raw_text.py"
|
||||
|
|
|
|||
|
|
@ -120,7 +120,6 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|||
from importlib.metadata import version as importlib_version
|
||||
from importlib.metadata import PackageNotFoundError
|
||||
|
||||
# Check for unsloth_zoo
|
||||
try:
|
||||
unsloth_zoo_version = importlib_version("unsloth_zoo")
|
||||
if Version(unsloth_zoo_version) < Version("2026.5.2"):
|
||||
|
|
@ -145,7 +144,7 @@ except:
|
|||
raise
|
||||
del PackageNotFoundError, importlib_version
|
||||
|
||||
# Try importing PyTorch and check version
|
||||
# Try importing PyTorch
|
||||
try:
|
||||
import torch
|
||||
except ModuleNotFoundError:
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ zephyr_ollama = _ollama_template("zephyr")
|
|||
|
||||
zephyr_eos_token = "eos_token"
|
||||
CHAT_TEMPLATES["zephyr"] = (zephyr_template, zephyr_eos_token, False, zephyr_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["zephyr"] = None # No system message in Zephyr
|
||||
DEFAULT_SYSTEM_MESSAGE["zephyr"] = None
|
||||
|
||||
# =========================================== ChatML
|
||||
# ChatML has no BOS and not EOS! Rather <|im_start|> and <|im_end|> acts as BOS / EOS.
|
||||
|
|
@ -117,7 +117,7 @@ chatml_ollama = _ollama_template("chatml")
|
|||
|
||||
chatml_eos_token = "<|im_end|>"
|
||||
CHAT_TEMPLATES["chatml"] = (chatml_template, chatml_eos_token, True, chatml_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["chatml"] = None # No system message in ChatML
|
||||
DEFAULT_SYSTEM_MESSAGE["chatml"] = None
|
||||
|
||||
# =========================================== Mistral-1
|
||||
# Mistral Instruct doesn't allow system prompts, so we append it to the user message.
|
||||
|
|
@ -149,7 +149,7 @@ mistral_ollama = _ollama_template("mistral")
|
|||
|
||||
mistral_eos_token = "eos_token"
|
||||
CHAT_TEMPLATES["mistral"] = (mistral_template, mistral_eos_token, False, mistral_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["mistral"] = None # No system message in Mistral
|
||||
DEFAULT_SYSTEM_MESSAGE["mistral"] = None
|
||||
|
||||
# =========================================== Llama-2
|
||||
# Adds BOS to every convo! And weird <<SYS>> system messages.
|
||||
|
|
@ -180,7 +180,7 @@ llama_ollama = _ollama_template("llama")
|
|||
|
||||
llama_eos_token = "eos_token"
|
||||
CHAT_TEMPLATES["llama"] = (llama_template, llama_eos_token, False, llama_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["llama"] = None # No system message in Llama
|
||||
DEFAULT_SYSTEM_MESSAGE["llama"] = None
|
||||
|
||||
# =========================================== Vicuna
|
||||
# https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template
|
||||
|
|
@ -304,7 +304,7 @@ gemma_ollama = _ollama_template("gemma")
|
|||
|
||||
gemma_eos_token = "<end_of_turn>"
|
||||
CHAT_TEMPLATES["gemma"] = (gemma_template, gemma_eos_token, True, gemma_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma"] = None # No system message in Gemma
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma"] = None
|
||||
|
||||
# =========================================== Gemma with ChatML instead
|
||||
# We find using <eos> is still more appropriate!
|
||||
|
|
@ -317,7 +317,7 @@ gemma_chatml_eos_token = (
|
|||
"<|im_end|>",
|
||||
)
|
||||
CHAT_TEMPLATES["gemma_chatml"] = (gemma_chatml_template, gemma_chatml_eos_token, True, gemma_chatml_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma_chatml"] = None # No system message in Gemma
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma_chatml"] = None
|
||||
|
||||
# =========================================== Gemma 2
|
||||
# Same as Gemma 1, but with sliding window attention!
|
||||
|
|
@ -326,14 +326,14 @@ gemma2_template = gemma_template
|
|||
gemma2_ollama = _ollama_template("gemma2")
|
||||
gemma2_eos_token = "<end_of_turn>"
|
||||
CHAT_TEMPLATES["gemma2"] = (gemma2_template, gemma2_eos_token, True, gemma2_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma2"] = None # No system message in Gemma 2
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma2"] = None
|
||||
|
||||
# =========================================== Gemma 2 with ChatML instead
|
||||
gemma2_chatml_template = gemma_chatml_template
|
||||
gemma2_chatml_ollama = _ollama_template("gemma2_chatml")
|
||||
gemma2_chatml_eos_token = gemma_chatml_eos_token
|
||||
CHAT_TEMPLATES["gemma2_chatml"] = (gemma2_chatml_template, gemma2_chatml_eos_token, True, gemma2_chatml_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma2_chatml"] = None # No system message in Gemma 2
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma2_chatml"] = None
|
||||
|
||||
# =========================================== Llama-3
|
||||
# Weirdly \n\n is needed?
|
||||
|
|
@ -358,10 +358,10 @@ llama3_ollama = _ollama_template("llama-3")
|
|||
llama3_template_eos_token = "eos_token"
|
||||
|
||||
CHAT_TEMPLATES["llama-3"] = (llama3_template, llama3_template_eos_token, False, llama3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["llama-3"] = None # No system message in Llama-3
|
||||
DEFAULT_SYSTEM_MESSAGE["llama-3"] = None
|
||||
|
||||
CHAT_TEMPLATES["llama3"] = (llama3_template, llama3_template_eos_token, False, llama3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["llama3"] = None # No system message in Llama-3
|
||||
DEFAULT_SYSTEM_MESSAGE["llama3"] = None
|
||||
|
||||
|
||||
# =========================================== Phi-3
|
||||
|
|
@ -385,13 +385,13 @@ phi3_ollama = _ollama_template("phi-3")
|
|||
|
||||
phi3_template_eos_token = "<|end|>"
|
||||
CHAT_TEMPLATES["phi-3"] = (phi3_template, phi3_template_eos_token, False, phi3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-3"] = None # No system message in Phi-3
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-3"] = None
|
||||
|
||||
CHAT_TEMPLATES["phi-35"] = CHAT_TEMPLATES["phi-3"]
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-35"] = None # No system message in Phi-3.5
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-35"] = None
|
||||
|
||||
CHAT_TEMPLATES["phi-3.5"] = CHAT_TEMPLATES["phi-3"]
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-3.5"] = None # No system message in Phi-3.5
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-3.5"] = None
|
||||
|
||||
# =========================================== Llama-3.1
|
||||
"""
|
||||
|
|
@ -596,16 +596,16 @@ qwen25_ollama = _ollama_template("qwen-2.5")
|
|||
qwen25_template_eos_token = "eos_token"
|
||||
qwen25_default_system_message = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
|
||||
CHAT_TEMPLATES["qwen-2.5"] = (qwen25_template, qwen25_template_eos_token, False, qwen25_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-2.5"] = qwen25_default_system_message # No system message in Qwen 2.5
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-2.5"] = qwen25_default_system_message
|
||||
|
||||
CHAT_TEMPLATES["qwen-25"] = (qwen25_template, qwen25_template_eos_token, False, qwen25_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-25"] = qwen25_default_system_message # No system message in Qwen 2.5
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-25"] = qwen25_default_system_message
|
||||
|
||||
CHAT_TEMPLATES["qwen25"] = (qwen25_template, qwen25_template_eos_token, False, qwen25_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen25"] = qwen25_default_system_message # No system message in Qwen 2.5
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen25"] = qwen25_default_system_message
|
||||
|
||||
CHAT_TEMPLATES["qwen2.5"] = (qwen25_template, qwen25_template_eos_token, False, qwen25_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen2.5"] = qwen25_default_system_message # No system message in Qwen 2.5
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen2.5"] = qwen25_default_system_message
|
||||
|
||||
# =========================================== Phi-4
|
||||
# "{{ bos_token }}"\ # Phi-4 removes BOS?
|
||||
|
|
@ -633,7 +633,7 @@ phi4_ollama = _ollama_template("phi-4")
|
|||
|
||||
phi4_template_eos_token = "<|im_end|>"
|
||||
CHAT_TEMPLATES["phi-4"] = (phi4_template, phi4_template_eos_token, False, phi4_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-4"] = None # No system message in Phi-4
|
||||
DEFAULT_SYSTEM_MESSAGE["phi-4"] = None
|
||||
|
||||
|
||||
# =========================================== Gemma-3
|
||||
|
|
@ -687,10 +687,10 @@ gemma3_ollama = _ollama_template("gemma-3")
|
|||
|
||||
gemma3_template_eos_token = "<end_of_turn>"
|
||||
CHAT_TEMPLATES["gemma-3"] = (gemma3_template, gemma3_template_eos_token, False, gemma3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma-3"] = None # No system message in Gemma-3
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma-3"] = None
|
||||
|
||||
CHAT_TEMPLATES["gemma3"] = (gemma3_template, gemma3_template_eos_token, False, gemma3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma3"] = None # No system message in Gemma-3
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma3"] = None
|
||||
|
||||
# =========================================== Qwen-3
|
||||
# Official Qwen-3 chat template (see https://ollama.com/library/qwen3/blobs/eb4402837c78)
|
||||
|
|
@ -799,10 +799,10 @@ qwen3_template = \
|
|||
qwen3_ollama = _ollama_template("qwen-3")
|
||||
qwen3_template_eos_token = "<|im_end|>"
|
||||
CHAT_TEMPLATES["qwen-3"] = (qwen3_template, qwen3_template_eos_token, False, qwen3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-3"] = None # No default system message for Qwen-3
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen-3"] = None
|
||||
|
||||
CHAT_TEMPLATES["qwen3"] = (qwen3_template, qwen3_template_eos_token, False, qwen3_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3"] = None # No default system message for Qwen-3
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3"] = None
|
||||
|
||||
# =========================================== Gemma-3n
|
||||
# Obtained via
|
||||
|
|
@ -856,10 +856,10 @@ gemma3n_template = \
|
|||
gemma3n_ollama = _ollama_template("gemma-3n")
|
||||
gemma3n_template_eos_token = "<end_of_turn>"
|
||||
CHAT_TEMPLATES["gemma-3n"] = (gemma3n_template, gemma3n_template_eos_token, False, gemma3n_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma-3n"] = None # No system message in Gemma-3n
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma-3n"] = None
|
||||
|
||||
CHAT_TEMPLATES["gemma3n"] = (gemma3n_template, gemma3n_template_eos_token, False, gemma3n_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma3n"] = None # No system message in Gemma-3n
|
||||
DEFAULT_SYSTEM_MESSAGE["gemma3n"] = None
|
||||
|
||||
# =========================================== Gemma-4
|
||||
# Gemma-4 uses <|turn>role\n...<turn|>\n format
|
||||
|
|
@ -1555,10 +1555,10 @@ PARAMETER top_p 1.0
|
|||
|
||||
gptoss_template_template_eos_token = "<|return|>"
|
||||
CHAT_TEMPLATES["gpt-oss"] = (gptoss_template, gptoss_template_template_eos_token, False, gptoss_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gpt-oss"] = None # No system message in GPT-oss
|
||||
DEFAULT_SYSTEM_MESSAGE["gpt-oss"] = None
|
||||
|
||||
CHAT_TEMPLATES["gptoss"] = (gptoss_template, gptoss_template_template_eos_token, False, gptoss_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["gptoss"] = None # No system message in GPT-oss
|
||||
DEFAULT_SYSTEM_MESSAGE["gptoss"] = None
|
||||
|
||||
# =========================================== Qwen3-Instruct
|
||||
qwen3_instruct_template = \
|
||||
|
|
@ -1651,7 +1651,7 @@ qwen3_instruct_template = \
|
|||
|
||||
qwen3_template_eos_token = "<|im_end|>"
|
||||
CHAT_TEMPLATES["qwen3-instruct"] = (qwen3_instruct_template, qwen3_template_eos_token, False, _ollama_template("qwen3-instruct"),)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3-instruct"] = None # No system message in Qwen3
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3-instruct"] = None
|
||||
|
||||
|
||||
# =========================================== Qwen3-Thinking
|
||||
|
|
@ -1749,7 +1749,7 @@ CHAT_TEMPLATES["qwen3-thinking"] = (
|
|||
False,
|
||||
_ollama_template("qwen3-thinking"),
|
||||
)
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3-thinking"] = None # No system message in Qwen3
|
||||
DEFAULT_SYSTEM_MESSAGE["qwen3-thinking"] = None
|
||||
|
||||
|
||||
# =========================================== Liquid-LFM2
|
||||
|
|
@ -1762,7 +1762,7 @@ liquid_lfm2_template = \
|
|||
|
||||
liquid_lfm2_template_eos_token = "<|im_end|>"
|
||||
CHAT_TEMPLATES["lfm-2"] = (liquid_lfm2_template, liquid_lfm2_template_eos_token, False, None)
|
||||
DEFAULT_SYSTEM_MESSAGE["lfm-2"] = None # No system message in Phi-3
|
||||
DEFAULT_SYSTEM_MESSAGE["lfm-2"] = None
|
||||
CHAT_TEMPLATES["lfm-2.5"] = (liquid_lfm2_template, liquid_lfm2_template_eos_token, False, None)
|
||||
DEFAULT_SYSTEM_MESSAGE["lfm-2.5"] = None
|
||||
|
||||
|
|
@ -1864,7 +1864,7 @@ def get_chat_template(
|
|||
# pass
|
||||
# pass
|
||||
|
||||
# We first check if the tokenizer is a fast one. If not, we cannot convert this!
|
||||
# Non-fast tokenizers cannot be converted
|
||||
is_fast_tokenizer = getattr(tokenizer, "is_fast", False)
|
||||
old_padding_side = tokenizer.padding_side
|
||||
|
||||
|
|
@ -1872,8 +1872,7 @@ def get_chat_template(
|
|||
type_chat_template = None
|
||||
|
||||
if type(chat_template) in (list, tuple,):
|
||||
# For changing system message later
|
||||
# Since it's not supported yet, we will raise an error first!
|
||||
# type_chat_template lets us swap the system message later
|
||||
type_chat_template = chat_template[0].lower()
|
||||
chat_template, stop_word = chat_template
|
||||
assert(type(chat_template) is str)
|
||||
|
|
@ -1881,12 +1880,10 @@ def get_chat_template(
|
|||
ollama_modelfile = None
|
||||
|
||||
elif type(chat_template) is str:
|
||||
# For changing system message later
|
||||
type_chat_template = chat_template.lower()
|
||||
|
||||
chat_template, stop_word, yes_map_eos_token, ollama_modelfile = CHAT_TEMPLATES[chat_template]
|
||||
|
||||
# Check mapping to eos_token
|
||||
if not map_eos_token and yes_map_eos_token: map_eos_token = True
|
||||
if not yes_map_eos_token and map_eos_token: map_eos_token = False
|
||||
|
||||
|
|
@ -1974,9 +1971,8 @@ def get_chat_template(
|
|||
old_unk_token = getattr(tokenizer, "unk_token", None)
|
||||
|
||||
string_vocab = tokenizer._tokenizer.to_str()
|
||||
# First check if new stop_word is in the tokenizer
|
||||
if stop_word in string_vocab:
|
||||
# We shall swap them around
|
||||
# swap stop_word and the old EOS around
|
||||
temporary_stop_token = "<|:__TEMP//STOP//TOKEN__:|>"
|
||||
string_vocab = string_vocab.replace(old_eos_token, temporary_stop_token)
|
||||
string_vocab = string_vocab.replace(stop_word, old_eos_token)
|
||||
|
|
@ -2047,7 +2043,7 @@ def get_chat_template(
|
|||
|
||||
tokenizer.chat_template = chat_template
|
||||
|
||||
# Also fix up other tokens
|
||||
# Restore the original special tokens
|
||||
old_pad_token = getattr(old_tokenizer, "pad_token", None)
|
||||
old_bos_token = getattr(old_tokenizer, "bos_token", None)
|
||||
old_unk_token = getattr(old_tokenizer, "unk_token", None)
|
||||
|
|
@ -2061,19 +2057,16 @@ def get_chat_template(
|
|||
|
||||
# stopping_criteria = create_stopping_criteria(tokenizer, stop_word)
|
||||
|
||||
# Patch saving functions
|
||||
if patch_saving:
|
||||
from .save import patch_saving_functions
|
||||
tokenizer = patch_saving_functions(tokenizer)
|
||||
|
||||
# Add Ollama
|
||||
tokenizer._ollama_modelfile = ollama_modelfile
|
||||
tokenizer._system_message = system_message
|
||||
return tokenizer#, stopping_criteria
|
||||
|
||||
|
||||
def remove_special_tokens(tokenizer, prompt):
|
||||
# Removes double BOS token
|
||||
if prompt.startswith(tokenizer.bos_token):
|
||||
prompt = prompt[len(tokenizer.bos_token):]
|
||||
return prompt
|
||||
|
|
@ -2096,24 +2089,20 @@ def _parse_combined_prompt(combined_prompt, dataset):
|
|||
|
||||
final_optional_prompts = []
|
||||
if len(optional_prompts) != 0:
|
||||
# Add left
|
||||
left = optional_prompts[0]
|
||||
l = left[0][0]
|
||||
if l != 0: final_optional_prompts.append(combined_prompt[:l])
|
||||
|
||||
# Add in between
|
||||
for left, right in zip(optional_prompts[:-1], optional_prompts[1:]):
|
||||
l, r = left[0][-1], right[0][0]
|
||||
final_optional_prompts.append(left)
|
||||
if l != r: final_optional_prompts.append(combined_prompt[l : r])
|
||||
final_optional_prompts.append(optional_prompts[-1])
|
||||
|
||||
# Add right
|
||||
right = optional_prompts[-1]
|
||||
r = right[0][1]
|
||||
if r != len(combined_prompt): final_optional_prompts.append(combined_prompt[r:])
|
||||
else:
|
||||
# Just add in the entire string
|
||||
final_optional_prompts.append(combined_prompt)
|
||||
|
||||
check_combined = "".join(x if type(x) is str else x[1] for x in final_optional_prompts)
|
||||
|
|
@ -2189,15 +2178,9 @@ def to_sharegpt(
|
|||
conversation_extension = 1,
|
||||
random_state = 3407,
|
||||
):
|
||||
"""
|
||||
Converts a dataset to ShareGPT style (1 input + 1 output field).
|
||||
Merge multiple columns into 1 input via `merged_prompt`; use
|
||||
`conversation_extension` to pack several convos into one.
|
||||
|
||||
merged_prompt = "", Prompt to merge columns into 1 input
|
||||
merged_column_name = "instruction", Final column name for the input field
|
||||
output_column_name = "output", Final column name for the output field
|
||||
conversation_extension = 1, Combines this many convos into 1
|
||||
"""Convert a dataset to ShareGPT style (1 input + 1 output field).
|
||||
`merged_prompt` merges multiple columns into the input;
|
||||
`conversation_extension` packs that many convos into one.
|
||||
"""
|
||||
if "conversations" in dataset.column_names:
|
||||
convo = dataset[0]["conversations"]
|
||||
|
|
@ -2229,11 +2212,10 @@ def to_sharegpt(
|
|||
__convert_to_sharegpt__,
|
||||
batched = True,
|
||||
desc = "Converting to ShareGPT",
|
||||
# Remove unused columns!
|
||||
remove_columns = dataset.column_names if remove_unused_columns else None,
|
||||
)
|
||||
|
||||
# Randomnly concat conversations to create a long stream!
|
||||
# Randomly concat conversations to create a long stream!
|
||||
from datasets import concatenate_datasets
|
||||
n_extensions = max(conversation_extension-1, 0)
|
||||
if n_extensions == 0: return dataset
|
||||
|
|
@ -2245,7 +2227,6 @@ def to_sharegpt(
|
|||
all_shuffled.append(shuffled)
|
||||
dataset = concatenate_datasets(all_shuffled, axis = 1)
|
||||
|
||||
# Combine them into 1
|
||||
n_extensions += 1
|
||||
conversation_columns = [f"conversations{j}" for j in range(n_extensions)]
|
||||
def __combine_conversations__(examples):
|
||||
|
|
@ -2262,7 +2243,6 @@ def to_sharegpt(
|
|||
__combine_conversations__,
|
||||
batched = True,
|
||||
desc = "Extending conversations",
|
||||
# Remove unused columns!
|
||||
remove_columns = dataset.column_names if remove_unused_columns else None,
|
||||
)
|
||||
return dataset
|
||||
|
|
@ -2272,15 +2252,12 @@ def get_ollama_eos_tokens(tokenizer, extra_eos_tokens = []):
|
|||
added_tokens_decoder = tokenizer.added_tokens_decoder.values()
|
||||
added_tokens_decoder = [str(x) for x in added_tokens_decoder]
|
||||
|
||||
# Remove added_tokens_decoder duplicates
|
||||
added_tokens_decoder = list(set(added_tokens_decoder) - set(extra_eos_tokens))
|
||||
|
||||
# Remove BOS
|
||||
if getattr(tokenizer, "bos_token", None) is not None:
|
||||
added_tokens_decoder = [x for x in added_tokens_decoder if x != tokenizer.bos_token]
|
||||
|
||||
repeatted_tokens = []
|
||||
# Join all vocab
|
||||
joined_text = "\x01\x00".join(added_tokens_decoder)
|
||||
for token in added_tokens_decoder:
|
||||
n = len(token)
|
||||
|
|
@ -2296,13 +2273,12 @@ def get_ollama_eos_tokens(tokenizer, extra_eos_tokens = []):
|
|||
repeatted_tokens.append(token[:j])
|
||||
break
|
||||
|
||||
# Remove duplicates
|
||||
splitted = joined_text.split("\x01\x00")
|
||||
final_eos_tokens = [old for old, new in zip(added_tokens_decoder, splitted) if old == new]
|
||||
final_eos_tokens += extra_eos_tokens
|
||||
final_eos_tokens += repeatted_tokens
|
||||
|
||||
# Remove new lines, spaces and HTML tags
|
||||
# Drop newline / space / short HTML-tag tokens
|
||||
filtered_eos_tokens = []
|
||||
for token in final_eos_tokens:
|
||||
if token.count("\n") == len(token): continue
|
||||
|
|
@ -2334,12 +2310,8 @@ default_system_message = \
|
|||
|
||||
extra_eos_tokens = None,
|
||||
):
|
||||
"""
|
||||
Creates an Ollama modelfile and a HF Jinja template from a custom
|
||||
template. You must provide 2x examples of an input & output.
|
||||
There is an optional system message as well.
|
||||
|
||||
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
|
||||
"""Build an Ollama modelfile and HF Jinja template from a custom template.
|
||||
Use {INPUT} and {OUTPUT} twice each; {SYSTEM} is optional.
|
||||
"""
|
||||
# Strip only the left: trailing whitespace can be part of the repeated example
|
||||
# (e.g. "{OUTPUT}\n"). Accidental trailing whitespace (#992) is retried on failure.
|
||||
|
|
@ -2374,14 +2346,11 @@ extra_eos_tokens = None,
|
|||
"Unsloth: Your tokenizer does not have an EOS token? Please provide one via extra_eos_tokens!"
|
||||
)
|
||||
|
||||
# Check tokenizer types
|
||||
tokenizer_name = tokenizer.name_or_path.lower()
|
||||
if tokenizer_name.startswith(("unsloth/llama-3-8b-instruct", "unsloth/llama-3-70b-instruct")):
|
||||
# Add <|eot_id|>
|
||||
extra_eos_tokens.append("<|eot_id|>")
|
||||
elif ("<|eot_id|>" in extra_eos_tokens or "<|eot_id|>" in chat_template) and \
|
||||
tokenizer_name.startswith(("unsloth/llama-3-8b", "unsloth/llama-3-70b")):
|
||||
# Warn
|
||||
logger.warning(
|
||||
"Unsloth: Base llama-3 models did not train <|eot_id|>.\n"\
|
||||
"Please use the instruct version or use <|end_of_text|>"
|
||||
|
|
@ -2412,7 +2381,6 @@ extra_eos_tokens = None,
|
|||
# Must be equivalent to left
|
||||
final_combined_check = True
|
||||
|
||||
# Repeatted text
|
||||
instruction_response = chat_template[j:]
|
||||
if instruction_response.count("{INPUT}") != 1 or instruction_response.count("{OUTPUT}") != 1:
|
||||
raise RuntimeError(error_msg)
|
||||
|
|
@ -2522,8 +2490,6 @@ extra_eos_tokens = None,
|
|||
eos = extra_eos_tokens[0]
|
||||
output_part = output_part + eos
|
||||
|
||||
# Ollama modelfile parts
|
||||
|
||||
# Check bos_token is in system prompt
|
||||
ollama_system = system_part
|
||||
has_bos_token = False
|
||||
|
|
@ -2541,14 +2507,12 @@ extra_eos_tokens = None,
|
|||
input_modelfile = "{{ if .Prompt }}" + input_part .replace("{INPUT}", "{{ .Prompt }}") + "{{ end }}"
|
||||
output_modelfile = output_part.replace("{OUTPUT}", "{{ .Response }}")
|
||||
|
||||
# Ollama EOS
|
||||
ollama_eos = get_ollama_eos_tokens(tokenizer, extra_eos_tokens)
|
||||
ollama_eos = '\n'.join(f'PARAMETER stop "{eos}"' for eos in ollama_eos)
|
||||
|
||||
# Add temperature and min_p to counteract gibberish
|
||||
ollama_eos += "\nPARAMETER temperature 1.5\nPARAMETER min_p 0.1"
|
||||
|
||||
# Ollama modelfile
|
||||
part = '"""'
|
||||
modelfile = 'FROM {__FILE_LOCATION__}\n\n'\
|
||||
'TEMPLATE ' + part + system_modelfile + input_modelfile + output_modelfile + \
|
||||
|
|
@ -2722,12 +2686,8 @@ default_system_message = \
|
|||
extra_eos_tokens = None,
|
||||
|
||||
):
|
||||
"""
|
||||
Creates an Ollama modelfile and a HF Jinja template from a custom
|
||||
template. You must provide 2x examples of an input & output.
|
||||
There is an optional system message as well.
|
||||
|
||||
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
|
||||
"""Apply a custom chat template to a dataset (builds the Ollama modelfile
|
||||
and HF Jinja template). Use {INPUT} and {OUTPUT} twice each; {SYSTEM} is optional.
|
||||
"""
|
||||
modelfile, jinja_template, input_part, output_part = construct_chat_template(
|
||||
tokenizer = tokenizer,
|
||||
|
|
@ -2884,10 +2844,7 @@ def test_chat_templates():
|
|||
|
||||
|
||||
def test_hf_gguf_equivalence(tokenizer, gguf_model = "./model-unsloth.F16.gguf"):
|
||||
"""
|
||||
Carefully checks the output of GGUF's tokenization and HF.
|
||||
Can catch all tokenization bugs.
|
||||
"""
|
||||
"""Check GGUF vs HF tokenization to catch tokenization bugs."""
|
||||
import subprocess
|
||||
import re
|
||||
messages = [
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ class RawTextDataLoader:
|
|||
self.return_tokenized = return_tokenized
|
||||
|
||||
def detect_format(self, file_path):
|
||||
"""Auto-detect file format and parse accordingly"""
|
||||
"""Auto-detect file format from extension."""
|
||||
extension = Path(file_path).suffix.lower()
|
||||
return SUPPORTED_FORMATS.get(extension, "plain_text")
|
||||
|
||||
|
|
@ -102,11 +102,9 @@ class RawTextDataLoader:
|
|||
def create_causal_dataset(self, chunks):
|
||||
"""Create dataset for causal language modeling"""
|
||||
if chunks and isinstance(chunks[0], dict):
|
||||
# Already-tokenized chunks: reshape for Dataset.from_dict
|
||||
input_ids = [chunk["input_ids"] for chunk in chunks]
|
||||
attention_mask = [chunk["attention_mask"] for chunk in chunks]
|
||||
# Labels == input_ids for causal LM
|
||||
labels = [list(ids) for ids in input_ids]
|
||||
labels = [list(ids) for ids in input_ids] # labels == input_ids for causal LM
|
||||
return Dataset.from_dict(
|
||||
{
|
||||
"input_ids": input_ids,
|
||||
|
|
@ -125,13 +123,7 @@ class RawTextDataLoader:
|
|||
stride,
|
||||
return_tokenized = True,
|
||||
):
|
||||
"""
|
||||
Intelligent chunking that:
|
||||
1. Respects sentence/paragraph boundaries
|
||||
2. Handles various text formats (.txt, .md, .json, etc.)
|
||||
3. Maintains context with stride overlap
|
||||
4. Returns tokenized chunks directly (more efficient) or text chunks
|
||||
"""
|
||||
"""Chunk text with stride overlap; return tokenized chunks or text."""
|
||||
# Tokenize the whole text once for accurate token counts
|
||||
tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False)
|
||||
tokens = tokenized["input_ids"]
|
||||
|
|
@ -141,11 +133,9 @@ class RawTextDataLoader:
|
|||
if hasattr(tokens[0], "__len__"):
|
||||
tokens = tokens[0]
|
||||
elif isinstance(tokens, int):
|
||||
# Tokenizer returned a count; build a range
|
||||
tokens = list(range(tokens))
|
||||
tokens = list(range(tokens)) # tokenizer returned a count
|
||||
|
||||
if len(tokens) <= chunk_size:
|
||||
# Fits in a single chunk
|
||||
if return_tokenized:
|
||||
eos_token_id = getattr(self.tokenizer, "eos_token_id", None)
|
||||
if eos_token_id is not None:
|
||||
|
|
@ -190,7 +180,6 @@ class RawTextDataLoader:
|
|||
|
||||
chunks.append(chunk_text)
|
||||
|
||||
# Advance with stride overlap
|
||||
if end_idx == len(tokens):
|
||||
break
|
||||
start_idx += chunk_size - stride
|
||||
|
|
@ -268,13 +257,7 @@ class TextPreprocessor:
|
|||
return text
|
||||
|
||||
def validate_dataset(self, dataset):
|
||||
"""
|
||||
Check for:
|
||||
- Minimum/maximum sequence lengths
|
||||
- Character encoding issues
|
||||
- Repeated content
|
||||
- Empty chunks
|
||||
"""
|
||||
"""Compute dataset stats: lengths, encoding issues, repeats, empties."""
|
||||
stats = {
|
||||
"total_samples": len(dataset),
|
||||
"empty_samples": 0,
|
||||
|
|
@ -295,31 +278,26 @@ class TextPreprocessor:
|
|||
stats["empty_samples"] += 1
|
||||
continue
|
||||
|
||||
# Check for encoding issues
|
||||
try:
|
||||
text.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
stats["encoding_issues"] += 1
|
||||
|
||||
# Calculate lengths
|
||||
length = len(text)
|
||||
text_lengths.append(length)
|
||||
stats["min_length"] = min(stats["min_length"], length)
|
||||
stats["max_length"] = max(stats["max_length"], length)
|
||||
|
||||
# Check for repeated content
|
||||
text_hash = hash(text.strip())
|
||||
if text_hash in seen_texts:
|
||||
stats["repeated_content"] += 1
|
||||
else:
|
||||
seen_texts.add(text_hash)
|
||||
|
||||
# Calculate average length
|
||||
if text_lengths:
|
||||
stats["avg_length"] = sum(text_lengths) / len(text_lengths)
|
||||
stats["min_length"] = stats["min_length"] if stats["min_length"] != float("inf") else 0
|
||||
|
||||
# Generate warnings
|
||||
if stats["empty_samples"] > 0:
|
||||
stats["warnings"].append(f"Found {stats['empty_samples']} empty samples")
|
||||
|
||||
|
|
|
|||
|
|
@ -217,12 +217,11 @@ class SyntheticDataKit:
|
|||
elif dtype_val == torch.float32:
|
||||
dtype_val = "float32"
|
||||
engine_args["dtype"] = dtype_val
|
||||
# Convert torch dtype to valid CLI string
|
||||
# torch dtype -> CLI string
|
||||
if hasattr(dtype_val, "name"):
|
||||
engine_args["dtype"] = dtype_val.name
|
||||
elif isinstance(dtype_val, str) and dtype_val.startswith("torch."):
|
||||
engine_args["dtype"] = dtype_val.split(".")[-1]
|
||||
# Only allow valid vLLM choices
|
||||
valid_dtypes = {"auto", "bfloat16", "float", "float16", "float32", "half"}
|
||||
if engine_args["dtype"] not in valid_dtypes:
|
||||
engine_args["dtype"] = "auto"
|
||||
|
|
@ -250,10 +249,8 @@ class SyntheticDataKit:
|
|||
"--" + flag,
|
||||
]
|
||||
elif which == "False":
|
||||
# Ignore flag
|
||||
pass
|
||||
elif which == "None":
|
||||
# Ignore flag
|
||||
pass
|
||||
else:
|
||||
subprocess_commands += [
|
||||
|
|
@ -285,7 +282,7 @@ class SyntheticDataKit:
|
|||
ready_regex = None,
|
||||
text = False,
|
||||
)
|
||||
# we don't print stderr to console but self.stderr_capture.tail(200) will print the last 200 lines
|
||||
# stderr not echoed; self.stderr_capture.tail(200) retrieves it
|
||||
|
||||
ready = self.stdout_capture.wait_for_ready(timeout = timeout)
|
||||
if not ready:
|
||||
|
|
@ -372,7 +369,6 @@ class SyntheticDataKit:
|
|||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Delete vLLM module as well
|
||||
if hasattr(self, "_delete_vllm"):
|
||||
self._delete_vllm(llm = None)
|
||||
|
||||
|
|
@ -386,7 +382,6 @@ class SyntheticDataKit:
|
|||
self.cleanup()
|
||||
|
||||
def chunk_data(self, filename = None):
|
||||
# Chunks data by max tokens and generation length
|
||||
assert filename is not None
|
||||
assert os.path.exists(filename)
|
||||
assert hasattr(self, "tokenizer")
|
||||
|
|
@ -405,7 +400,6 @@ class SyntheticDataKit:
|
|||
raise RuntimeError("Generation length is way too long!")
|
||||
input_ids = self.tokenizer(text, add_special_tokens = False).input_ids
|
||||
|
||||
# Get left and right boundaries
|
||||
length = len(input_ids)
|
||||
n_chunks = int(np.ceil(length / (max_tokens - self.overlap)))
|
||||
boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks)).astype(int)
|
||||
|
|
|
|||
|
|
@ -64,7 +64,6 @@ def get_device_type():
|
|||
return "cuda"
|
||||
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
||||
return "xpu"
|
||||
# Check torch.accelerator
|
||||
if hasattr(torch, "accelerator"):
|
||||
if not torch.accelerator.is_available():
|
||||
raise NotImplementedError("Unsloth cannot find any torch accelerator? You need a GPU.")
|
||||
|
|
|
|||
|
|
@ -104,12 +104,11 @@ except Exception:
|
|||
|
||||
@contextlib.contextmanager
|
||||
def suppress_cuda_printf():
|
||||
"""Suppress CUDA device-side printf by redirecting stdout/stderr fds to /dev/null.
|
||||
"""Suppress CUDA device-side printf by redirecting fds 1/2 to /dev/null.
|
||||
|
||||
CUDA device printf (e.g. CUTLASS "Arch conditional MMA" errors on Blackwell)
|
||||
writes to fd 1 at the C level, bypassing Python's sys.stdout, so the
|
||||
HidePrintMessage filter can't catch it. Redirect fd 1 and 2 at the OS level,
|
||||
sync CUDA, then restore.
|
||||
CUDA device printf (e.g. CUTLASS "Arch conditional MMA" on Blackwell) writes
|
||||
to fd 1 at the C level, bypassing sys.stdout, so HidePrintMessage can't catch
|
||||
it. Redirect fds 1/2 at the OS level, sync CUDA, then restore.
|
||||
"""
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
|
@ -598,7 +597,7 @@ def patch_ipykernel_hf_xet():
|
|||
|
||||
|
||||
def patch_trackio():
|
||||
# Set some environment variables to customize the Trackio dashboard for experiment tracking
|
||||
# Customize the Trackio dashboard via environment variables
|
||||
# See https://github.com/unslothai/notebooks/pull/110
|
||||
os.environ["TRACKIO_LOGO_LIGHT_URL"] = (
|
||||
"https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png"
|
||||
|
|
@ -643,8 +642,8 @@ def check_fbgemm_gpu_version():
|
|||
|
||||
|
||||
def patch_enable_input_require_grads():
|
||||
"""Patch PreTrainedModel.enable_input_require_grads to tolerate vision models
|
||||
that raise NotImplementedError from get_input_embeddings()."""
|
||||
"""Patch enable_input_require_grads to tolerate vision models that raise
|
||||
NotImplementedError from get_input_embeddings()."""
|
||||
import inspect
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
|
|
@ -699,11 +698,10 @@ def patch_enable_input_require_grads():
|
|||
|
||||
def patch_unsafe_trainer_rng_load():
|
||||
"""Harden Trainer._load_rng_state against CVE-2026-1839 (RCE from a malicious
|
||||
rng_state.pth on resume). Hardens only the rng torch.load, via a thread-local
|
||||
flag, so it forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD)
|
||||
and refuses torch < 2.6 (CVE-2025-32434), while rng-less resumes and unrelated
|
||||
torch.load calls are untouched. No-op if transformers is absent or already
|
||||
guards the load (>= 5.0.0rc3)."""
|
||||
rng_state.pth on resume). Via a thread-local flag, hardens only the rng
|
||||
torch.load: forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD)
|
||||
and refuses torch < 2.6 (CVE-2025-32434); other torch.load calls untouched.
|
||||
No-op if transformers is absent or already guards the load (>= 5.0.0rc3)."""
|
||||
if importlib.util.find_spec("transformers") is None:
|
||||
return
|
||||
try:
|
||||
|
|
@ -767,11 +765,10 @@ def patch_unsafe_trainer_rng_load():
|
|||
|
||||
|
||||
def _is_custom_torch_build(raw_version_str):
|
||||
"""Check if a raw version string indicates a custom or source build.
|
||||
"""True if a raw version string indicates a custom/source build.
|
||||
|
||||
Operates on the raw importlib_version() string (our Version() strips local
|
||||
identifiers). Standard releases use +cu124/+rocm6.3/+cpu/+xpu; custom builds
|
||||
use +gitXXXX or other suffixes.
|
||||
Operates on the raw importlib_version() string. Standard releases use
|
||||
+cu124/+rocm6.3/+cpu/+xpu; custom builds use +gitXXXX or other suffixes.
|
||||
"""
|
||||
if "+" not in raw_version_str:
|
||||
return False
|
||||
|
|
@ -785,13 +782,11 @@ def _is_custom_torch_build(raw_version_str):
|
|||
|
||||
|
||||
def _infer_required_torchvision(torch_major, torch_minor):
|
||||
"""Infer the minimum required torchvision minor version from torch version.
|
||||
"""Min required torchvision (tv_major, tv_minor) from torch version, or None.
|
||||
|
||||
The torch -> torchvision minor version mapping follows a consistent formula:
|
||||
torch 1.x -> torchvision 0.(x + 1) (verified: torch 1.7 through 1.13)
|
||||
torch 2.x -> torchvision 0.(x + 15) (verified: torch 2.0 through 2.9)
|
||||
|
||||
Returns (tv_major, tv_minor) or None if the major version is unrecognized.
|
||||
Mapping formula:
|
||||
torch 1.x -> torchvision 0.(x + 1) (verified: torch 1.7 - 1.13)
|
||||
torch 2.x -> torchvision 0.(x + 15) (verified: torch 2.0 - 2.9)
|
||||
"""
|
||||
if torch_major == 1 and torch_minor >= 7:
|
||||
return (0, torch_minor + 1)
|
||||
|
|
@ -1010,19 +1005,13 @@ def fix_huggingface_hub():
|
|||
|
||||
|
||||
def fix_triton_compiled_kernel_missing_attrs():
|
||||
"""
|
||||
Triton 3.6.0+ removed direct `num_ctas` and `cluster_dims` attributes from
|
||||
CompiledKernel, but torch 2.9.x Inductor still expects them in
|
||||
torch/_inductor/runtime/triton_heuristics.py make_launcher() (line ~1757).
|
||||
"""Re-add num_ctas/cluster_dims to triton CompiledKernel for torch.compile.
|
||||
|
||||
The scope dict eagerly evaluates:
|
||||
binary.metadata.num_ctas, *binary.metadata.cluster_dims
|
||||
when hasattr(binary, "metadata") is True, but metadata lacks cluster_dims.
|
||||
This crashes before reaching the new launch path that doesn't need cta_args.
|
||||
|
||||
Upstream fix: pytorch/pytorch@97bd4db added hasattr guards.
|
||||
We monkey-patch CompiledKernel.__init__ to inject the missing attributes
|
||||
so the older hasattr(binary, "num_ctas") branch succeeds instead.
|
||||
Triton 3.6.0+ dropped the direct `num_ctas`/`cluster_dims` attrs, but torch
|
||||
2.9.x Inductor's make_launcher() still eagerly reads
|
||||
binary.metadata.num_ctas/*cluster_dims (metadata lacks cluster_dims), crashing
|
||||
before the new launch path. Upstream fix pytorch/pytorch@97bd4db added hasattr
|
||||
guards; we instead patch CompiledKernel.__init__ to inject the missing attrs.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
|
@ -1058,16 +1047,11 @@ def fix_triton_compiled_kernel_missing_attrs():
|
|||
|
||||
|
||||
def patch_trunc_normal_precision_issue():
|
||||
"""
|
||||
Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
|
||||
"""Patch torch.nn.init.trunc_normal_ to run fp16/bf16 init in fp32.
|
||||
|
||||
torch.nn.init.trunc_normal_ can saturate at truncation bounds in fp16/bf16 on
|
||||
some versions/backends. This was observed in TorchTitan investigations where
|
||||
low-precision truncation produced boundary-heavy initialization behavior:
|
||||
https://github.com/pytorch/torchtitan/pull/2342
|
||||
|
||||
To avoid that failure mode, initialize into a temporary fp32 tensor, then copy
|
||||
back to the original dtype.
|
||||
trunc_normal_ can saturate at truncation bounds in fp16/bf16 on some
|
||||
versions/backends (https://github.com/pytorch/torchtitan/pull/2342). Avoid
|
||||
it by initializing into a temporary fp32 tensor, then copying back.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
|
@ -1138,16 +1122,11 @@ def patch_trunc_normal_precision_issue():
|
|||
|
||||
|
||||
def check_vllm_torch_sm100_compatibility():
|
||||
"""
|
||||
Check for incompatible vLLM + torch < 2.9.0 + SM100 (Blackwell) combination.
|
||||
"""Raise a helpful error for the vLLM + torch < 2.9.0 + SM100 combination.
|
||||
|
||||
vLLM's distributed module (device_communicators) crashes with std::bad_alloc
|
||||
when imported on SM100 GPUs (B200/B100) with torch < 2.9.0. This is due to
|
||||
C++ code in vLLM's NCCL/distributed layer being incompatible with older
|
||||
torch versions on the newer Blackwell architecture.
|
||||
|
||||
This check runs early (before vLLM import) to provide a helpful error message
|
||||
instead of a cryptic std::bad_alloc crash.
|
||||
vLLM's distributed module crashes with std::bad_alloc when imported on SM100
|
||||
GPUs (B200/B100) with torch < 2.9.0. Runs early (before vLLM import) to give a
|
||||
clear message instead of the cryptic crash.
|
||||
"""
|
||||
# vLLM installed? (without importing it)
|
||||
if importlib.util.find_spec("vllm") is None:
|
||||
|
|
@ -1202,14 +1181,11 @@ def check_vllm_torch_sm100_compatibility():
|
|||
|
||||
|
||||
def fix_vllm_pdl_blackwell():
|
||||
"""
|
||||
Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100).
|
||||
"""Fix vLLM PDL (Programmatic Dependent Launch) bug on SM100 (Blackwell).
|
||||
|
||||
The issue: vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL
|
||||
optimization on SM90+ GPUs. This fails on SM100 (B200/B100) during CUDA graph
|
||||
capture because Triton's pipeliner can't handle gdc_wait in complex kernels.
|
||||
|
||||
See: https://github.com/vllm-project/vllm/issues/30872
|
||||
vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL on SM90+, but
|
||||
it fails on SM100 (B200/B100) during CUDA graph capture (Triton's pipeliner
|
||||
can't handle gdc_wait). See https://github.com/vllm-project/vllm/issues/30872
|
||||
"""
|
||||
if importlib.util.find_spec("vllm") is None:
|
||||
return
|
||||
|
|
@ -1327,12 +1303,9 @@ def fix_vllm_pdl_blackwell():
|
|||
def patch_openspiel_env_async():
|
||||
"""Apply nest_asyncio for OpenEnv EnvClient async compatibility.
|
||||
|
||||
OpenEnv's EnvClient uses async methods (reset/step). In Jupyter notebooks
|
||||
these work via top-level await, but converted scripts need
|
||||
asyncio.get_event_loop().run_until_complete() wrappers. Applying nest_asyncio
|
||||
ensures nested event loop calls work in all contexts without replacing the
|
||||
original async methods (which would break scripts that already have their own
|
||||
sync wrappers).
|
||||
OpenEnv's EnvClient uses async reset/step. nest_asyncio makes nested event
|
||||
loop calls work in both notebooks and converted scripts without replacing the
|
||||
original async methods (which would break existing sync wrappers).
|
||||
"""
|
||||
try:
|
||||
import inspect
|
||||
|
|
@ -1365,10 +1338,9 @@ def patch_torchcodec_audio_decoder():
|
|||
def disable_torchcodec_if_broken():
|
||||
"""Make broken torchcodec behave as if uninstalled (#5446).
|
||||
|
||||
transformers and datasets both detect torchcodec via find_spec, which
|
||||
returns True even when the native libs cannot dlopen. We flip their
|
||||
flags and seat a sys.modules sentinel so downstream imports fall through
|
||||
their existing except ImportError handlers cleanly.
|
||||
transformers and datasets detect torchcodec via find_spec, which returns True
|
||||
even when the native libs can't dlopen. We flip their flags and seat a
|
||||
sys.modules sentinel so downstream imports hit their except ImportError paths.
|
||||
"""
|
||||
try:
|
||||
import importlib.util
|
||||
|
|
@ -1421,18 +1393,11 @@ def disable_torchcodec_if_broken():
|
|||
def disable_broken_wandb():
|
||||
"""Disable wandb if it's installed but cannot actually import.
|
||||
|
||||
wandb can fail to import when there's a protobuf version mismatch
|
||||
(e.g., wandb < 0.19.11 with protobuf >= 6.0). This causes cascading
|
||||
import failures through trl -> transformers/accelerate -> wandb that
|
||||
crash unsloth's import chain.
|
||||
|
||||
There are two separate is_wandb_available() functions used by trl:
|
||||
- transformers.integrations.integration_utils.is_wandb_available
|
||||
(used by most trl trainers)
|
||||
- accelerate.utils.imports.is_wandb_available
|
||||
(used by trl/trainer/callbacks.py)
|
||||
|
||||
Both must be patched to fully prevent broken wandb imports.
|
||||
wandb can fail to import on a protobuf mismatch (e.g. wandb < 0.19.11 with
|
||||
protobuf >= 6.0), cascading through trl -> transformers/accelerate -> wandb.
|
||||
trl uses two separate is_wandb_available() functions
|
||||
(transformers.integrations.integration_utils and accelerate.utils.imports);
|
||||
both must be patched.
|
||||
"""
|
||||
if importlib.util.find_spec("wandb") is None:
|
||||
return # wandb not installed, nothing to do
|
||||
|
|
@ -1545,9 +1510,9 @@ def _install_transformers_conversion_mapping_stub():
|
|||
def _install_transformers_core_model_loading_stub():
|
||||
"""Stub the 8 symbols peft 0.19.x imports from this module at top level.
|
||||
|
||||
``Concatenate`` and ``ConversionOps`` MUST be real classes (peft
|
||||
subclasses them at module top); the rest only appear in runtime
|
||||
``isinstance`` / construction calls gated behind ``is_transformers_ge_v5``."""
|
||||
``Concatenate``/``ConversionOps`` MUST be real classes (peft subclasses them
|
||||
at module top); the rest only appear in runtime calls gated behind
|
||||
``is_transformers_ge_v5``."""
|
||||
name = "transformers.core_model_loading"
|
||||
existing = sys.modules.get(name)
|
||||
if existing is not None and getattr(existing, _UNSLOTH_STUB_SENTINEL, False):
|
||||
|
|
@ -1636,15 +1601,11 @@ def _install_transformers_core_model_loading_stub():
|
|||
def fix_peft_transformers_weight_conversion_import():
|
||||
"""Make ``from peft.utils import transformers_weight_conversion`` import
|
||||
cleanly on (peft 0.19.x, transformers 4.x) by stubbing the two missing
|
||||
transformers-v5 submodules. See header block above for details.
|
||||
transformers-v5 submodules (see header block above).
|
||||
|
||||
Must run BEFORE ``patch_peft_weight_converter_compatibility`` -- that
|
||||
function's bare ``except (ImportError, AttributeError): return`` would
|
||||
otherwise silently no-op.
|
||||
|
||||
No-op if peft / transformers missing, or if the peft module already
|
||||
imports cleanly. Idempotent and strictly additive (never overwrites a
|
||||
real ``transformers.conversion_mapping`` / ``core_model_loading``).
|
||||
Must run BEFORE ``patch_peft_weight_converter_compatibility``, whose bare
|
||||
``except (ImportError, AttributeError): return`` would otherwise silently
|
||||
no-op. Idempotent and strictly additive (never overwrites real submodules).
|
||||
|
||||
Returns True if patched, False if no action needed, None if peft absent."""
|
||||
if importlib.util.find_spec("peft") is None:
|
||||
|
|
@ -2601,11 +2562,10 @@ def _disable_transformers_causal_conv1d():
|
|||
|
||||
|
||||
def disable_broken_causal_conv1d():
|
||||
"""Disable causal_conv1d dynamically when its shared library is ABI-broken.
|
||||
"""Disable causal_conv1d when its shared library is ABI-broken.
|
||||
|
||||
This mirrors Unsloth's FlashAttention fallback behavior: if importing causal_conv1d
|
||||
fails with a known binary symbol error, we disable it at startup so model imports do
|
||||
not hard-fail.
|
||||
Mirrors the FlashAttention fallback: if import fails with a known binary
|
||||
symbol error, disable it at startup so model imports don't hard-fail.
|
||||
"""
|
||||
global CAUSAL_CONV1D_BROKEN
|
||||
if CAUSAL_CONV1D_BROKEN:
|
||||
|
|
@ -2684,16 +2644,13 @@ def _detect_installed_bnb_rocm_version():
|
|||
def maybe_set_windows_rocm_bnb_version():
|
||||
"""Pin ``BNB_ROCM_VERSION`` from the installed wheel on Windows + ROCm torch.
|
||||
|
||||
AMD's Windows wheel ships one ``libbitsandbytes_rocm<NN>.dll`` whose
|
||||
suffix can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll),
|
||||
breaking the native 4-bit/8-bit paths. Pin the installed suffix before
|
||||
bitsandbytes is first imported.
|
||||
AMD's Windows wheel ships one ``libbitsandbytes_rocm<NN>.dll`` whose suffix
|
||||
can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll), breaking the
|
||||
native 4/8-bit paths; pin the installed suffix before bitsandbytes is imported.
|
||||
|
||||
No-op unless ALL of: Windows, a real HIP torch build (env hints like
|
||||
HIP_PATH do not count), a ROCm DLL installed, and no explicit user value.
|
||||
Linux is untouched. Values seeded by Studio's venv sitecustomize.py
|
||||
(marked ``UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize``) are
|
||||
redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
|
||||
No-op unless ALL of: Windows, a real HIP torch build (env hints don't count),
|
||||
a ROCm DLL installed, and no explicit user value. sitecustomize-seeded values
|
||||
are redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
|
||||
opts out and drops a seeded default. Returns the value set, else None.
|
||||
"""
|
||||
if sys.platform != "win32":
|
||||
|
|
@ -2725,13 +2682,11 @@ def maybe_set_windows_rocm_bnb_version():
|
|||
|
||||
|
||||
def patch_accelerate_recursively_apply():
|
||||
"""
|
||||
Make Accelerate's recursive utilities tolerate Unsloth's EmptyLogits
|
||||
sentinel. recursively_apply returns the sentinel unchanged instead of
|
||||
raising TypeError, and find_device skips it while still finding real
|
||||
tensors, falling back to PartialState().device only for sentinel-only
|
||||
payloads. Both wrappers are idempotent and are propagated to every
|
||||
already imported accelerate namespace.
|
||||
"""Make Accelerate's recursive utilities tolerate Unsloth's EmptyLogits
|
||||
sentinel: recursively_apply returns it unchanged (no TypeError), and
|
||||
find_device skips it while still finding real tensors, falling back to
|
||||
PartialState().device only for sentinel-only payloads. Both wrappers are
|
||||
idempotent and propagated to every already-imported accelerate namespace.
|
||||
"""
|
||||
try:
|
||||
import accelerate.utils.operations as acc_ops
|
||||
|
|
|
|||
|
|
@ -14,9 +14,7 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""
|
||||
Auto-tuning cache system for MoE kernels to ensure tuning runs only once at training start.
|
||||
"""
|
||||
"""Auto-tuning cache for MoE kernels so tuning runs only once at training start."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
|
|
@ -42,7 +40,7 @@ def _get_cache_key(
|
|||
device_capability: Tuple[int, int],
|
||||
seq_len: int = 8192, # Default sequence length for tuning
|
||||
) -> str:
|
||||
"""Generate a unique cache key based on model configuration."""
|
||||
"""Unique cache key from model configuration."""
|
||||
key_data = {
|
||||
"num_experts": num_experts,
|
||||
"hidden_dim": hidden_dim,
|
||||
|
|
@ -57,7 +55,7 @@ def _get_cache_key(
|
|||
|
||||
|
||||
def _get_cache_file_path(cache_key: str) -> str:
|
||||
"""Get the file path for the cache file."""
|
||||
"""Path to the cache file for this key."""
|
||||
cache_dir = os.path.expanduser("~/.cache/unsloth/moe_autotune")
|
||||
os.makedirs(cache_dir, exist_ok = True)
|
||||
return os.path.join(cache_dir, f"{cache_key}.json")
|
||||
|
|
@ -131,21 +129,8 @@ def get_or_autotune_moe_kernels(
|
|||
force_autotune: bool = False,
|
||||
seq_len: int = 8192,
|
||||
) -> Tuple[Any, Any, Any]:
|
||||
"""
|
||||
Get cached kernel configurations or run auto-tuning.
|
||||
|
||||
Args:
|
||||
num_experts: Number of experts in the MoE layer
|
||||
hidden_dim: Hidden dimension of the model
|
||||
intermediate_dim: Intermediate dimension for MoE MLP
|
||||
top_k: Number of experts to route to
|
||||
dtype: Data type for computation
|
||||
force_autotune: Force re-running autotuning even if cache exists
|
||||
seq_len: Sequence length to use for tuning benchmarks
|
||||
|
||||
Returns:
|
||||
Tuple of (config_fwd, config_bwd_dx, config_bwd_dw)
|
||||
"""
|
||||
"""Return cached MoE kernel configs (config_fwd, config_bwd_dx, config_bwd_dw),
|
||||
running auto-tuning if needed. force_autotune ignores existing caches."""
|
||||
device_capability = torch.cuda.get_device_capability()
|
||||
cache_key = _get_cache_key(
|
||||
num_experts,
|
||||
|
|
@ -167,7 +152,6 @@ def get_or_autotune_moe_kernels(
|
|||
logger.info(f"Using in-memory cached MoE kernel configs: {cache_key}")
|
||||
return _kernel_config_cache[cache_key]
|
||||
|
||||
# Try to load from disk
|
||||
if not force_autotune:
|
||||
cached_data = load_cached_config(cache_key)
|
||||
if cached_data is not None:
|
||||
|
|
@ -206,7 +190,6 @@ def get_or_autotune_moe_kernels(
|
|||
_kernel_config_cache[cache_key] = configs
|
||||
_autotune_completed[cache_key] = True
|
||||
|
||||
# Save to disk
|
||||
config_fwd, config_bwd_dx, config_bwd_dw = configs
|
||||
save_cached_config(
|
||||
cache_key,
|
||||
|
|
@ -242,9 +225,8 @@ def _run_moe_autotuning(
|
|||
seq_len: int,
|
||||
) -> Tuple[Any, Any, Any]:
|
||||
"""Run the actual auto-tuning for MoE kernels."""
|
||||
|
||||
device = "cuda"
|
||||
# Fixed token count avoids OOMs and seq_len dependency; we ignore the passed seq_len here
|
||||
# Fixed token count avoids OOMs and seq_len dependency; passed seq_len is ignored
|
||||
num_tokens = 4096
|
||||
total_tokens = num_tokens * top_k
|
||||
|
||||
|
|
@ -260,7 +242,6 @@ def _run_moe_autotuning(
|
|||
# Dummy routing data
|
||||
m_sizes = torch.randint(1, total_tokens // num_experts + 1, (num_experts,), device = device)
|
||||
m_sizes = m_sizes * (total_tokens // m_sizes.sum().item())
|
||||
# Adjust to exact total
|
||||
diff = total_tokens - m_sizes.sum().item()
|
||||
if diff != 0:
|
||||
m_sizes[0] += diff
|
||||
|
|
@ -268,7 +249,7 @@ def _run_moe_autotuning(
|
|||
gather_indices = torch.arange(total_tokens, device = device)
|
||||
torch.randperm(total_tokens, out = gather_indices)
|
||||
|
||||
# Autotune via the interface function with autotune=True (lets triton tune)
|
||||
# autotune=True lets triton tune the kernels
|
||||
from .grouped_gemm.interface import (
|
||||
grouped_gemm_forward,
|
||||
grouped_gemm_dX,
|
||||
|
|
@ -309,7 +290,6 @@ def _run_moe_autotuning(
|
|||
use_tma_store = triton_config_fwd.kwargs.get("USE_TMA_STORE", False),
|
||||
)
|
||||
|
||||
# Autotune backward dX kernel
|
||||
logger.info("Autotuning backward dX kernel...")
|
||||
dummy_grad = torch.randn(total_tokens, 2 * intermediate_dim, device = device, dtype = dtype)
|
||||
_ = grouped_gemm_dX(
|
||||
|
|
@ -335,7 +315,6 @@ def _run_moe_autotuning(
|
|||
use_tma_store = triton_config_bwd_dx.kwargs.get("USE_TMA_STORE", False),
|
||||
)
|
||||
|
||||
# Autotune backward dW kernel
|
||||
logger.info("Autotuning backward dW kernel...")
|
||||
_ = grouped_gemm_dW(
|
||||
X = hidden_states,
|
||||
|
|
@ -366,10 +345,7 @@ def _run_moe_autotuning(
|
|||
|
||||
|
||||
def _get_heuristic_configs() -> Tuple[Any, Any, Any]:
|
||||
"""
|
||||
Get 'Safe Heuristic' kernel configurations.
|
||||
These are verified to be safe on A100 (SM80) and provide ~9x speedup on H100/B200.
|
||||
"""
|
||||
"""'Safe Heuristic' kernel configs: safe on A100 (SM80), ~9x speedup on H100/B200."""
|
||||
from .grouped_gemm.kernels.tuning import (
|
||||
KernelConfigForward,
|
||||
KernelConfigBackward_dX,
|
||||
|
|
@ -386,7 +362,7 @@ def _get_heuristic_configs() -> Tuple[Any, Any, Any]:
|
|||
permute_x = True,
|
||||
permute_y = True,
|
||||
use_tma_load_x = False,
|
||||
use_tma_load_w = False, # TMA loads might need alignment checks, safer to disable for heuristic
|
||||
use_tma_load_w = False, # TMA loads may need alignment checks; disabled for heuristic
|
||||
use_tma_store = False,
|
||||
)
|
||||
|
||||
|
|
@ -421,7 +397,7 @@ def _get_heuristic_configs() -> Tuple[Any, Any, Any]:
|
|||
|
||||
|
||||
def _get_default_configs() -> Tuple[Any, Any, Any]:
|
||||
"""Get default kernel configurations as fallback."""
|
||||
"""Default fallback kernel configurations."""
|
||||
from .grouped_gemm.kernels.tuning import (
|
||||
KernelConfigForward,
|
||||
KernelConfigBackward_dX,
|
||||
|
|
|
|||
|
|
@ -55,7 +55,6 @@ def run_benchmark_forward(
|
|||
|
||||
X = torch.randn(bs, seqlen, hidden_size, dtype = dtype, device = device, requires_grad = True)
|
||||
|
||||
# Forward
|
||||
bench_forward_ref = lambda: ref_model(X) # noqa: E731
|
||||
bench_forward_fused = lambda: tt_model(X) # noqa: E731
|
||||
|
||||
|
|
@ -104,7 +103,6 @@ def run_benchmark_backward(
|
|||
]
|
||||
test_output, _ = tt_model(X_test)
|
||||
|
||||
# Bench
|
||||
grad_output = torch.randn_like(output)
|
||||
bench_backward_ref = lambda: output.backward(grad_output, retain_graph = True) # noqa: E731
|
||||
bench_backward_fused = lambda: test_output.backward(grad_output, retain_graph = True) # noqa: E731
|
||||
|
|
@ -136,7 +134,7 @@ def setup_model(
|
|||
if isinstance(config, Qwen3MoeConfig):
|
||||
ref_model = Qwen3MoeSparseMoeBlock(config).to(device, dtype)
|
||||
|
||||
# Triton kernel grouped gemm version of MoE Block -- this is what we're testing
|
||||
# Triton grouped-gemm MoE block under test.
|
||||
tt_model = Qwen3MoeFusedGroupedGEMMBlock.from_hf(
|
||||
ref_model,
|
||||
permute_x = permute_x,
|
||||
|
|
@ -275,13 +273,13 @@ if __name__ == "__main__":
|
|||
)
|
||||
parser.add_argument(
|
||||
"--use_tma_load_w", action = "store_true"
|
||||
) # Auto-parametrized per kernel config; no need to specify
|
||||
) # Auto-parametrized per kernel config
|
||||
parser.add_argument(
|
||||
"--use_tma_load_x", action = "store_true"
|
||||
) # Auto-parametrized per kernel config; no need to specify
|
||||
) # Auto-parametrized per kernel config
|
||||
parser.add_argument(
|
||||
"--use_tma_load_dy", action = "store_true"
|
||||
) # Auto-parametrized per kernel config; no need to specify
|
||||
) # Auto-parametrized per kernel config
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type = str,
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ def create_merged_results(
|
|||
test_config_cols = list(test_config_dict.keys())
|
||||
for col in test_config_cols:
|
||||
df[col] = test_config_dict[col]
|
||||
# Reorder columns so that test config cols are first
|
||||
# Put test config cols first
|
||||
df = df[test_config_cols + kernel_result_cols]
|
||||
return df
|
||||
|
||||
|
|
|
|||
|
|
@ -33,17 +33,14 @@ ch.setFormatter(formatter)
|
|||
logger.addHandler(ch)
|
||||
|
||||
|
||||
# Precompute TMA support to avoid graph breaks
|
||||
# TMA requires both:
|
||||
# 1. NVIDIA GPU with capability >= 9 (Hopper+)
|
||||
# 2. Triton version with TMA API (make_tensor_descriptor or _experimental_make_tensor_descriptor)
|
||||
# Precomputed to avoid graph breaks. TMA needs GPU capability >= 9 (Hopper+) and a Triton TMA API.
|
||||
def _check_tma_support():
|
||||
if DEVICE_TYPE in ("xpu", "hip"):
|
||||
return False
|
||||
import triton.language as tl
|
||||
|
||||
gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9
|
||||
# Support both old experimental and new stable API names
|
||||
# Old experimental and new stable API names
|
||||
triton_has_tma_api = hasattr(tl, "make_tensor_descriptor") or hasattr(
|
||||
tl, "_experimental_make_tensor_descriptor"
|
||||
)
|
||||
|
|
@ -52,7 +49,7 @@ def _check_tma_support():
|
|||
|
||||
_SUPPORTS_TMA = _check_tma_support()
|
||||
|
||||
# Check if triton.set_allocator is available (Triton 3.0+)
|
||||
# triton.set_allocator exists on Triton 3.0+
|
||||
_HAS_SET_ALLOCATOR = hasattr(triton, "set_allocator")
|
||||
|
||||
|
||||
|
|
@ -68,10 +65,9 @@ except ImportError:
|
|||
|
||||
def _is_tracing(*tensors):
|
||||
"""
|
||||
True if tensors are fake tensors used during torch.compile tracing (Triton can't run).
|
||||
|
||||
NOTE: We do NOT use torch.compiler.is_compiling() because it returns True during both
|
||||
tracing AND execution; we only want to skip kernels during tracing on fake tensors.
|
||||
True if tensors are fake tensors from torch.compile tracing (Triton can't run).
|
||||
We avoid torch.compiler.is_compiling() since it's True during tracing AND
|
||||
execution; we only want to skip kernels during tracing on fake tensors.
|
||||
"""
|
||||
for t in tensors:
|
||||
name = type(t).__name__
|
||||
|
|
@ -118,13 +114,12 @@ def grouped_gemm_forward(
|
|||
m_sizes: torch.Tensor,
|
||||
gather_indices: torch.Tensor = None,
|
||||
topk_weights: torch.Tensor = None,
|
||||
# Fusions
|
||||
permute_x: bool = False,
|
||||
permute_y: bool = False,
|
||||
fuse_mul_post: bool = False,
|
||||
# Autotuning -- overrides manual kernel params when True
|
||||
# overrides manual kernel params when True
|
||||
autotune: bool = False,
|
||||
# Kernel tuning params (must be tuned, else poor performance)
|
||||
# must be tuned, else poor performance
|
||||
BLOCK_SIZE_M: int = 32,
|
||||
BLOCK_SIZE_N: int = 32,
|
||||
BLOCK_SIZE_K: int = 32,
|
||||
|
|
@ -135,32 +130,28 @@ def grouped_gemm_forward(
|
|||
use_tma_store: bool = False,
|
||||
# software pipelining; no effect until loop is re-written
|
||||
flatten: bool = True,
|
||||
# debugging
|
||||
debug: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Grouped GEMM forward pass for MoE MLPs.
|
||||
Grouped GEMM forward for MoE MLPs. Returns y: (total_tokens, N).
|
||||
|
||||
The implementation offers a number of fusions specific to MoE:
|
||||
- `permute_x`: fuse the permutation of hidden states from token order (original order) to grouped expert order, typically only needed for the first grouped GEMM in an MoE MLP.
|
||||
- When `permute_x` is True, `X` is expected to be of shape (num_tokens, K).
|
||||
- When `permute_x` is False, `X` is expected to be of shape (total_tokens, K) where `total_tokens = num_tokens * topk` AND already permuted to grouped expert order, i.e., hidden states are sorted such that tokens assigned to each expert are contiguous.
|
||||
- `permute_y`: fused the permutation of the output from expert grouped order back to original token order, typically only needed for the second grouped GEMM in an MoE MLP.
|
||||
- `fuse_mul_pre`: fuse the multiplication of the routed input with topk_weights, only done in the first grouped GEMM in an MoE MLP as for Llama4. Do not use, since results in performance regression as it interrupts the GEMM mainloop.
|
||||
- `fuse_mul_post`: fuse the multiplication of the routed output with topk_weights, used only when `permute_y` is True. NOTE: this should only be used when using this kernel for inference, not for training.
|
||||
MoE-specific fusions:
|
||||
- permute_x: fuse the token->grouped-expert-order permute of X (first GEMM).
|
||||
True: X is (num_tokens, K). False: X is (total_tokens, K), total_tokens =
|
||||
num_tokens * topk, already sorted so each expert's tokens are contiguous.
|
||||
- permute_y: fuse the grouped-expert-order->token-order permute of the output
|
||||
(second GEMM).
|
||||
- fuse_mul_post: fuse multiply of routed output by topk_weights; only with
|
||||
permute_y, and inference only (not training).
|
||||
|
||||
X: (M, K) hidden states where M is the num_tokens if `permute_x` is True, otherwise `total_tokens` where `total_tokens = num_tokens * topk`.
|
||||
W: (E, N, K) expert weights, where E is number of experts, N in the intermediate (output) dim, and K is the reduction dim
|
||||
m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
|
||||
gather_indices: (total_tokens,) indices of tokens assigned to each expert. E.g., slicing gather_indices by cumsum of m_sizes gives the indices of tokens assigned to each expert.
|
||||
topk_weights: (total_tokens,) weights to multiply routed output by in expert MLP calculation, used only when `fuse_mul_post` is True (see note on `fuse_mul_post`).
|
||||
use_fast_accum: currently unused; trade off faster accumulation dtype in GEMM for less precision.
|
||||
use_tma_load_x: use TMA for loading activations, incompatible with permute_x. TODO: add TMA gather / scatter support for Blackwell+.
|
||||
use_tma_load_w: use TMA for loading weights. If TMA supported, this should always be enabled as it is faster than global memory load.
|
||||
use_tma_store: use TMA for storing output, incompatible with permute_y. TODO: add TMA scatter support for Blackwell+.
|
||||
|
||||
Returns:
|
||||
y: (total_tokens, N) output of grouped GEMM
|
||||
X: (M, K) hidden states; M = num_tokens if permute_x else total_tokens.
|
||||
W: (E, N, K) expert weights (E experts, N output dim, K reduction dim).
|
||||
m_sizes: tokens per expert = the M of each per-expert GEMM.
|
||||
gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
|
||||
topk_weights: (total_tokens,) routed-output weights, used only if fuse_mul_post.
|
||||
use_tma_load_x: TMA load of activations, incompatible with permute_x.
|
||||
use_tma_load_w: TMA load of weights; prefer when TMA is supported (faster).
|
||||
use_tma_store: TMA store of output, incompatible with permute_y.
|
||||
"""
|
||||
|
||||
assert X.device.type == "cuda", "X and W must be on CUDA"
|
||||
|
|
@ -170,12 +161,11 @@ def grouped_gemm_forward(
|
|||
W = W.contiguous()
|
||||
m_sizes = m_sizes.contiguous()
|
||||
|
||||
# Preconditions
|
||||
assert not (permute_x and permute_y), "Cannot permute both X and Y"
|
||||
assert not (permute_y and use_tma_store), "Cannot use both TMA store and permute_y"
|
||||
|
||||
if use_tma_load_x:
|
||||
# TMA load for activations, TMA gather only supported on Blackwell+
|
||||
# TMA gather only supported on Blackwell+
|
||||
assert not permute_x, "Cannot use both use_tma_load_x and permute_x"
|
||||
|
||||
use_tma = use_tma_load_w or use_tma_load_x or use_tma_store
|
||||
|
|
@ -264,27 +254,21 @@ def grouped_gemm_forward(
|
|||
print(f"DEBUG::GROUPED_GEMM {m_sizes.tolist()} {(gather_indices // topk).tolist()}")
|
||||
|
||||
kernel_args = {
|
||||
# Inputs
|
||||
"x_ptr": X,
|
||||
"w_ptr": W,
|
||||
"m_sizes_ptr": m_sizes,
|
||||
"gather_indices_ptr": gather_indices,
|
||||
"topk_weights_ptr": topk_weights,
|
||||
# Output
|
||||
"y_ptr": y,
|
||||
# Problem shapes
|
||||
"NUM_TOKENS": num_tokens,
|
||||
"NUM_EXPERTS": num_experts,
|
||||
"TOPK": topk,
|
||||
"N": N,
|
||||
"K": K,
|
||||
"NUM_SMS": NUM_SMS,
|
||||
# Gather / Scatter
|
||||
"PERMUTE_X": permute_x,
|
||||
"PERMUTE_Y": permute_y,
|
||||
# TopK weight merging
|
||||
"FUSE_MUL_POST": fuse_mul_post,
|
||||
# Loop pipelining
|
||||
"FLATTEN": flatten,
|
||||
}
|
||||
if not autotune:
|
||||
|
|
@ -338,24 +322,21 @@ def grouped_gemm_dX(
|
|||
autotune: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
dX backward kernel
|
||||
grad_output: (M, N)
|
||||
gather_indices: (total_tokens,), indices of tokens assigned to each expert. E.g., slicing gather_indices by cumsum of m_sizes gives the indices of tokens assigned to each expert.
|
||||
m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
|
||||
topk: number of experts chosen per token.
|
||||
`permute_x`: whether X was permuted on load in the forward pass, typically only used for the first grouped GEMM in an MoE MLP to group tokens by expert.
|
||||
- In the forward pass, if we permuted X on load, we need to permute store in the backward pass
|
||||
- Shapes
|
||||
- the forward pass input X shape is [NUM_TOKENS, K], reduce across K, output y is [NUM_TOKENS * TOPK, K]
|
||||
- the backward pass input dy shape is [NUM_TOKENS * TOPK, N], reduce across N, output dX is [NUM_TOKENS * TOPK, K]
|
||||
- Note that in the backward pass, the output size is still [NUM_TOKENS * TOPK, K] since we still need to accumulate gradients for each expert chosen by the token in a post-processing step.
|
||||
`permute_y`: whether the output was permuted on store in the forward pass, typically only used for the second grouped GEMM in an MoE MLP to restore to the original token order.
|
||||
- In the forward pass, if we permuted output on store (e.g., in the second grouped GEMM in fused MoE MLP), we need to permute on load to get from token order to expert grouped order
|
||||
- We still store in contiguous order since we are writing out dX which will be the input to the backwards pass of the first grouped GEMM
|
||||
`fuse_mul_{pre,post}`: always set to False since this should only be used for inference.
|
||||
use_tma_load_dy: use TMA for loading dy. use_tma_load_dy is incompatible with permute_y. TODO: add TMA gather / scatter support for Blackwell+ which will enable permute_y and use_tma_load_dy.
|
||||
use_tma_load_w: use TMA for loading weights. If TMA supported, this should always be enabled as it is faster than global memory load.
|
||||
use_tma_store: use TMA for storing dX. Incompatible with permute_x. TODO: add TMA gather / scatter support for Blackwell+ which will enable permute_x and use_tma_store.
|
||||
dX backward kernel. Shapes: dy is (NUM_TOKENS*TOPK, N) reduced over N,
|
||||
output dX is (NUM_TOKENS*TOPK, K) (per-expert grads are accumulated in a
|
||||
post-processing step).
|
||||
|
||||
gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
|
||||
m_sizes: tokens per expert = the M of each per-expert GEMM.
|
||||
topk: experts chosen per token.
|
||||
permute_x: whether X was permuted on load in the forward (first GEMM); if so we
|
||||
permute on store here.
|
||||
permute_y: whether output was permuted on store in the forward (second GEMM); if
|
||||
so we permute on load here. dX is always stored contiguous.
|
||||
fuse_mul_{pre,post}: must be False (inference only).
|
||||
use_tma_load_dy: TMA load of dy, incompatible with permute_y.
|
||||
use_tma_load_w: TMA load of weights; prefer when TMA is supported (faster).
|
||||
use_tma_store: TMA store of dX, incompatible with permute_x.
|
||||
"""
|
||||
assert not fuse_mul_pre, "fuse_mul_pre should only be used for inference, not for training"
|
||||
assert not fuse_mul_post, "fuse_mul_post should only be used for inference, not for training"
|
||||
|
|
@ -364,10 +345,8 @@ def grouped_gemm_dX(
|
|||
assert m_sizes.is_contiguous()
|
||||
assert m_sizes.ndim == 1
|
||||
|
||||
# Preconditions
|
||||
assert not (permute_x and permute_y), "Cannot permute both X and Y"
|
||||
# Note that this is flipped from the forward pass
|
||||
# If we permuted y in the forward, we need to permute on load in the backward
|
||||
# Flipped from the forward: permuting y on store means permuting on load here
|
||||
assert not (permute_y and use_tma_load_dy), "Cannot use both TMA load and permute_y"
|
||||
assert not (permute_x and use_tma_store), "Cannot use both TMA store and permute_x"
|
||||
|
||||
|
|
@ -409,8 +388,7 @@ def grouped_gemm_dX(
|
|||
total_tokens = gather_indices.shape[0]
|
||||
assert total_tokens == M_total, f"Total tokens ({total_tokens}) must match M_total ({M_total})"
|
||||
|
||||
# Note that the output shape is [NUM_TOKENS * TOPK, K] even when `permute_x` is True since we need to accumulate gradients across all experts chosen by the token.
|
||||
# This will be done in a post-processing step reduction step.
|
||||
# Output stays [NUM_TOKENS * TOPK, K] even when permute_x: per-expert grads are reduced in a later step.
|
||||
output_shape = (total_tokens, K)
|
||||
dX = torch.zeros(output_shape, device = dY.device, dtype = dY.dtype)
|
||||
|
||||
|
|
@ -431,21 +409,17 @@ def grouped_gemm_dX(
|
|||
print(f"DEBUG::GROUPED_GEMM {m_sizes.tolist()}")
|
||||
|
||||
kernel_args = {
|
||||
# Inputs
|
||||
"dY_ptr": dY,
|
||||
"w_ptr": W,
|
||||
"gather_indices_ptr": gather_indices,
|
||||
"m_sizes_ptr": m_sizes,
|
||||
# Output
|
||||
"dX_ptr": dX,
|
||||
# Problem sizes
|
||||
"NUM_EXPERTS": num_experts,
|
||||
"NUM_TOKENS": num_tokens,
|
||||
"TOPK": topk,
|
||||
"N": N,
|
||||
"K": K,
|
||||
"NUM_SMS": NUM_SMS,
|
||||
# Gather / Scatter
|
||||
"PERMUTE_X": permute_x,
|
||||
"PERMUTE_Y": permute_y,
|
||||
"FLATTEN": flatten,
|
||||
|
|
@ -500,22 +474,20 @@ def grouped_gemm_dW(
|
|||
debug: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
X: (M, K) hidden states where M is the num_tokens if `permute_x` is True, otherwise `total_tokens` where `total_tokens = num_tokens * topk`.
|
||||
dY: (M, N)
|
||||
topk: number of experts to choose per token.
|
||||
m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
|
||||
gather_indices: (total_tokens,) indices of tokens assigned to each expert. E.g., slicing gather_indices by cumsum of m_sizes gives the indices of tokens assigned to each expert.
|
||||
permute_x: whether X was permuted on load in the forward pass, typically only used for the first grouped GEMM in an MoE MLP to group tokens by expert.
|
||||
- for the first grouped GEMM, we permuted on load -> X was [num_tokens, K] and stored y in expert grouped order [num_tokens * topk, K]
|
||||
- in the backwards pass, we need to permute on load of X while loading dy in contiguous (expert grouped) order
|
||||
- since we are writing out dW, there is no need to permute on store
|
||||
permute_y: whether the output was permuted on store in the forward pass, typically only used for the second grouped GEMM in an MoE MLP to restore to the original token order.
|
||||
- for the second grouped GEMM, we permuted on store -> y was permuted from expert grouped order to token order while X was loaded in expert grouped order since it was the output of the first grouped GEMM
|
||||
- in the backwards pass, we need to permute on load of dy to get from token order to expert grouped order to match the order of X
|
||||
- since we are writing out dW, there is no need to permute on store
|
||||
use_tma_load_dy: use TMA for loading dy. use_tma_load_dy is incompatible with permute_y. TODO: add TMA gather / scatter support for Blackwell+ which will enable permute_y and use_tma_load_dy.
|
||||
use_tma_load_x: use TMA for loading x. use_tma_load_x is incompatible with permute_x. TODO: add TMA gather / scatter support for Blackwell+ which will enable permute_x and use_tma_load_x.
|
||||
use_tma_store: use TMA for storing dW. If TMA supported, this should always be enabled as it is faster than global memory store.
|
||||
dW backward kernel.
|
||||
|
||||
X: (M, K) hidden states; M = num_tokens if permute_x else total_tokens.
|
||||
dY: (M, N).
|
||||
topk: experts chosen per token.
|
||||
m_sizes: tokens per expert = the M of each per-expert GEMM.
|
||||
gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
|
||||
permute_x: whether X was permuted on load in the forward (first GEMM); if so we
|
||||
permute X on load here. dW never needs permute on store.
|
||||
permute_y: whether output was permuted on store in the forward (second GEMM); if
|
||||
so we permute dy on load here to match X's order.
|
||||
use_tma_load_dy: TMA load of dy, incompatible with permute_y.
|
||||
use_tma_load_x: TMA load of x, incompatible with permute_x.
|
||||
use_tma_store: TMA store of dW; prefer when TMA is supported (faster).
|
||||
"""
|
||||
assert not fuse_mul_pre, "fuse_mul_pre not supported"
|
||||
assert not fuse_mul_post, "fuse_mul_post not supported"
|
||||
|
|
@ -524,7 +496,6 @@ def grouped_gemm_dW(
|
|||
dY = dY.contiguous()
|
||||
m_sizes = m_sizes.contiguous()
|
||||
|
||||
# Preconditions
|
||||
assert not (permute_x and permute_y), "Cannot permute both X and Y"
|
||||
assert not (permute_y and use_tma_load_dy), "Cannot use both TMA load and permute_y"
|
||||
assert not (permute_x and use_tma_load_x), "Cannot use both TMA load and permute_x"
|
||||
|
|
@ -561,7 +532,6 @@ def grouped_gemm_dW(
|
|||
num_tokens = total_tokens // topk
|
||||
|
||||
num_experts = m_sizes.shape[0]
|
||||
# Get dimensions
|
||||
_, K = X.shape
|
||||
M_grad, N = dY.shape
|
||||
|
||||
|
|
@ -598,24 +568,19 @@ def grouped_gemm_dW(
|
|||
m_start += m_sizes[i]
|
||||
|
||||
kernel_args = {
|
||||
# Inputs
|
||||
"x_ptr": X,
|
||||
"dY_ptr": dY,
|
||||
"m_sizes_ptr": m_sizes,
|
||||
"gather_indices_ptr": gather_indices,
|
||||
# Output
|
||||
"dW_ptr": dW,
|
||||
# Problem sizes
|
||||
"NUM_TOKENS": num_tokens,
|
||||
"TOPK": topk,
|
||||
"NUM_EXPERTS": num_experts,
|
||||
"N": N,
|
||||
"K": K,
|
||||
"NUM_SMS": NUM_SMS,
|
||||
# Gather / Scatter
|
||||
"PERMUTE_X": permute_x,
|
||||
"PERMUTE_Y": permute_y,
|
||||
# Loop pipelining
|
||||
"FLATTEN": flatten,
|
||||
}
|
||||
|
||||
|
|
@ -678,7 +643,7 @@ class GroupedGemm(torch.autograd.Function):
|
|||
ctx.dX_only = dX_only
|
||||
ctx.dW_only = dW_only
|
||||
|
||||
# NOTE: we don't save topk_weights for backward since we do not support training with fused_mul
|
||||
# topk_weights not saved: training with fused_mul is unsupported
|
||||
ctx.save_for_backward(X, W, m_sizes, gather_indices)
|
||||
|
||||
fwd_config = {}
|
||||
|
|
@ -702,9 +667,8 @@ class GroupedGemm(torch.autograd.Function):
|
|||
permute_x = permute_x,
|
||||
permute_y = permute_y,
|
||||
fuse_mul_post = fuse_mul_post,
|
||||
# Autotune -- this will override the manual kernel config if true
|
||||
# overrides the manual kernel config when True
|
||||
autotune = autotune,
|
||||
# Manual kernel config
|
||||
**fwd_config,
|
||||
)
|
||||
|
||||
|
|
@ -755,9 +719,8 @@ class GroupedGemm(torch.autograd.Function):
|
|||
topk = topk,
|
||||
permute_x = permute_x,
|
||||
permute_y = permute_y,
|
||||
# Autotune -- this will override the manual kernel config if true
|
||||
# overrides the manual kernel config when True
|
||||
autotune = autotune,
|
||||
# Manual kernel config
|
||||
**bwd_dW_config,
|
||||
)
|
||||
else:
|
||||
|
|
@ -783,9 +746,8 @@ class GroupedGemm(torch.autograd.Function):
|
|||
topk = topk,
|
||||
permute_x = permute_x,
|
||||
permute_y = permute_y,
|
||||
# Autotune -- this will override the manual kernel config if true
|
||||
# overrides the manual kernel config when True
|
||||
autotune = autotune,
|
||||
# Manual kernel config
|
||||
**bwd_dX_config,
|
||||
)
|
||||
|
||||
|
|
@ -865,7 +827,7 @@ def check_valid_config_bwd_dX(
|
|||
fuse_mul_post,
|
||||
is_first_gemm,
|
||||
):
|
||||
"""Check if the configuration is valid for the backward pass of dW."""
|
||||
"""Check if the configuration is valid for the backward pass of dX."""
|
||||
is_second_gemm = not is_first_gemm
|
||||
if fuse_mul_post:
|
||||
assert False, "Cannot fuse_mul is not supported for backward pass"
|
||||
|
|
@ -895,27 +857,26 @@ def grouped_gemm(
|
|||
dW_only: bool = False,
|
||||
):
|
||||
"""
|
||||
Grouped GEMM for MoE MLPs.
|
||||
Grouped GEMM for MoE MLPs (autograd-aware wrapper over the fwd/bwd kernels).
|
||||
|
||||
The implementation offers a number of fusions specific to MoE:
|
||||
- `permute_x`: fuse the permutation of hidden states from token order (original order) to grouped expert order, typically only needed for the first grouped GEMM in an MoE MLP.
|
||||
- When `permute_x` is True, `X` is expected to be of shape (num_tokens, K).
|
||||
- When `permute_x` is False, `X` is expected to be of shape (total_tokens, K) where `total_tokens = num_tokens * topk` AND already permuted to grouped expert order, i.e., hidden states are sorted such that tokens assigned to each expert are contiguous.
|
||||
- `permute_y`: fused the permutation of the output from expert grouped order back to original token order, typically only needed for the second grouped GEMM in an MoE MLP.
|
||||
- `fuse_mul`: fuse the multiplication of the routed output with topk_weights, used only when `permute_y` is True. NOTE: this should only be used when using this kernel for inference, not for training.
|
||||
|
||||
X: (M, K) hidden states where M is the num_tokens if `permute_x` is True, otherwise `total_tokens` where `total_tokens = num_tokens * topk`.
|
||||
W: (E, N, K) expert weights, where E is number of experts, N in the intermediate (output) dim, and K is the reduction dim
|
||||
m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
|
||||
gather_indices: (total_tokens,) indices of tokens assigned to each expert. E.g., slicing gather_indices by cumsum of m_sizes gives the indices of tokens assigned to each expert. Needed when either `permute_x` or `permute_y` is True.
|
||||
topk_weights: (total_tokens,) weights to multiply routed output by in expert MLP calculation, used only when `fuse_mul` is True (see note on `fuse_mul`).
|
||||
kernel_config_fwd: KernelConfigForward for forward pass.
|
||||
kernel_config_bwd_dX: KernelConfigBackward_dX for backward pass of dX.
|
||||
kernel_config_bwd_dW: KernelConfigBackward_dW for backward pass of dW.
|
||||
autotune: whether to autotune the kernel, if yes, kernel_config_fwd, kernel_config_bwd_dX, and kernel_config_bwd_dW will be ignored.
|
||||
is_first_gemm: whether this is the first grouped GEMM in an MoE MLP. This is needed to check whether kernel configs are valid. `permute_x` should only be used for first gemm; `permute_y` should only be used for second gemm.
|
||||
This will impact whether TMA can be used for loading and storing.
|
||||
MoE-specific fusions:
|
||||
- permute_x: fuse the token->grouped-expert-order permute of X (first GEMM).
|
||||
True: X is (num_tokens, K). False: X is (total_tokens, K), total_tokens =
|
||||
num_tokens * topk, already sorted by expert.
|
||||
- permute_y: fuse the grouped-expert-order->token-order permute of the output
|
||||
(second GEMM).
|
||||
- fuse_mul_post: fuse multiply of routed output by topk_weights; only with
|
||||
permute_y, and inference only (not training).
|
||||
|
||||
X: (M, K) hidden states; M = num_tokens if permute_x else total_tokens.
|
||||
W: (E, N, K) expert weights (E experts, N output dim, K reduction dim).
|
||||
m_sizes: tokens per expert = the M of each per-expert GEMM.
|
||||
gather_indices: (total_tokens,) token indices per expert (needed if permute_x/y).
|
||||
topk_weights: (total_tokens,) routed-output weights, used only if fuse_mul_post.
|
||||
kernel_config_fwd/bwd_dX/bwd_dW: per-pass kernel configs.
|
||||
autotune: if True, the kernel_config_* args are ignored.
|
||||
is_first_gemm: validates configs; permute_x is first-GEMM only, permute_y
|
||||
second-GEMM only, and this gates whether TMA can be used.
|
||||
"""
|
||||
if not autotune:
|
||||
assert (
|
||||
|
|
|
|||
|
|
@ -14,9 +14,7 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""
|
||||
Autotuning utils
|
||||
"""
|
||||
"""Autotuning utils."""
|
||||
|
||||
import logging
|
||||
from itertools import product
|
||||
|
|
@ -53,7 +51,7 @@ def _triton_supports_tma():
|
|||
"""Check if current Triton version supports TMA API."""
|
||||
import triton.language as tl
|
||||
|
||||
# Check for both old experimental and new stable API names
|
||||
# Old experimental and new stable API names
|
||||
return hasattr(tl, "make_tensor_descriptor") or hasattr(
|
||||
tl, "_experimental_make_tensor_descriptor"
|
||||
)
|
||||
|
|
@ -68,14 +66,13 @@ def get_forward_configs(
|
|||
BLOCK_M = DEFAULT_M_BLOCK_SIZES,
|
||||
BLOCK_N = DEFAULT_N_BLOCK_SIZES,
|
||||
BLOCK_K = DEFAULT_K_BLOCK_SIZES,
|
||||
TMA_LOAD_X = None, # Auto-detect if not specified
|
||||
TMA_LOAD_W = None, # Auto-detect if not specified
|
||||
TMA_STORE = False, # NOTE: TMA_STORE is disabled for now
|
||||
TMA_LOAD_X = None, # Auto-detect if None
|
||||
TMA_LOAD_W = None, # Auto-detect if None
|
||||
TMA_STORE = False, # disabled for now
|
||||
num_warps = DEFAULT_NUM_WARPS,
|
||||
num_stages = DEFAULT_NUM_STAGES,
|
||||
num_ctas = DEFAULT_NUM_CTAS,
|
||||
):
|
||||
# Auto-detect TMA support
|
||||
if TMA_LOAD_X is None:
|
||||
TMA_LOAD_X = _TRITON_HAS_TMA
|
||||
if TMA_LOAD_W is None:
|
||||
|
|
@ -149,14 +146,13 @@ def get_dX_kernel_configs(
|
|||
BLOCK_M = DEFAULT_M_BLOCK_SIZES,
|
||||
BLOCK_N = DEFAULT_N_BLOCK_SIZES,
|
||||
BLOCK_K = DEFAULT_K_BLOCK_SIZES,
|
||||
TMA_LOAD_dY = None, # Auto-detect if not specified
|
||||
TMA_LOAD_W = None, # Auto-detect if not specified
|
||||
TMA_STORE = False, # NOTE: TMA_STORE is disabled for now
|
||||
TMA_LOAD_dY = None, # Auto-detect if None
|
||||
TMA_LOAD_W = None, # Auto-detect if None
|
||||
TMA_STORE = False, # disabled for now
|
||||
num_warps = DEFAULT_NUM_WARPS,
|
||||
num_stages = DEFAULT_NUM_STAGES,
|
||||
num_ctas = DEFAULT_NUM_CTAS,
|
||||
):
|
||||
# Auto-detect TMA support
|
||||
if TMA_LOAD_dY is None:
|
||||
TMA_LOAD_dY = _TRITON_HAS_TMA
|
||||
if TMA_LOAD_W is None:
|
||||
|
|
@ -232,11 +228,10 @@ def get_dW_kernel_configs(
|
|||
num_warps = DEFAULT_NUM_WARPS,
|
||||
num_stages = DEFAULT_NUM_STAGES,
|
||||
num_ctas = DEFAULT_NUM_CTAS,
|
||||
TMA_LOAD_dY = None, # Auto-detect if not specified
|
||||
TMA_LOAD_X = None, # Auto-detect if not specified
|
||||
TMA_LOAD_dY = None, # Auto-detect if None
|
||||
TMA_LOAD_X = None, # Auto-detect if None
|
||||
TMA_STORE = False,
|
||||
):
|
||||
# Auto-detect TMA support
|
||||
if TMA_LOAD_dY is None:
|
||||
TMA_LOAD_dY = _TRITON_HAS_TMA
|
||||
if TMA_LOAD_X is None:
|
||||
|
|
@ -376,13 +371,12 @@ def prune_kernel_configs_fwd(configs: list[triton.Config], args, **kwargs):
|
|||
|
||||
pruned_configs = []
|
||||
for config in configs:
|
||||
# disable TMA if gpu does not support it
|
||||
maybe_disable_tma(config)
|
||||
|
||||
if common_prune_criteria(config, kwargs, dtype):
|
||||
continue
|
||||
if config.kwargs["USE_TMA_LOAD_X"] and kwargs["PERMUTE_X"]:
|
||||
# Dynamically disable TMA_LOAD_X for permuted X
|
||||
# TMA load incompatible with permuted X
|
||||
config.kwargs["USE_TMA_LOAD_X"] = False
|
||||
if config.kwargs["USE_TMA_STORE"] and kwargs["PERMUTE_Y"]:
|
||||
continue
|
||||
|
|
@ -403,7 +397,7 @@ def prune_dX_configs(configs: List[triton.Config], args, **kwargs):
|
|||
if common_prune_criteria(config, kwargs, dtype):
|
||||
continue
|
||||
if config.kwargs["USE_TMA_LOAD_dY"] and kwargs["PERMUTE_Y"]:
|
||||
# dynamically disable TMA_LOAD_dY for permuted Y
|
||||
# TMA load incompatible with permuted Y
|
||||
config.kwargs["USE_TMA_LOAD_dY"] = False
|
||||
if config.kwargs["USE_TMA_STORE"] and kwargs["PERMUTE_X"]:
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -77,10 +77,9 @@ def _grouped_gemm_dX_kernel(
|
|||
tl.static_assert(N % BLOCK_SIZE_N == 0, "N must be divisible by BLOCK_SIZE_N")
|
||||
tl.static_assert(K % BLOCK_SIZE_K == 0, "K must be divisible by BLOCK_SIZE_K")
|
||||
|
||||
# Create TMA descriptors for loading sorted tokens
|
||||
# When using TMA load, we don't permute_x, so shape should be [TOTAL_TOKENS, K]
|
||||
# Also, we are defining a single global descriptor with single block shape
|
||||
# Need to check that this does not result in errors when crossing expert boundaries
|
||||
# TMA descriptors for loading sorted tokens. With TMA load we don't permute_x, so shape is
|
||||
# [TOTAL_TOKENS, K]. Single global descriptor with one block shape -- verify this doesn't error
|
||||
# when crossing expert boundaries.
|
||||
if USE_TMA_LOAD_dY:
|
||||
dY_desc = tl.make_tensor_descriptor(
|
||||
dY_ptr,
|
||||
|
|
@ -110,10 +109,9 @@ def _grouped_gemm_dX_kernel(
|
|||
m_end = m_start + m_size
|
||||
|
||||
if m_size > 0:
|
||||
# Advance n offset to the weights for that respective expert
|
||||
# Advance n offset to this expert's weights
|
||||
n_start = expert_idx * N
|
||||
# N_start_offset = g.to(tl.int64) * N
|
||||
# tiles for this group's GEMM
|
||||
num_m_tiles = tl.cdiv(m_size, BLOCK_SIZE_M)
|
||||
num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
|
||||
num_tiles_per_expert = num_m_tiles * num_k_tiles
|
||||
|
|
@ -133,7 +131,6 @@ def _grouped_gemm_dX_kernel(
|
|||
while tidx >= processed_tiles and tidx < (processed_tiles + num_tiles_per_expert):
|
||||
group_index = tidx - processed_tiles
|
||||
|
||||
# Output tile for this thread block for this expert group
|
||||
tile_m_idx = group_index % num_m_tiles
|
||||
tile_k_idx = group_index // num_m_tiles
|
||||
|
||||
|
|
@ -246,10 +243,8 @@ def _grouped_gemm_dX_kernel(
|
|||
mask = store_mask,
|
||||
)
|
||||
|
||||
# Move to the next tile within this expert group
|
||||
tidx += NUM_SMS
|
||||
|
||||
# Update the total tiles count for the next expert group
|
||||
processed_tiles += num_tiles_per_expert
|
||||
|
||||
|
||||
|
|
@ -348,15 +343,12 @@ def _grouped_gemm_dW_kernel(
|
|||
)
|
||||
|
||||
for tile_idx in range(tidx, output_tiles_per_expert, NUM_SMS): # , flatten=FLATTEN):
|
||||
# Output tile index
|
||||
tile_n_idx = tile_idx % num_n_tiles
|
||||
tile_k_idx = tile_idx // num_n_tiles
|
||||
|
||||
# Output tile offsets
|
||||
n_offset = tile_n_idx * BLOCK_SIZE_N
|
||||
k_offset = tile_k_idx * BLOCK_SIZE_K
|
||||
|
||||
# For storing
|
||||
# TODO: Check whether the k mask is needed since we statically check that K is divisible by BLOCK_SIZE_K in the forward kernel
|
||||
# ditto for n_mask
|
||||
n_mask = block_range_n + n_offset < N
|
||||
|
|
@ -397,7 +389,6 @@ def _grouped_gemm_dW_kernel(
|
|||
m_block_size = tl.minimum(BLOCK_SIZE_M, m_size - tile_m_idx)
|
||||
|
||||
if m_block_size > 0:
|
||||
# Global offset for this chunk
|
||||
m_global_offset = m_start + tile_m_idx
|
||||
m_offsets = m_global_offset + block_range_m
|
||||
|
||||
|
|
|
|||
|
|
@ -11,15 +11,10 @@ from .autotuning import (
|
|||
)
|
||||
|
||||
|
||||
#
|
||||
# PERMUTE_X -> permute tokens so that they are ordered by expert
|
||||
# PERMUTE_Y -> permute output so that they are ordered by token
|
||||
# These are effectively the same thing: the former loads in permuted order, the latter stores in permuted order => we only need to define the permutation indices once
|
||||
# In the former, we use these row indices when loading X
|
||||
# For the latter, we use these row indices when storing Y
|
||||
# FUSE_MUL -> multiply routed outputs by their respective weights
|
||||
# topk_weights are in token order
|
||||
# Only account for the case when X is in expert order and we are permuting Y when fusing mul -- this precondition is checked in the interface
|
||||
# PERMUTE_X -> permute X to expert order on load; PERMUTE_Y -> permute Y to token
|
||||
# order on store. Same permutation indices either way (load vs store).
|
||||
# FUSE_MUL -> multiply routed outputs by topk_weights (token order).
|
||||
# Fusing mul assumes X in expert order while permuting Y -- checked in the interface.
|
||||
@triton.jit
|
||||
def _grouped_gemm_forward_kernel(
|
||||
x_ptr,
|
||||
|
|
@ -61,10 +56,8 @@ def _grouped_gemm_forward_kernel(
|
|||
tidx = tl.program_id(0)
|
||||
output_dtype: tl.dtype = y_ptr.dtype.element_ty
|
||||
|
||||
# Create TMA descriptors for loading sorted tokens
|
||||
# When using TMA load, we don't permute_x, so shape should be [TOTAL_TOKENS, K]
|
||||
# Also, we are defining a single global descriptor with single block shape
|
||||
# Need to check that this does not result in errors when crossing expert boundaries
|
||||
# TMA load implies no permute_x, so descriptor shape is [TOTAL_TOKENS, K].
|
||||
# Single global descriptor; may need checking across expert boundaries.
|
||||
if USE_TMA_LOAD_X:
|
||||
x_desc = tl.make_tensor_descriptor(
|
||||
x_ptr,
|
||||
|
|
@ -98,7 +91,7 @@ def _grouped_gemm_forward_kernel(
|
|||
num_n_tiles = tl.cdiv(N, BLOCK_SIZE_N)
|
||||
num_tiles_per_expert = num_m_tiles * num_n_tiles
|
||||
|
||||
# Need to create tma_store within loop since we need to predicate stores based on m_size
|
||||
# tma_store must be created in-loop to predicate stores on m_size.
|
||||
if USE_TMA_STORE:
|
||||
y_desc = tl.make_tensor_descriptor(
|
||||
y_ptr, # + m_start * N,
|
||||
|
|
@ -107,16 +100,14 @@ def _grouped_gemm_forward_kernel(
|
|||
block_shape = [BLOCK_SIZE_M, BLOCK_SIZE_N],
|
||||
)
|
||||
|
||||
# Process tiles for this expert
|
||||
while tidx >= processed_tiles and tidx < processed_tiles + num_tiles_per_expert:
|
||||
tile_idx = tidx - processed_tiles
|
||||
|
||||
# Check if L2 cache re-use for this order is optimal
|
||||
# [TODO] Check if this tile order gives optimal L2 reuse.
|
||||
tile_m_idx = tile_idx % num_m_tiles
|
||||
tile_n_idx = tile_idx // num_m_tiles
|
||||
|
||||
if SHOULD_PERMUTE_OR_FUSE:
|
||||
# These will be used for loading and storing in permuted order
|
||||
gather_offsets = tile_m_idx * BLOCK_SIZE_M + m_block_range
|
||||
indices_to_gather = m_start + tl.max_contiguous(
|
||||
tl.multiple_of(gather_offsets % m_size, BLOCK_SIZE_M),
|
||||
|
|
@ -128,26 +119,23 @@ def _grouped_gemm_forward_kernel(
|
|||
)
|
||||
expert_token_offsets = expert_token_idx[:, None]
|
||||
|
||||
# Masks for permuted load and store
|
||||
|
||||
row_mask = gather_offsets < m_size
|
||||
row_mask = row_mask[:, None]
|
||||
|
||||
# row_mask = indices_to_gather < m_end
|
||||
# row_mask = row_mask[:, None]
|
||||
|
||||
# We only take into account the following two cases: (PERMUTE_X and NOT PERMUTE_Y) and (NOT PERMUTE_X and PERMUTE_Y)
|
||||
# Hence, we can make the following simplifying assumptions when loading and storing
|
||||
# Note the different strides between the two cases: the offsets for loading and storing are flipped and the strides must also be adjusted
|
||||
# Only two cases supported: (PERMUTE_X, not PERMUTE_Y) and (not PERMUTE_X, PERMUTE_Y).
|
||||
# Between them the load/store offsets and strides are flipped.
|
||||
if PERMUTE_X:
|
||||
load_idx = (
|
||||
(expert_token_offsets // TOPK) * K
|
||||
) # Permute on load from token -> expert order, divide by TOPK to index from original number of tokens
|
||||
store_idx = indices_to_gather[:, None] * N # Store in contiguous order
|
||||
expert_token_offsets // TOPK
|
||||
) * K # token -> expert order; //TOPK indexes the original tokens
|
||||
store_idx = indices_to_gather[:, None] * N # contiguous store
|
||||
else:
|
||||
off_am = tile_m_idx * BLOCK_SIZE_M
|
||||
if not PERMUTE_Y:
|
||||
# These will already be computed if permuting y
|
||||
# Already computed above when permuting y.
|
||||
offs_am = off_am + m_block_range
|
||||
row_mask = offs_am[:, None] < m_size
|
||||
row_idx = m_start + offs_am[:, None]
|
||||
|
|
@ -166,10 +154,8 @@ def _grouped_gemm_forward_kernel(
|
|||
expert_token_offsets * N
|
||||
) # Permute on store from expert -> token order
|
||||
|
||||
# We always load topk weights in expert order
|
||||
# In the pre-multiplication case, we multiply permuted hidden states by weights before the first gemm
|
||||
# In the post-multiplication case, we multiply permuted hidden states by weights after the second gemm
|
||||
# In either case, the hidden states are grouped by expert, so we always permute on load of topk weights
|
||||
# Hidden states are grouped by expert, so topk weights are always loaded in expert order
|
||||
# (pre-mul: before first gemm; post-mul: after second gemm).
|
||||
if SHOULD_FUSE_MUL:
|
||||
topk_load_idx = expert_token_offsets
|
||||
|
||||
|
|
@ -194,7 +180,6 @@ def _grouped_gemm_forward_kernel(
|
|||
x = x_desc.load([m_start + off_am, k_offset])
|
||||
|
||||
if FUSE_MUL_PRE:
|
||||
# Check for correct broadcasting
|
||||
topk_weights = tl.load(topk_weights_ptr + topk_load_idx, mask = row_mask)
|
||||
x *= topk_weights.to(x.dtype)
|
||||
|
||||
|
|
@ -218,7 +203,6 @@ def _grouped_gemm_forward_kernel(
|
|||
# NOTE: order of fusing multiplication is important
|
||||
# Fusing before accumulator dtype conversion results in numerical diffs
|
||||
if FUSE_MUL_POST:
|
||||
# Check for correct broadcasting
|
||||
topk_weights = tl.load(topk_weights_ptr + topk_load_idx, mask = row_mask)
|
||||
y *= topk_weights.to(output_dtype)
|
||||
|
||||
|
|
|
|||
|
|
@ -272,5 +272,5 @@ class TritonTuningContext:
|
|||
f"Error running Triton grouped GEMM for kernel config: {self.kernel_config}: {exc_value}"
|
||||
)
|
||||
self.success = False
|
||||
# Return False to propagate exceptions, True to suppress them
|
||||
# True suppresses the exception, False propagates it
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -189,17 +189,14 @@ class Llama4GroupedGemmTextMoe(Llama4TextMoe):
|
|||
hidden_states = hidden_states.sum(dim = 1)
|
||||
hidden_states_after_weight_merge = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
# Token counts per expert + gather indices (token->expert order).
|
||||
# Auxiliary structs; not recorded in the autograd graph.
|
||||
# Auxiliary structs (token->expert order); not in the autograd graph.
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
||||
# Permute tokens into expert order
|
||||
hidden_states = permute(hidden_states_after_weight_merge, gather_indices, self.top_k)
|
||||
assert hidden_states.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Start expert computation
|
||||
first_gemm = torch_grouped_gemm(
|
||||
X = hidden_states, W = self.experts.gate_up_proj, m_sizes = token_counts_by_expert
|
||||
)
|
||||
|
|
@ -208,13 +205,11 @@ class Llama4GroupedGemmTextMoe(Llama4TextMoe):
|
|||
intermediate = self.act_and_mul(first_gemm)
|
||||
assert intermediate.shape == (total_tokens, self.experts.expert_dim)
|
||||
|
||||
# See comment above
|
||||
second_gemm = torch_grouped_gemm(
|
||||
X = intermediate, W = self.experts.down_proj, m_sizes = token_counts_by_expert
|
||||
)
|
||||
assert second_gemm.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Post-processing
|
||||
hidden_states_unpermute = unpermute(second_gemm, gather_indices)
|
||||
assert hidden_states_unpermute.shape == (total_tokens, hidden_dim)
|
||||
# grouped_gemm_out = hidden_states.view(batch_size, sequence_length, hidden_dim)
|
||||
|
|
@ -365,17 +360,14 @@ class Llama4TritonTextMoe(Llama4GroupedGemmTextMoe):
|
|||
hidden_states = hidden_states.sum(dim = 1)
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
# Token counts per expert + gather indices (token->expert order).
|
||||
# Auxiliary structs; not recorded in the autograd graph.
|
||||
# Auxiliary structs (token->expert order); not in the autograd graph.
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
||||
# Permute tokens into expert order
|
||||
hidden_states = permute(hidden_states, gather_indices, self.top_k)
|
||||
assert hidden_states.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Start expert computation
|
||||
hidden_states = grouped_gemm(
|
||||
X = hidden_states,
|
||||
W = self.experts.gate_up_proj,
|
||||
|
|
@ -410,7 +402,6 @@ class Llama4TritonTextMoe(Llama4GroupedGemmTextMoe):
|
|||
dX_only = self.dX_only,
|
||||
)
|
||||
|
||||
# Unpermute from expert order back to token order
|
||||
if not self.permute_y:
|
||||
hidden_states = unpermute(hidden_states, gather_indices)
|
||||
hidden_states += shared_expert_out
|
||||
|
|
|
|||
|
|
@ -70,10 +70,8 @@ class Qwen3MoeGroupedGEMMBlock(torch.nn.Module):
|
|||
config.moe_intermediate_size,
|
||||
)
|
||||
|
||||
# gating
|
||||
self.gate = torch.nn.Parameter(gate)
|
||||
|
||||
# experts
|
||||
self.gate_up_proj = torch.nn.Parameter(gate_up_proj, requires_grad = True)
|
||||
self.down_proj = torch.nn.Parameter(down_proj, requires_grad = True)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
|
@ -132,7 +130,6 @@ class Qwen3MoeGroupedGEMMBlock(torch.nn.Module):
|
|||
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim = -1)
|
||||
if self.norm_topk_prob: # only diff with mixtral sparse moe block!
|
||||
routing_weights /= routing_weights.sum(dim = -1, keepdim = True)
|
||||
# we cast back to the input dtype
|
||||
routing_weights = routing_weights.to(hidden_states.dtype)
|
||||
|
||||
return router_logits, routing_weights, selected_experts
|
||||
|
|
@ -157,8 +154,7 @@ class Qwen3MoeGroupedGEMMBlock(torch.nn.Module):
|
|||
|
||||
router_logits, routing_weights, selected_experts = self.run_router(hidden_states)
|
||||
|
||||
# Token counts per expert + gather indices (token->expert order).
|
||||
# Auxiliary structs; not recorded in the autograd graph.
|
||||
# Token counts + gather indices (token->expert order); aux structs, not in the autograd graph.
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
|
@ -167,7 +163,6 @@ class Qwen3MoeGroupedGEMMBlock(torch.nn.Module):
|
|||
hidden_states = permute(hidden_states, gather_indices, self.top_k)
|
||||
assert hidden_states.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Start expert computation
|
||||
first_gemm = torch_grouped_gemm(
|
||||
X = hidden_states, W = self.gate_up_proj, m_sizes = token_counts_by_expert
|
||||
)
|
||||
|
|
@ -274,8 +269,7 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
router_logits, routing_weights, selected_experts = self.run_router(hidden_states)
|
||||
# Token counts per expert + gather indices (token->expert order).
|
||||
# Auxiliary structs; not recorded in the autograd graph.
|
||||
# Token counts + gather indices (token->expert order); aux structs, not in the autograd graph.
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
|
@ -283,7 +277,6 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
# When permute_x is set, the permute fuses into the first gemm prologue
|
||||
if not self.permute_x:
|
||||
hidden_states = permute(hidden_states, gather_indices, self.top_k)
|
||||
# Start expert computation
|
||||
hidden_states = grouped_gemm(
|
||||
X = hidden_states,
|
||||
W = self.gate_up_proj,
|
||||
|
|
|
|||
|
|
@ -96,17 +96,15 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
router_logits, routing_weights, selected_experts = self.run_router(hidden_states)
|
||||
# Pre-processing
|
||||
# 1. Compute tokens per expert and indices for gathering tokes from token order to expert order
|
||||
# NOTE: these are auxiliary data structs which don't need to be recorded in autograd graph
|
||||
# Tokens per expert + token->expert gather indices.
|
||||
# Auxiliary structs; not recorded in the autograd graph.
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
||||
# 2. permute_x -> permutation will be fused in prologue of first grouped gemm
|
||||
# With permute_x the permutation is fused into the first grouped gemm prologue
|
||||
if not self.permute_x:
|
||||
hidden_states = permute(hidden_states, gather_indices, self.top_k)
|
||||
# Start expert computation
|
||||
hidden_states = grouped_gemm(
|
||||
X = hidden_states,
|
||||
W = self.gate_up_proj,
|
||||
|
|
@ -141,12 +139,11 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
dX_only = self.dX_only,
|
||||
)
|
||||
|
||||
# Post-processing
|
||||
# 1. Unpermute from expert order to token order
|
||||
# Unpermute from expert order back to token order
|
||||
if not self.permute_y:
|
||||
hidden_states = unpermute(hidden_states, gather_indices)
|
||||
|
||||
# 2. Merge topk weights
|
||||
# Merge topk weights
|
||||
hidden_states = (
|
||||
hidden_states.view(num_tokens, self.top_k, hidden_dim) * routing_weights[..., None]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,21 +6,13 @@ import torch.nn.functional as F
|
|||
|
||||
|
||||
def permute(X: torch.Tensor, gather_indices: torch.Tensor, topk: int):
|
||||
"""
|
||||
Scatters X to a new tensor with shape [total_tokens, hidden_dim] where total_tokens is num_tokens * topk,
|
||||
permuting the tokens according to sorted_token_idx.
|
||||
"""Reorder tokens by expert for grouped gemm.
|
||||
|
||||
Helper for grouped gemm where hidden states need be ordered by expert.
|
||||
X: [num_tokens, hidden_dim]
|
||||
sorted_token_idx: [num_tokens * topk]
|
||||
topk: int
|
||||
|
||||
Returns:
|
||||
[total_tokens, hidden_dim]
|
||||
X: [num_tokens, hidden_dim], gather_indices: [num_tokens * topk].
|
||||
Returns [total_tokens, hidden_dim] where total_tokens = num_tokens * topk.
|
||||
"""
|
||||
assert gather_indices.ndim == 1
|
||||
X = X.view(-1, X.shape[-1])
|
||||
# Shortcut for topk == 1
|
||||
if topk == 1:
|
||||
return X[gather_indices]
|
||||
|
||||
|
|
@ -42,12 +34,8 @@ def calculate_topk(
|
|||
pre_act: bool = True,
|
||||
post_act: bool = False,
|
||||
):
|
||||
"""
|
||||
If post_act is True, then activation function is run AFTER topk
|
||||
If post_act is False, then activation function is run BEFORE topk
|
||||
|
||||
This is to align with triton_bench implementation (post_act) whereas most models use pre_act (e.g. llama4, deepseek)
|
||||
"""
|
||||
"""Run activation before topk (pre_act, e.g. llama4/deepseek) or after
|
||||
(post_act, aligns with triton_bench)."""
|
||||
assert pre_act ^ post_act, "only one of pre_act or post_act can be True"
|
||||
|
||||
def _activation(gating_output: torch.Tensor):
|
||||
|
|
@ -80,21 +68,16 @@ def get_routing_indices(
|
|||
num_experts,
|
||||
return_scatter_indices: bool = False,
|
||||
):
|
||||
"""Returns token_counts_by_expert [num_experts], gather_indices [num_tokens],
|
||||
and optionally scatter_indices [bs*seqlen*top_k] to unpermute back to token order.
|
||||
"""
|
||||
Returns:
|
||||
token_counts_by_expert: [num_experts]
|
||||
gather_indices: [num_tokens]
|
||||
scatter_indices [Optional] (torch.Tensor):
|
||||
Indices for unpermuting gathered inputs back to token order, shape ``(bs * seqlen * top_k,)``.
|
||||
"""
|
||||
# group tokens together by expert indices from 0 to num_experts and pass that to experts forward
|
||||
token_counts_by_expert = torch.histc(
|
||||
selected_experts.view(-1),
|
||||
bins = num_experts,
|
||||
min = 0,
|
||||
max = num_experts,
|
||||
)
|
||||
# token_indices_experts_sorted shape (bs*slen*top_k,)
|
||||
# Sort tokens by expert so each expert gets a contiguous slice. Stable keeps token order within an expert.
|
||||
gather_indices = torch.argsort(selected_experts.view(-1), stable = True)
|
||||
if return_scatter_indices:
|
||||
scatter_indices = gather_indices.argsort()
|
||||
|
|
@ -109,14 +92,8 @@ def torch_grouped_gemm(
|
|||
m_sizes,
|
||||
transpose = True,
|
||||
):
|
||||
"""
|
||||
X: [M, K] if forward, else [M, N]
|
||||
W: [E, N, K]
|
||||
m_sizes: [E]
|
||||
|
||||
Returns:
|
||||
Y: [M, N] if forward, else [M, K]
|
||||
"""
|
||||
"""X: [M, K] (fwd) else [M, N]; W: [E, N, K]; m_sizes: [E].
|
||||
Returns Y: [M, N] (fwd) else [M, K]."""
|
||||
X = X.view(-1, X.shape[-1])
|
||||
M, K = X.shape
|
||||
|
||||
|
|
@ -136,11 +113,8 @@ def torch_grouped_gemm(
|
|||
if m_size > 0:
|
||||
m_end = m_start + m_size
|
||||
|
||||
# Extract group input
|
||||
# m_size x K
|
||||
X_g = X[m_start:m_end]
|
||||
# N x K
|
||||
W_g = W[g]
|
||||
X_g = X[m_start:m_end] # [m_size, K]
|
||||
W_g = W[g] # [N, K]
|
||||
|
||||
# Y_g = X_g @ W_g.T -> [m_size, N]
|
||||
W_g = W_g.T if transpose else W_g
|
||||
|
|
|
|||
|
|
@ -120,19 +120,17 @@ def assert_close(
|
|||
Compare reference values against obtained values.
|
||||
"""
|
||||
|
||||
# cast to float32:
|
||||
ref = ref.to(torch.float32).detach()
|
||||
tri = tri.to(torch.float32).detach()
|
||||
assert ref.shape == tri.shape, f"Tensors must have same size {ref.shape = } {tri.shape = }"
|
||||
|
||||
# deal with infinite elements:
|
||||
inf_mask_ref = torch.isinf(ref)
|
||||
inf_mask_tri = torch.isinf(tri)
|
||||
assert torch.equal(inf_mask_ref, inf_mask_tri), "Tensor must have same infinite elements"
|
||||
refn = torch.where(inf_mask_ref, 0, ref)
|
||||
trin = torch.where(inf_mask_tri, 0, tri)
|
||||
|
||||
# normalise so that RMS calculation doesn't overflow:
|
||||
# normalise so RMS calculation doesn't overflow
|
||||
eps = 1.0e-30
|
||||
multiplier = 1.0 / (torch.max(torch.abs(refn)) + eps)
|
||||
refn *= multiplier
|
||||
|
|
@ -243,7 +241,6 @@ def remove_feature_flags(
|
|||
):
|
||||
pruned_configs = []
|
||||
for config in kernel_configs:
|
||||
# Remove permute flags first:
|
||||
if permute_x and config.permute_x:
|
||||
continue
|
||||
if permute_y and config.permute_y:
|
||||
|
|
|
|||
|
|
@ -34,7 +34,6 @@ def rebind_experts_to_shared_buffer(moe_block: Qwen3MoeSparseMoeBlock, config: Q
|
|||
buffer_gate = torch.empty(num_experts, interm_size, hidden_size, device = device, dtype = dtype)
|
||||
buffer_down = torch.empty(num_experts, hidden_size, interm_size, device = device, dtype = dtype)
|
||||
|
||||
# Copy existing expert weights into buffers
|
||||
for i, expert in enumerate(moe_block.experts):
|
||||
buffer_up[i].copy_(expert.up_proj.weight.data)
|
||||
buffer_gate[i].copy_(expert.gate_proj.weight.data)
|
||||
|
|
@ -75,7 +74,7 @@ class ForwardResult:
|
|||
output: torch.Tensor
|
||||
router_logits: torch.Tensor
|
||||
X: torch.Tensor
|
||||
# When using grouped gemm MoE implementation to additional debugging / checking of intermediate results
|
||||
# Intermediate results from the grouped gemm MoE path, for debugging / checks
|
||||
grouped_gemm_result: GroupedGEMMResult = None
|
||||
|
||||
|
||||
|
|
@ -117,13 +116,11 @@ def check_gate_up_proj_grad(
|
|||
assert ref_gate_proj_grad is not None
|
||||
assert ref_up_proj_grad is not None
|
||||
|
||||
# Extract gradients
|
||||
test_gate_proj_grad = grouped_gemm_block.gate_up_proj.grad[i, :moe_intermediate_size]
|
||||
test_up_proj_grad = grouped_gemm_block.gate_up_proj.grad[i, moe_intermediate_size:]
|
||||
assert test_gate_proj_grad is not None
|
||||
assert test_up_proj_grad is not None
|
||||
|
||||
# Sanity check shapes
|
||||
assert (
|
||||
ref_gate_proj_grad.shape == test_gate_proj_grad.shape
|
||||
), f"{ref_gate_proj_grad.shape} != {test_gate_proj_grad.shape}"
|
||||
|
|
@ -131,7 +128,6 @@ def check_gate_up_proj_grad(
|
|||
ref_up_proj_grad.shape == test_up_proj_grad.shape
|
||||
), f"{ref_up_proj_grad.shape} != {test_up_proj_grad.shape}"
|
||||
|
||||
# Check gradients
|
||||
diff = (ref_gate_proj_grad - test_gate_proj_grad).abs().max()
|
||||
if not torch.allclose(ref_gate_proj_grad, test_gate_proj_grad, atol = atol, rtol = rtol):
|
||||
print(f"expert {i} gate_proj_grad_diff: {diff.detach().cpu().item():.6f}")
|
||||
|
|
@ -199,7 +195,6 @@ def check_expert_grads(
|
|||
ref_grads.shape == test_grads.shape
|
||||
), f"{field}: {ref_grads.shape} != {test_grads.shape}"
|
||||
|
||||
# Test each expert
|
||||
for i in range(ref_grads.shape[0]):
|
||||
ref_grad = ref_grads[i]
|
||||
test_grad = test_grads[i]
|
||||
|
|
@ -208,7 +203,6 @@ def check_expert_grads(
|
|||
ref_grad, test_grad, atol = atol, rtol = rtol
|
||||
), f"{field}[{i}] diff: {diff.detach().cpu().item():.6f}"
|
||||
|
||||
# Test all experts
|
||||
diff = (ref_grads - test_grads).abs().max()
|
||||
if verbose:
|
||||
print(f"{field} diff: {diff.detach().cpu().item():.6f}")
|
||||
|
|
@ -238,7 +232,6 @@ def check_fwd(
|
|||
rtol: float,
|
||||
verbose: bool = False,
|
||||
):
|
||||
# First check hidden states (output)
|
||||
ref_output = ref_result.output
|
||||
test_output = test_result.output
|
||||
diff = (ref_output - test_output).abs().max()
|
||||
|
|
@ -248,7 +241,6 @@ def check_fwd(
|
|||
ref_output, test_output, atol = atol, rtol = rtol
|
||||
), f"output diff: {diff.detach().cpu().item():.6f}"
|
||||
|
||||
# Check router logits
|
||||
ref_router_logits = ref_result.router_logits
|
||||
test_router_logits = test_result.router_logits
|
||||
diff = (ref_router_logits - test_router_logits).abs().max()
|
||||
|
|
@ -272,8 +264,8 @@ def check_grouped_gemm_results(
|
|||
test_value = getattr(fused_result, field.name)
|
||||
diff = (ref_value - test_value).abs().max()
|
||||
|
||||
# second_gemm in torch grouped gemm is not yet unpermuted so comparing the fused unpermuted second_gemm will result in error
|
||||
# instead the hidden_states_unpermute should match since hidden_states_unpermute for the fused result is the same as second_gemm
|
||||
# torch second_gemm is still permuted vs the fused one; compare via
|
||||
# hidden_states_unpermute instead (equals second_gemm for the fused result).
|
||||
if field.name == "second_gemm" and permute_y:
|
||||
continue
|
||||
|
||||
|
|
@ -332,11 +324,9 @@ def run_backward(
|
|||
|
||||
|
||||
class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
||||
"""Reference MoE block using triton grouped gemm.
|
||||
|
||||
Like Qwen3MoeGroupedGEMMBlock but with triton (not torch-native) grouped gemm.
|
||||
NOT for production: it saves intermediate results and runs extra checks for
|
||||
debugging. See grouped_gemm/reference/moe_block.py for a cleaner version.
|
||||
"""Reference MoE block like Qwen3MoeGroupedGEMMBlock but with triton (not torch-native) grouped
|
||||
gemm. NOT for production: saves intermediates and runs extra debug checks. See
|
||||
grouped_gemm/reference/moe_block.py for a cleaner version.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -404,8 +394,8 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
router_logits, routing_weights, selected_experts = self.run_router(hidden_states)
|
||||
# Pre-processing: token counts per expert + token-order -> expert-order
|
||||
# gather indices (auxiliary, not recorded in the autograd graph).
|
||||
# Token counts per expert + token-order -> expert-order gather indices
|
||||
# (auxiliary, not recorded in the autograd graph).
|
||||
token_counts_by_expert, gather_indices = self.get_token_counts_and_gather_indices(
|
||||
selected_experts
|
||||
)
|
||||
|
|
@ -415,7 +405,6 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
hidden_states = permute(hidden_states, gather_indices, self.top_k)
|
||||
assert hidden_states.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Start expert computation
|
||||
first_gemm = grouped_gemm(
|
||||
X = hidden_states,
|
||||
W = self.gate_up_proj,
|
||||
|
|
@ -449,7 +438,7 @@ class Qwen3MoeFusedGroupedGEMMBlock(Qwen3MoeGroupedGEMMBlock):
|
|||
)
|
||||
assert second_gemm.shape == (total_tokens, hidden_dim)
|
||||
|
||||
# Post-processing: unpermute expert order -> token order
|
||||
# Unpermute expert order -> token order
|
||||
if not self.permute_y:
|
||||
hidden_states_unpermute = unpermute(second_gemm, gather_indices)
|
||||
assert hidden_states_unpermute.shape == (total_tokens, hidden_dim)
|
||||
|
|
|
|||
|
|
@ -43,8 +43,7 @@ from .common import (
|
|||
SEED = 0
|
||||
|
||||
|
||||
# Only certain (permute_x, permute_y, use_W1) combinations are valid; see the
|
||||
# module string below for the full rationale.
|
||||
# Only certain (permute_x, permute_y, use_W1) combos are valid; see module string below for rationale.
|
||||
def check_valid_config(
|
||||
permute_x,
|
||||
permute_y,
|
||||
|
|
@ -106,7 +105,6 @@ def _test_grouped_gemm_forward(
|
|||
use_W1: bool, # W1 -> first grouped GEMM in a fused MoE MLP, not W1 -> second grouped GEMM in a fused MoE MLP
|
||||
fuse_mul_post: bool = False,
|
||||
flatten: bool = True,
|
||||
# Manually tuned parameters
|
||||
use_tma_load_w: bool = False,
|
||||
use_tma_load_x: bool = False,
|
||||
use_tma_store: bool = False,
|
||||
|
|
@ -115,10 +113,8 @@ def _test_grouped_gemm_forward(
|
|||
BLOCK_SIZE_K: int = None,
|
||||
num_warps: int = None,
|
||||
num_stages: int = None,
|
||||
# Autotuning parameters
|
||||
autotune: bool = False,
|
||||
num_autotune_configs: int = None,
|
||||
# Flag to manually enable TMA store
|
||||
allow_tma_store: bool = False,
|
||||
use_autograd: bool = False,
|
||||
):
|
||||
|
|
@ -189,7 +185,7 @@ def _test_grouped_gemm_forward(
|
|||
else:
|
||||
X_test = Xperm
|
||||
|
||||
# No need to run all configs for tests, otherwise takes too long
|
||||
# Limit configs so tests don't take too long
|
||||
if autotune:
|
||||
from grouped_gemm.kernels.forward import _autotuned_grouped_gemm_forward_kernel
|
||||
if num_autotune_configs is not None:
|
||||
|
|
@ -197,7 +193,6 @@ def _test_grouped_gemm_forward(
|
|||
_autotuned_grouped_gemm_forward_kernel.configs[:num_autotune_configs]
|
||||
)
|
||||
|
||||
# Use autograd.Function interface
|
||||
if use_autograd:
|
||||
from grouped_gemm.interface import grouped_gemm
|
||||
kernel_config_fwd = KernelConfigForward(
|
||||
|
|
@ -228,7 +223,6 @@ def _test_grouped_gemm_forward(
|
|||
autotune = autotune,
|
||||
is_first_gemm = use_W1,
|
||||
)
|
||||
# Use manual interface
|
||||
else:
|
||||
test_output = grouped_gemm_forward(
|
||||
X = X_test,
|
||||
|
|
@ -257,8 +251,7 @@ def _test_grouped_gemm_forward(
|
|||
if permute_y:
|
||||
ref_output = unpermute(ref_output, gather_indices)
|
||||
if fuse_mul_post:
|
||||
# if we don't permute_y, then test output is permuted with topk weights applied
|
||||
# the ref output needs to be unpermuted before multiplying by topk weights since topk weights are in token order
|
||||
# topk weights are in token order, so unpermute both before multiplying when permute_y is False
|
||||
if not permute_y:
|
||||
ref_output = unpermute(ref_output, gather_indices)
|
||||
test_output = unpermute(test_output, gather_indices)
|
||||
|
|
@ -269,7 +262,7 @@ def _test_grouped_gemm_forward(
|
|||
), f"Grouped gemm forward failed: {(ref_output - test_output).abs().max().item():.6f}"
|
||||
|
||||
|
||||
# NOTE: Fuse multiplication of topk weights is only supported for inference and not training, although this may change in the future; not currently tested.
|
||||
# Fused topk-weight mul is inference-only; not currently tested.
|
||||
@pytest.mark.parametrize(
|
||||
"kernel_config",
|
||||
KERNEL_CONFIGS_FWD,
|
||||
|
|
@ -537,7 +530,7 @@ def _test_grouped_gemm_backward_dX(
|
|||
ref_grad = Xperm.grad
|
||||
|
||||
if autotune:
|
||||
# No need to run all configs for autotuning
|
||||
# Limit configs to speed up autotuning
|
||||
from grouped_gemm.kernels.backward import _autotuned_grouped_gemm_dX_kernel
|
||||
if num_autotune_configs is not None:
|
||||
_autotuned_grouped_gemm_dX_kernel.configs = _autotuned_grouped_gemm_dX_kernel.configs[
|
||||
|
|
@ -650,8 +643,7 @@ def _test_grouped_gemm_backward_dX(
|
|||
# debug=True,
|
||||
)
|
||||
|
||||
# if permute_x and use_W1 (first grouped GEMM) then the kernel should have unpermuted the dX
|
||||
# therefore we need to unpermute the ref_grad to compare to the output of the kernel
|
||||
# For the first GEMM with permute_x the kernel unpermutes dX, so unpermute ref_grad to match
|
||||
if permute_x and use_W1:
|
||||
ref_grad = unpermute(ref_grad, gather_indices)
|
||||
|
||||
|
|
@ -665,12 +657,9 @@ def _test_grouped_gemm_backward_dX(
|
|||
), f"Grouped gemm manual backward_dX outputs mismatch: {diff:.6f}"
|
||||
|
||||
if permute_x and use_W1:
|
||||
# Show that reduction results in diffs
|
||||
# First calculate X.grad manually by backpropping through unpermuted ref_grad
|
||||
# Show that the topk reduction introduces diffs vs autograd
|
||||
dX_ref_check = ref_grad.view(num_tokens, topk, K).sum(dim = 1)
|
||||
# Do the same for the actual output of the kernel
|
||||
dX_test_check = dX_test.view(num_tokens, topk, K).sum(dim = 1)
|
||||
# Show diffs for each combination
|
||||
diff_ref_check = (X.grad - dX_ref_check).abs().max().item()
|
||||
diff_test_check = (X.grad - dX_test_check).abs().max().item()
|
||||
diff_check_test = (dX_ref_check - dX_test_check).abs().max().item()
|
||||
|
|
@ -679,8 +668,7 @@ def _test_grouped_gemm_backward_dX(
|
|||
)
|
||||
|
||||
|
||||
# NOTE: We reduce the size of the Llama4 model configs to prevent OOM
|
||||
# Important to note that for the full model size (5120, 8192), the tests do result in diffs on the order of 1e-2.
|
||||
# Llama4 configs are shrunk to avoid OOM; the full size (5120, 8192) shows diffs ~1e-2.
|
||||
@pytest.mark.parametrize(
|
||||
"kernel_config",
|
||||
KERNEL_CONFIGS_BWD_dX,
|
||||
|
|
@ -759,7 +747,6 @@ def test_grouped_gemm_backward_dX_autotune(
|
|||
use_W1: bool,
|
||||
num_autotune_configs: int,
|
||||
):
|
||||
# TMA loads / stores will be autotuned
|
||||
_test_grouped_gemm_backward_dX(
|
||||
data_config = data_config,
|
||||
model_config = model_config,
|
||||
|
|
@ -792,7 +779,6 @@ def test_grouped_gemm_backward_dX_autotune_autograd(
|
|||
use_W1: bool,
|
||||
num_autotune_configs: int,
|
||||
):
|
||||
# TMA loads / stores will be autotuned
|
||||
_test_grouped_gemm_backward_dX(
|
||||
data_config = data_config,
|
||||
model_config = model_config,
|
||||
|
|
@ -903,8 +889,7 @@ def _test_grouped_gemm_backward_dW(
|
|||
ref_output = torch_grouped_gemm(X = Xperm, W = W, m_sizes = expert_token_counts)
|
||||
assert ref_output.shape == output_shape
|
||||
|
||||
# if permute_y then the assumption is that the output of grouped_gemm was unpermuted on store
|
||||
# Therefore we have to unpermute before backpropping to ensure proper alignment
|
||||
# permute_y means grouped_gemm unpermuted on store, so unpermute before backprop to align
|
||||
if permute_y:
|
||||
ref_output = unpermute(ref_output, gather_indices)
|
||||
|
||||
|
|
@ -913,7 +898,6 @@ def _test_grouped_gemm_backward_dW(
|
|||
assert X.grad is not None
|
||||
assert W.grad is not None
|
||||
|
||||
# Test backward kernel directly
|
||||
X_ = X_test if permute_x else Xperm_test
|
||||
|
||||
if debug:
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ LLAMA4_SCOUT_ID = "meta-llama/Llama-4-Scout-17B-16E"
|
|||
SEED = 42
|
||||
SEQ_LENS = [1024]
|
||||
DTYPES = [torch.bfloat16]
|
||||
# Reduce the number of autotuning configs to prevent excessive runtime
|
||||
# Cap autotuning configs to keep runtime reasonable
|
||||
NUM_AUTOTUNE_CONFIGS = 50
|
||||
|
||||
|
||||
|
|
@ -162,7 +162,7 @@ def test_llama4_ref(
|
|||
permute_x: bool,
|
||||
permute_y: bool,
|
||||
overlap_router_shared: bool,
|
||||
model_config: Llama4TextConfig, # test fixture
|
||||
model_config: Llama4TextConfig,
|
||||
bs: int = 1,
|
||||
device = "cuda",
|
||||
precision = ".6f",
|
||||
|
|
@ -180,7 +180,6 @@ def test_llama4_ref(
|
|||
# Reference op -- HF
|
||||
llama4_ref = Llama4TextMoe(model_config).to(dtype = dtype, device = device)
|
||||
|
||||
# Torch grouped gemm impl
|
||||
llama4_gg_ref = Llama4GroupedGemmTextMoe(
|
||||
model_config, overlap_router_shared = overlap_router_shared
|
||||
).to(dtype = dtype, device = device)
|
||||
|
|
|
|||
|
|
@ -113,9 +113,7 @@ def test_qwen3_moe(
|
|||
permute_y: bool,
|
||||
autotune: bool,
|
||||
):
|
||||
torch.manual_seed(
|
||||
SEED
|
||||
) # Should not be needed when running using pytest -- autouse fixture in conftest.py
|
||||
torch.manual_seed(SEED) # Redundant under pytest -- conftest.py has an autouse fixture
|
||||
device = "cuda"
|
||||
hidden_size = config.hidden_size
|
||||
bs = 1
|
||||
|
|
@ -123,7 +121,7 @@ def test_qwen3_moe(
|
|||
# Reference op -- HF
|
||||
moe_block = Qwen3MoeSparseMoeBlock(config).to(device, dtype)
|
||||
|
||||
# Torch-native grouped gemm version of MoE Block -- for sanity checking
|
||||
# Torch-native grouped gemm version -- for sanity checking
|
||||
grouped_gemm_block = Qwen3MoeGroupedGEMMBlock.from_hf(moe_block).to(device, dtype)
|
||||
grouped_gemm_block.check_weights(moe_block)
|
||||
|
||||
|
|
@ -153,7 +151,7 @@ def test_qwen3_moe(
|
|||
kernel_config_bwd_dW = None
|
||||
kernel_config_bwd_dX = None
|
||||
|
||||
# Triton kernel grouped gemm version of MoE Block -- this is what we're testing
|
||||
# Triton kernel grouped gemm version -- this is what we're testing
|
||||
fused_gemm_block = Qwen3MoeFusedGroupedGEMMBlock.from_hf(
|
||||
moe_block,
|
||||
permute_x = permute_x,
|
||||
|
|
@ -186,7 +184,7 @@ def test_qwen3_moe(
|
|||
with annotated_context(
|
||||
"Checking torch grouped gemm MoE vs fused grouped gemm MoE forward outputs..."
|
||||
):
|
||||
# We implement a custom check for grouped gemm results to test each of the intermediate results for easier debugging
|
||||
# Custom check so each intermediate result is compared for easier debugging
|
||||
check_grouped_gemm_results(
|
||||
grouped_result.grouped_gemm_result,
|
||||
fused_result.grouped_gemm_result,
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from .sentence_transformer import FastSentenceTransformer
|
|||
try:
|
||||
from .falcon_h1 import FastFalconH1Model
|
||||
except:
|
||||
# falcon_h1 absent before transformers 4.53.0; skip
|
||||
# falcon_h1 needs transformers >= 4.53.0
|
||||
pass
|
||||
from .dpo import PatchDPOTrainer, PatchKTOTrainer
|
||||
from ._utils import is_bfloat16_supported, is_vLLM_available, __version__
|
||||
|
|
|
|||
|
|
@ -196,17 +196,9 @@ from unsloth_zoo.temporary_patches import (
|
|||
|
||||
def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_length, dtype):
|
||||
"""
|
||||
Apply gradient checkpointing with smart heuristics.
|
||||
|
||||
For seq < 512, gc="unsloth" offloading overhead isn't worth it; standard gc is faster.
|
||||
|
||||
Args:
|
||||
use_gradient_checkpointing: "unsloth", True, False, or None
|
||||
max_seq_length: The maximum sequence length
|
||||
dtype: The model dtype for patching
|
||||
|
||||
Returns:
|
||||
The effective use_gradient_checkpointing value (may change from "unsloth" to True)
|
||||
Apply gradient checkpointing with smart heuristics, returning the effective
|
||||
setting (may downgrade "unsloth" to True). For seq < 512, "unsloth" offloading
|
||||
overhead isn't worth it, so standard gc is used instead.
|
||||
"""
|
||||
if use_gradient_checkpointing == "unsloth":
|
||||
# Offloading not worth it below ~512; standard gc is faster (crossover ~384-512).
|
||||
|
|
@ -733,8 +725,8 @@ class HideLoggingMessage(logging.Filter):
|
|||
# Replace warning messages (analogous to HideLoggingMessage but for warnings.warn)
|
||||
class ReplaceWarningMessage:
|
||||
"""
|
||||
Intercepts warnings.warn calls and replaces matching messages with Unsloth branded ones.
|
||||
Uses a list of registered (match_text, replacement, category) rules checked in order.
|
||||
Intercept warnings.warn and replace matching messages with Unsloth ones, via
|
||||
registered (match_text, replacement, category) rules checked in order.
|
||||
"""
|
||||
|
||||
_rules = []
|
||||
|
|
@ -1077,9 +1069,7 @@ from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled
|
|||
|
||||
|
||||
def extract_quant_model_param_count(model):
|
||||
"""
|
||||
Calculate quant model param count based on difference in param class. Returns int for param count.
|
||||
"""
|
||||
"""Param count of a quantized model (Params4bit counted as 2x numel)."""
|
||||
count: int = 0
|
||||
for name, p in model.named_parameters():
|
||||
if p.__class__.__name__ == "Params4bit":
|
||||
|
|
@ -1090,9 +1080,7 @@ def extract_quant_model_param_count(model):
|
|||
|
||||
|
||||
def get_model_param_count(model, trainable_only = False):
|
||||
"""
|
||||
Calculate model's total param count. If trainable_only is True then count only those requiring grads
|
||||
"""
|
||||
"""Total model param count; if trainable_only, count only params requiring grads."""
|
||||
if is_deepspeed_zero3_enabled():
|
||||
|
||||
def numel(p):
|
||||
|
|
@ -1145,8 +1133,7 @@ def patch_mistral_nemo_config(config):
|
|||
|
||||
|
||||
try:
|
||||
# Some Config files use layer_type_validation
|
||||
# for eg Gemma-2, so we must import it to stop errors.
|
||||
# Needed for configs that use layer_type_validation (e.g. Gemma-2).
|
||||
from transformers.configuration_utils import layer_type_validation
|
||||
except:
|
||||
pass
|
||||
|
|
@ -1175,11 +1162,9 @@ model_architectures = [
|
|||
"falcon_h1",
|
||||
]
|
||||
|
||||
# Transformers 5.x uses class-level annotations with @strict, @auto_docstring,
|
||||
# and interval() in config classes. exec(inspect.getsource(...)) fails because
|
||||
# those symbols are not in scope. Skip the exec-based config patching for 5.x
|
||||
# since those configs already use rope_parameters (the v5 replacement for
|
||||
# rope_scaling).
|
||||
# Skip exec-based config patching on transformers 5.x: its @strict/@auto_docstring/
|
||||
# interval() config symbols aren't in scope for exec(getsource(...)), and v5 configs
|
||||
# already use rope_parameters (the rope_scaling replacement).
|
||||
_skip_config_exec_patch = Version(transformers_version) >= Version("5.0.0")
|
||||
|
||||
for model_name in model_architectures:
|
||||
|
|
@ -1187,7 +1172,7 @@ for model_name in model_architectures:
|
|||
break
|
||||
config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"
|
||||
model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
|
||||
config_filename = f"{model_name.title().replace('_','')}Config" # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
|
||||
config_filename = f"{model_name.title().replace('_','')}Config" # e.g. qwen3_moe folder but Qwen3Config class; strip underscores
|
||||
try:
|
||||
exec(f"from {config_filepath} import {config_filename}", globals())
|
||||
except:
|
||||
|
|
@ -1728,10 +1713,8 @@ import psutil
|
|||
|
||||
|
||||
def _get_statistics(statistics = None, force_download = True):
|
||||
# We log some basic stats about which environment is being used.
|
||||
# We simply download a README.md file from HF - all data is made public.
|
||||
# This is simply so we can check if some envs are broken or not.
|
||||
# You can disable this by commenting the below out
|
||||
# Log basic env stats by downloading a public README.md from HF (checks for broken/down envs).
|
||||
# Disable by commenting the below out.
|
||||
n_cpus = psutil.cpu_count(logical = False)
|
||||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
# Check modelscope for down detection
|
||||
|
|
@ -1834,11 +1817,8 @@ def _get_statistics(statistics = None, force_download = True):
|
|||
|
||||
|
||||
def get_statistics(local_files_only = False):
|
||||
# We log some basic stats about which environment is being used.
|
||||
# This is also to check if HuggingFace is down or not!
|
||||
# We simply download a README.md file from HF - all data is made public.
|
||||
# This is simply so we can check if some envs are broken or not.
|
||||
# You can disable this by setting UNSLOTH_DISABLE_STATISTICS
|
||||
# Log basic env stats by downloading a public README.md from HF (also detects if HF is down).
|
||||
# Disable via UNSLOTH_DISABLE_STATISTICS.
|
||||
import os
|
||||
|
||||
if (
|
||||
|
|
@ -2194,7 +2174,6 @@ def patch_llama_rope_scaling(
|
|||
|
||||
|
||||
def create_boolean_mask(n = 4096, sliding_window = 2048):
|
||||
# Creates a boolean mask for attention
|
||||
mask = torch.ones(n, n, dtype = torch.bool)
|
||||
if sliding_window == 0:
|
||||
return torch.triu(mask, diagonal = 1, out = mask)
|
||||
|
|
@ -2251,12 +2230,10 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):
|
|||
if "num_items_in_batch" in kwargs:
|
||||
num_items_in_batch = kwargs["num_items_in_batch"]
|
||||
if num_items_in_batch is None:
|
||||
# Remove it since the model does not support it!
|
||||
kwargs.pop("num_items_in_batch")
|
||||
elif "num_items_in_batch" not in inputs:
|
||||
inputs["num_items_in_batch"] = num_items_in_batch
|
||||
|
||||
# Get gradient accumulation steps if possible
|
||||
if (
|
||||
num_items_in_batch is None
|
||||
and getattr(getattr(self, "args", self), "gradient_accumulation_steps", 1) != 1
|
||||
|
|
@ -2865,9 +2842,7 @@ class TorchAOConfig:
|
|||
|
||||
def _untie_input_output_embeddings(model: torch.nn.Module) -> None:
|
||||
"""
|
||||
Utility to untie input/output embeddings in a HuggingFace model.
|
||||
This is useful if we want to quantize the input/ouput embeddings differently.
|
||||
Model is modified in-place.
|
||||
Untie input/output embeddings in-place (so they can be quantized differently).
|
||||
"""
|
||||
|
||||
# 1) Persist setting in config
|
||||
|
|
@ -2904,10 +2879,7 @@ def _untie_input_output_embeddings(model: torch.nn.Module) -> None:
|
|||
def _filter_fn_to_fqns(
|
||||
model: torch.nn.Module, filter_fn: Callable[[torch.nn.Module, str], bool]
|
||||
) -> Iterator[str]:
|
||||
"""
|
||||
Given a model and a filter function (m, fqn) -> bool,
|
||||
yield fully qualified names (FQNs) of modules that match.
|
||||
"""
|
||||
"""Yield FQNs of modules matching filter_fn(module, fqn) -> bool."""
|
||||
for fqn, module in model.named_modules():
|
||||
if filter_fn(module, fqn):
|
||||
yield fqn
|
||||
|
|
@ -2952,14 +2924,10 @@ def _prepare_model_for_qat(
|
|||
model: torch.nn.Module, qat_scheme: Union[str, TorchAOConfig]
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Transform a model for Quantization-Aware Training (QAT) during fine-tuning.
|
||||
|
||||
On a high level, this means fake quantizing the base (frozen) model during training.
|
||||
Fake quantization refers to simulating quantization numerics in high precision (e.g. bf16).
|
||||
This helps mitigate quantization degradations when the model is quantized after training.
|
||||
|
||||
QAT can be optionally combined with LoRA fine-tuning to for additional throughput improvement.
|
||||
For more details: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700
|
||||
Transform a model for Quantization-Aware Training (QAT) during fine-tuning, i.e.
|
||||
fake-quantize the frozen base model (simulate quant numerics in high precision)
|
||||
to reduce post-training quantization degradation. Combinable with LoRA.
|
||||
See https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700
|
||||
"""
|
||||
try:
|
||||
from torchao.quantization import PerRow, quantize_
|
||||
|
|
@ -3183,15 +3151,10 @@ def verify_fp8_support_if_applicable(model_config):
|
|||
|
||||
def _get_inference_mode_context_manager(model: torch.nn.Module):
|
||||
"""
|
||||
If the state dict was quantized using torchao, we will run into
|
||||
the following error when calling ops like aten.t() in inference mode.
|
||||
This is a bug in PyTorch that affects all tensor subclasses.
|
||||
|
||||
Cannot set version_counter for inference tensor
|
||||
|
||||
For now, we work around this issue by using `torch.no_grad()` in this case.
|
||||
See https://github.com/pytorch/pytorch/issues/164872 for more details.
|
||||
Otherwise, just return `torch.inference_mode()`.
|
||||
For torchao-quantized models, return torch.no_grad() instead of
|
||||
torch.inference_mode(), since ops like aten.t() on tensor subclasses hit a
|
||||
PyTorch bug ("Cannot set version_counter for inference tensor").
|
||||
See https://github.com/pytorch/pytorch/issues/164872
|
||||
"""
|
||||
torchao_config = getattr(model, "torchao_config", None)
|
||||
if torchao_config is not None and torchao_config.qat_scheme is None:
|
||||
|
|
@ -3223,15 +3186,7 @@ def hf_login(token: Optional[str] = None) -> Optional[str]:
|
|||
|
||||
|
||||
def is_moe_model(model) -> bool:
|
||||
"""
|
||||
Detect if a model is a Mixture of Experts (MoE) model.
|
||||
|
||||
Args:
|
||||
model: The model to check (can be HF model or config)
|
||||
|
||||
Returns:
|
||||
True if the model is an MoE model, False otherwise
|
||||
"""
|
||||
"""Detect if a model (or config) is a Mixture of Experts (MoE) model."""
|
||||
config = getattr(model, "config", model)
|
||||
|
||||
# Different MoE models use different config attribute names:
|
||||
|
|
@ -3256,11 +3211,9 @@ def is_moe_model(model) -> bool:
|
|||
|
||||
def _resolve_moe_parameter_name(model, default_name: str, alternate_name: str) -> str:
|
||||
"""
|
||||
Resolve the actual parameter path for MoE expert weights.
|
||||
|
||||
Most current Unsloth MoE models expose expert weights under
|
||||
``mlp.experts.*``. Gemma4 stores them directly under ``experts.*``.
|
||||
Prefer the path that exists on the loaded module when possible.
|
||||
Resolve the parameter path for MoE expert weights. Most models use
|
||||
``mlp.experts.*``; Gemma4 uses ``experts.*``. Prefer whichever exists on the
|
||||
loaded module.
|
||||
"""
|
||||
if hasattr(model, "named_parameters"):
|
||||
try:
|
||||
|
|
@ -3304,24 +3257,12 @@ def _moe_target_set_from_string(target_modules: str) -> set[str]:
|
|||
|
||||
def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str]]:
|
||||
"""
|
||||
Get the target_parameters for MoE expert layers if applicable.
|
||||
|
||||
For MoE models, returns the parameter paths for expert weights
|
||||
(gate_up_proj, down_proj) that should be targeted by PEFT's
|
||||
target_parameters for LoRA on nn.Parameter. The exact parameter path
|
||||
depends on the model layout, for example ``mlp.experts.*`` or
|
||||
``experts.*``.
|
||||
|
||||
Only includes MoE parameters that match what's in target_modules:
|
||||
- If "down_proj" is in target_modules -> includes "mlp.experts.down_proj"
|
||||
- If "gate_proj" or "up_proj" is in target_modules -> includes "mlp.experts.gate_up_proj"
|
||||
|
||||
Args:
|
||||
model: The model to get target parameters for
|
||||
target_modules: List/tuple of target module names to match against
|
||||
|
||||
Returns:
|
||||
List of parameter paths for MoE experts, or None if not an MoE model
|
||||
Return the MoE expert parameter paths to pass to PEFT's target_parameters for
|
||||
LoRA on nn.Parameter, or None if not an MoE model. Only paths matching
|
||||
target_modules are included:
|
||||
- "down_proj" -> "<prefix>.experts.down_proj"
|
||||
- "gate_proj"/"up_proj"/"gate_up_proj" -> "<prefix>.experts.gate_up_proj"
|
||||
The prefix depends on layout (``mlp.experts.*`` or ``experts.*``).
|
||||
"""
|
||||
if not is_moe_model(model):
|
||||
return None
|
||||
|
|
@ -3383,14 +3324,10 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str
|
|||
|
||||
|
||||
def make_fast_generate_wrapper(original_generate):
|
||||
"""
|
||||
Creates a wrapper around model.generate that checks for incorrect
|
||||
vLLM-style usage when fast_inference=False.
|
||||
"""
|
||||
"""Wrap model.generate to reject vLLM-style usage when fast_inference=False."""
|
||||
|
||||
@functools.wraps(original_generate)
|
||||
def _fast_generate_wrapper(*args, **kwargs):
|
||||
# Check for vLLM-specific arguments
|
||||
if "sampling_params" in kwargs:
|
||||
raise ValueError(
|
||||
"Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). "
|
||||
|
|
@ -3432,7 +3369,6 @@ def make_fast_generate_wrapper(original_generate):
|
|||
" )"
|
||||
)
|
||||
|
||||
# Call original generate
|
||||
return original_generate(*args, **kwargs)
|
||||
|
||||
return _fast_generate_wrapper
|
||||
|
|
|
|||
|
|
@ -89,7 +89,6 @@ def CohereAttention_fast_forward(
|
|||
*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
|
||||
|
|
@ -129,7 +128,6 @@ def CohereAttention_fast_forward(
|
|||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
rope_position_ids = position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
|
|
@ -137,7 +135,6 @@ def CohereAttention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
backend = select_attention_backend(use_varlen)
|
||||
attention_config = AttentionConfig(
|
||||
|
|
@ -193,7 +190,6 @@ def CohereDecoderLayer_fast_forward(
|
|||
device = f"{DEVICE_TYPE_TORCH}:0",
|
||||
)
|
||||
|
||||
# Self Attention
|
||||
residual = hidden_states
|
||||
hidden_states = fast_layernorm_inference(self.input_layernorm, hidden_states, out_weight)
|
||||
hidden_states_attention, self_attn_weights, present_key_value = self.self_attn(
|
||||
|
|
@ -208,7 +204,6 @@ def CohereDecoderLayer_fast_forward(
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states_mlp = fast_swiglu_inference(self.mlp, hidden_states)
|
||||
residual += hidden_states_attention
|
||||
residual += hidden_states_mlp
|
||||
|
|
@ -228,7 +223,6 @@ def CohereDecoderLayer_fast_forward(
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states_mlp = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states_attention + hidden_states_mlp
|
||||
|
||||
|
|
@ -242,7 +236,7 @@ def CohereDecoderLayer_fast_forward(
|
|||
|
||||
from math import sqrt as math_sqrt
|
||||
|
||||
KV_CACHE_INCREMENT = 256 # KV Cache update size
|
||||
KV_CACHE_INCREMENT = 256
|
||||
torch_nn_functional_softmax = torch.nn.functional.softmax
|
||||
torch_matmul = torch.matmul
|
||||
|
||||
|
|
@ -403,7 +397,6 @@ def CohereAttention_fast_forward_inference(
|
|||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
|
||||
# Attention
|
||||
if bsz == 1:
|
||||
Qn *= (
|
||||
self.scalar
|
||||
|
|
|
|||
|
|
@ -11,13 +11,12 @@
|
|||
# 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.
|
||||
"""
|
||||
FastDiffusionModel: a transformers-only slow path for text-diffusion models (e.g. DiffusionGemma).
|
||||
"""FastDiffusionModel: transformers-only slow path for text-diffusion models (e.g. DiffusionGemma).
|
||||
|
||||
These models use a block-diffusion sampling loop (custom generate) and a novel backbone, so we skip
|
||||
Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay
|
||||
bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the
|
||||
(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land.
|
||||
These models use a block-diffusion sampling loop and a novel backbone, so we skip Unsloth's
|
||||
autoregressive kernel/compile patching and load the unmodified HF model (outputs stay bit-identical to
|
||||
transformers), keeping only safe conveniences: 4bit/8bit loading, PEFT LoRA, the (model, tokenizer)
|
||||
API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
|
@ -213,7 +212,7 @@ class FastDiffusionModel:
|
|||
if not return_tokenizer:
|
||||
return model, None
|
||||
|
||||
# Prefer the processor (chat template + tokenizer); fall back to a bare tokenizer. Returned as
|
||||
# Prefer the processor (chat template + tokenizer), else a bare tokenizer; returned as
|
||||
# "tokenizer" to match the Unsloth (model, tokenizer) contract.
|
||||
try:
|
||||
tokenizer = AutoProcessor.from_pretrained(
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ def FalconH1Attention_fast_forward(
|
|||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, hidden_states.device)
|
||||
|
||||
# Falcon H1 multiplies key states by a multiplier
|
||||
# Falcon H1 scales key states by key_multiplier
|
||||
K = K * self.config.key_multiplier
|
||||
|
||||
Q = Q.transpose(1, 2)
|
||||
|
|
@ -125,7 +125,7 @@ def FalconH1Attention_fast_forward(
|
|||
cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
rope_position_ids = position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
# Useful for LongRoPE
|
||||
# Needed for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
|
|
@ -133,7 +133,6 @@ def FalconH1Attention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
window = (-1, -1)
|
||||
use_varlen = (
|
||||
attention_mask is None
|
||||
|
|
@ -190,33 +189,12 @@ def FalconH1Attention_fast_forward_inference(
|
|||
attention_mask = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
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
|
||||
"""Fast inference using the KV cache.
|
||||
|
||||
[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.
|
||||
QK^T splits into 4 chunks; the mask zeroes Qk^T and softmax is row-wise, so
|
||||
softmax(QK^T)V is just the prior step's attention. We therefore only compute
|
||||
the final row: pass one row of Q while remembering K and V (the KV cache).
|
||||
Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
|
||||
"""
|
||||
Xn = hidden_states
|
||||
bsz, _, hd = hidden_states.size()
|
||||
|
|
@ -294,8 +272,7 @@ def FalconH1Attention_fast_forward_inference(
|
|||
# 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
|
||||
# Extend 2 steps ahead to avoid errors on short KV cache
|
||||
self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
|
||||
# Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
|
||||
|
|
@ -348,7 +325,6 @@ def FalconH1Attention_fast_forward_inference(
|
|||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
|
||||
# Attention
|
||||
if bsz == 1:
|
||||
Qn *= (
|
||||
self.scalar
|
||||
|
|
@ -389,19 +365,7 @@ def FalconH1DecoderLayer_fast_forward(
|
|||
*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
|
||||
"""
|
||||
"""FalconH1 decoder layer: mamba + attention mixer, then SwiGLU MLP, with residuals."""
|
||||
if use_cache and hasattr(self, "_flag_for_generation"):
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
|
||||
|
|
@ -464,7 +428,6 @@ def FalconH1DecoderLayer_fast_forward(
|
|||
|
||||
hidden_states = mamba_hidden_states + attention_hidden_states
|
||||
|
||||
# residual connection after attention + Mamba
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
|
|
@ -535,7 +498,7 @@ def _FalconH1_fast_forward_inference(
|
|||
next_decoder_cache = []
|
||||
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual.copy_(X) # residual = X
|
||||
residual.copy_(X)
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.input_layernorm,
|
||||
X,
|
||||
|
|
@ -563,7 +526,7 @@ def _FalconH1_fast_forward_inference(
|
|||
|
||||
X += residual
|
||||
|
||||
residual.copy_(X) # residual = X
|
||||
residual.copy_(X)
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.pre_ff_layernorm,
|
||||
X,
|
||||
|
|
@ -672,7 +635,6 @@ def _fast_prepare_inputs_for_generation(
|
|||
|
||||
|
||||
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
|
||||
|
||||
|
|
|
|||
|
|
@ -100,7 +100,6 @@ def GemmaDecoderLayer_fast_forward(
|
|||
device = f"{DEVICE_TYPE_TORCH}:0",
|
||||
)
|
||||
|
||||
# Self Attention
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(
|
||||
self.input_layernorm, hidden_states, out_weight
|
||||
|
|
@ -118,7 +117,6 @@ def GemmaDecoderLayer_fast_forward(
|
|||
)
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(
|
||||
self.post_attention_layernorm, hidden_states, out_weight
|
||||
|
|
@ -141,7 +139,6 @@ def GemmaDecoderLayer_fast_forward(
|
|||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
|
|
@ -458,7 +455,6 @@ class FastGemmaModel(FastLlamaModel):
|
|||
else:
|
||||
param.requires_grad_(False)
|
||||
|
||||
# Patch RMS Layernorm
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, GemmaRMSNorm):
|
||||
# Must be in float32
|
||||
|
|
|
|||
|
|
@ -68,7 +68,6 @@ if HAS_FLASH_ATTENTION_SOFTCAPPING:
|
|||
from flash_attn import flash_attn_func
|
||||
|
||||
|
||||
# Logit softcapping
|
||||
def Gemma2Attention_fast_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
|
|
@ -82,7 +81,7 @@ def Gemma2Attention_fast_forward(
|
|||
*args,
|
||||
**kwargs,
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
# Clear inference
|
||||
# Clear cached inference buffers
|
||||
if hasattr(self, "paged_attention"):
|
||||
del self.paged_attention_K
|
||||
del self.paged_attention_V
|
||||
|
|
@ -127,7 +126,6 @@ def Gemma2Attention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Only enable if the attention_mask is True
|
||||
use_sliding_window = kwargs.get("use_sliding_window")
|
||||
has_sliding_window = (
|
||||
use_sliding_window
|
||||
|
|
@ -215,7 +213,6 @@ def Gemma2DecoderLayer_fast_forward(
|
|||
device = f"{DEVICE_TYPE_TORCH}:0",
|
||||
)
|
||||
|
||||
# Self Attention
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(
|
||||
self.input_layernorm, hidden_states, out_weight
|
||||
|
|
@ -237,7 +234,6 @@ def Gemma2DecoderLayer_fast_forward(
|
|||
)
|
||||
hidden_states += residual
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(
|
||||
self.pre_feedforward_layernorm, hidden_states, out_weight
|
||||
|
|
@ -264,7 +260,6 @@ def Gemma2DecoderLayer_fast_forward(
|
|||
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm(
|
||||
self.pre_feedforward_layernorm, hidden_states, gemma = True
|
||||
|
|
@ -403,7 +398,6 @@ def Gemma2Attention_fast_forward_inference(
|
|||
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 = self.config.sliding_window
|
||||
if use_sliding_window and kv_seq_len > sliding_window:
|
||||
start = kv_seq_len - sliding_window
|
||||
|
|
@ -420,7 +414,6 @@ def Gemma2Attention_fast_forward_inference(
|
|||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
|
||||
# Attention
|
||||
# [TODO] Gemma2 uses manual matmul for all batch sizes since SDPA lacks
|
||||
# softcapping (tanh logit scaling). If PyTorch adds a softcap param to
|
||||
# SDPA, consider SDPA for bsz > 1 to match the llama/qwen3 pattern.
|
||||
|
|
@ -500,8 +493,7 @@ def Gemma2Model_fast_forward_inference(
|
|||
GA = attention_mask
|
||||
next_decoder_cache = []
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
# For pipeline parallelism, we need to move all tensors to the same device
|
||||
# note that this movement is once per GPU in PP
|
||||
# Pipeline parallelism: move tensors to this layer's device (once per GPU)
|
||||
device_index = getattr(decoder_layer, "_per_layer_device_index", 0)
|
||||
hidden_states, position_ids = move_to_device(device_index, hidden_states, position_ids)
|
||||
|
||||
|
|
@ -609,7 +601,6 @@ class FastGemma2Model(FastLlamaModel):
|
|||
else:
|
||||
param.requires_grad_(False)
|
||||
|
||||
# Patch RMS Layernorm
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, Gemma2RMSNorm):
|
||||
# Must be in float32
|
||||
|
|
|
|||
|
|
@ -12,15 +12,10 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
GLM-4.7 Flash (GLM4 MoE Lite) optimized implementation using grouped GEMM.
|
||||
"""GLM-4.7 Flash (GLM4 MoE Lite) optimized implementation using grouped GEMM.
|
||||
|
||||
Key architecture differences from Qwen3 MoE:
|
||||
- Router uses sigmoid activation (not softmax)
|
||||
- Has routed_scaling_factor of 1.8
|
||||
- Has 1 shared expert that processes all tokens
|
||||
- Uses group-based selection before topk
|
||||
- Uses MLA (Multi-head Latent Attention)
|
||||
Differences from Qwen3 MoE: sigmoid router (not softmax), routed_scaling_factor 1.8, 1 shared expert
|
||||
processing all tokens, group-based selection before topk, and MLA (Multi-head Latent Attention).
|
||||
"""
|
||||
|
||||
from .llama import *
|
||||
|
|
@ -54,8 +49,7 @@ try:
|
|||
if _moe_path not in sys.path:
|
||||
sys.path.insert(0, _moe_path)
|
||||
|
||||
# Import first to apply the TMA compatibility shim (patches triton.language
|
||||
# for both old and new TMA API names)
|
||||
# Import first to apply the TMA compatibility shim (old + new TMA API names)
|
||||
import grouped_gemm # noqa: F401 - triggers TMA compatibility shim
|
||||
|
||||
from grouped_gemm.interface import grouped_gemm
|
||||
|
|
@ -88,7 +82,7 @@ try:
|
|||
except ImportError:
|
||||
HAS_GLM4_MOE = False
|
||||
|
||||
# Create dummy classes for type checking
|
||||
# Dummy classes for type checking
|
||||
class Glm4MoeLiteAttention:
|
||||
pass
|
||||
|
||||
|
|
@ -118,15 +112,7 @@ torch_nn_functional_silu = torch.nn.functional.silu
|
|||
|
||||
|
||||
def Glm4MoeLiteMoE_fast_forward(self, hidden_states):
|
||||
"""
|
||||
Optimized MoE forward pass using grouped GEMM.
|
||||
|
||||
GLM4 MoE specifics:
|
||||
- Uses sigmoid router activation (not softmax)
|
||||
- Has routed_scaling_factor of 1.8
|
||||
- Has 1 shared expert that always processes all tokens
|
||||
- Uses group-based selection with topk_group
|
||||
"""
|
||||
"""Optimized MoE forward pass using grouped GEMM (sigmoid router + 1 shared expert)."""
|
||||
residuals = hidden_states
|
||||
orig_shape = hidden_states.shape
|
||||
batch_size, seq_len, hidden_dim = orig_shape
|
||||
|
|
@ -185,7 +171,7 @@ def Glm4MoeLiteMoE_fast_forward(self, hidden_states):
|
|||
else:
|
||||
hidden_states = self.experts(hidden_states, topk_indices, topk_weights)
|
||||
|
||||
# Add shared expert output
|
||||
# Add shared expert (processes all tokens)
|
||||
hidden_states = hidden_states + self.shared_experts(residuals.view(-1, hidden_dim))
|
||||
|
||||
return hidden_states.view(*orig_shape)
|
||||
|
|
@ -194,17 +180,8 @@ def Glm4MoeLiteMoE_fast_forward(self, hidden_states):
|
|||
def Glm4MoeLiteNaiveMoe_fast_forward(
|
||||
self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Optimized expert forward using grouped GEMM.
|
||||
|
||||
Args:
|
||||
hidden_states: [num_tokens, hidden_dim]
|
||||
top_k_index: [num_tokens, top_k] indices of selected experts
|
||||
top_k_weights: [num_tokens, top_k] weights for selected experts
|
||||
|
||||
Returns:
|
||||
[num_tokens, hidden_dim] output after weighted sum of expert outputs
|
||||
"""
|
||||
"""Optimized expert forward using grouped GEMM. hidden_states [num_tokens, hidden_dim],
|
||||
top_k_index/top_k_weights [num_tokens, top_k] -> [num_tokens, hidden_dim]."""
|
||||
num_tokens, hidden_dim = hidden_states.shape
|
||||
top_k = top_k_index.shape[1]
|
||||
top_k_weights = top_k_weights.to(hidden_states.dtype)
|
||||
|
|
@ -244,7 +221,6 @@ def Glm4MoeLiteNaiveMoe_fast_forward(
|
|||
# Under autocast hidden_states may be fp32 while weights are bf16
|
||||
hidden_states = hidden_states.to(self.gate_up_proj.dtype)
|
||||
|
||||
# First grouped GEMM: gate_up_proj
|
||||
intermediate = grouped_gemm(
|
||||
X = hidden_states,
|
||||
W = self.gate_up_proj,
|
||||
|
|
@ -261,7 +237,6 @@ def Glm4MoeLiteNaiveMoe_fast_forward(
|
|||
gate, up = intermediate.chunk(2, dim = -1)
|
||||
intermediate = self.act_fn(gate) * up
|
||||
|
||||
# Second grouped GEMM: down_proj
|
||||
expert_output = grouped_gemm(
|
||||
X = intermediate,
|
||||
W = self.down_proj,
|
||||
|
|
@ -293,13 +268,10 @@ def Glm4MoeLiteDecoderLayer_fast_forward(
|
|||
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Optimized decoder layer forward with fast RMS layernorm.
|
||||
"""
|
||||
"""Optimized decoder layer forward with fast RMS layernorm."""
|
||||
is_inference = use_cache and hasattr(self, "_flag_for_generation")
|
||||
|
||||
if is_inference:
|
||||
# Self-attention with fast inference path
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
|
||||
hidden_states, _ = self.self_attn(
|
||||
|
|
@ -345,21 +317,13 @@ def Glm4MoeLiteDecoderLayer_fast_forward(
|
|||
|
||||
|
||||
def Glm4MoeLiteMLP_fast_forward(self, x):
|
||||
"""
|
||||
Optimized MLP forward using fused SwiGLU.
|
||||
"""
|
||||
"""Optimized MLP forward using fused SwiGLU."""
|
||||
return fast_swiglu_inference(self, x)
|
||||
|
||||
|
||||
class FastGLM47Model(FastLlamaModel):
|
||||
"""
|
||||
Fast GLM-4.7 Flash (GLM4 MoE Lite) model with grouped GEMM optimization.
|
||||
|
||||
This provides 2-3x throughput improvement for MoE layers by:
|
||||
- Replacing sequential expert loops with grouped GEMM operations
|
||||
- Fusing permutation operations into the GEMM kernels
|
||||
- Using optimized RMS LayerNorm and SwiGLU implementations
|
||||
"""
|
||||
"""Fast GLM-4.7 Flash (GLM4 MoE Lite) model with grouped GEMM optimization (2-3x MoE throughput
|
||||
via grouped GEMM, fused permutation, and optimized RMS LayerNorm / SwiGLU)."""
|
||||
|
||||
@staticmethod
|
||||
def pre_patch():
|
||||
|
|
@ -369,8 +333,7 @@ class FastGLM47Model(FastLlamaModel):
|
|||
"Please upgrade with: pip install --upgrade transformers"
|
||||
)
|
||||
|
||||
# Patch MoE forward with grouped GEMM (TMA compat handled by
|
||||
# grouped_gemm/__init__.py)
|
||||
# Patch MoE forward with grouped GEMM (TMA compat in grouped_gemm/__init__.py)
|
||||
if HAS_GROUPED_GEMM:
|
||||
Glm4MoeLiteNaiveMoe.forward = Glm4MoeLiteNaiveMoe_fast_forward
|
||||
Glm4MoeLiteMoE.forward = Glm4MoeLiteMoE_fast_forward
|
||||
|
|
|
|||
|
|
@ -112,7 +112,7 @@ def GraniteAttention_fast_forward(
|
|||
cos, sin = position_embeddings
|
||||
rope_position_ids = position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
if rope_position_ids is not None:
|
||||
# Useful for LongRoPE
|
||||
# Needed for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
else:
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
|
|
@ -122,7 +122,6 @@ def GraniteAttention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
use_varlen = attention_mask is None and seq_info is not None and past_key_value is None
|
||||
|
||||
backend = SDPA if attention_mask is not None else select_attention_backend(use_varlen)
|
||||
|
|
@ -255,7 +254,7 @@ def GraniteDecoderLayer_fast_forward(
|
|||
|
||||
from math import sqrt as math_sqrt
|
||||
|
||||
KV_CACHE_INCREMENT = 256 # KV Cache update size
|
||||
KV_CACHE_INCREMENT = 256
|
||||
torch_nn_functional_softmax = torch.nn.functional.softmax
|
||||
torch_matmul = torch.matmul
|
||||
torch_tanh = torch.tanh
|
||||
|
|
@ -493,8 +492,7 @@ class GraniteRotaryEmbedding(LlamaRotaryEmbedding):
|
|||
|
||||
def patched_init(original_init):
|
||||
def new_init(self, *args, **kwargs):
|
||||
# GraniteModel_fast_forward_inference can't reach residual_multiplier/config,
|
||||
# so stash the whole config here to pass it around. See:
|
||||
# Stash config so GraniteModel_fast_forward_inference can reach residual_multiplier. See:
|
||||
# https://github.com/huggingface/transformers/blob/e5fd865ebae062b7cf03a81b8c6affeb39f30bec/src/transformers/models/granite/modeling_granite.py#L243
|
||||
config = kwargs.get("config", args[0] if args else None)
|
||||
if config is not None:
|
||||
|
|
@ -538,14 +536,13 @@ class FastGraniteModel(FastLlamaModel):
|
|||
tokenizer,
|
||||
correct_dtype = None,
|
||||
):
|
||||
# Torch.compile fails on embedding matrix??
|
||||
# Workaround randomnly fixes it for torch versions < 2.2
|
||||
# Workaround for torch.compile failing on the embedding matrix (torch < 2.2)
|
||||
model.model.embed_tokens = torch.nn.Embedding.from_pretrained(
|
||||
model.model.embed_tokens.weight
|
||||
)
|
||||
model.config.update({"unsloth_version": __version__})
|
||||
|
||||
# We also do this for the lm_head
|
||||
# Same for lm_head
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.lm_head.weight
|
||||
|
|
@ -562,8 +559,7 @@ class FastGraniteModel(FastLlamaModel):
|
|||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
|
||||
# Also patch all dtypes - BnB seems to not allocate the correct type?
|
||||
# BnB default dtype seems to be float16!
|
||||
# Patch all dtypes - BnB defaults to float16 instead of the correct type
|
||||
correct_dtype = lm_head.weight.dtype
|
||||
|
||||
for name, module in model.named_modules():
|
||||
|
|
@ -572,12 +568,11 @@ class FastGraniteModel(FastLlamaModel):
|
|||
quant_state = weight.quant_state
|
||||
|
||||
if type(quant_state) is list:
|
||||
# BnB seems to have float16 as default!
|
||||
module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
|
||||
module.weight.quant_state[2] = correct_dtype # BnB defaults to float16
|
||||
else:
|
||||
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
|
||||
quant_state.dtype = correct_dtype
|
||||
# Downcast RoPE embedding to correct data type
|
||||
# Downcast RoPE embedding to correct dtype
|
||||
if name.endswith("rotary_emb") or hasattr(module, "cos_cached"):
|
||||
if hasattr(module, "cos_cached") and (module.cos_cached.dtype != correct_dtype):
|
||||
module.cos_cached = module.cos_cached.to(correct_dtype)
|
||||
|
|
@ -589,7 +584,6 @@ class FastGraniteModel(FastLlamaModel):
|
|||
module.short_cos_cached = module.short_cos_cached.to(correct_dtype)
|
||||
module.short_sin_cached = module.short_sin_cached.to(correct_dtype)
|
||||
|
||||
# Clear deleted GPU items
|
||||
import gc
|
||||
|
||||
for _ in range(3):
|
||||
|
|
|
|||
|
|
@ -158,9 +158,8 @@ def _offload_frozen_module_for_training(
|
|||
) -> None:
|
||||
"""Move the trainable copy to ``device_type`` and offload the frozen original.
|
||||
|
||||
float16 is promoted to float32 for GPU compatibility (e.g. Tesla T4).
|
||||
``offload_device`` currently only supports "cpu"; None leaves the frozen
|
||||
module in place. Modifies ``module`` in-place.
|
||||
float16 is promoted to float32 (Tesla T4). ``offload_device`` only supports
|
||||
"cpu"; None leaves the frozen module in place. Modifies ``module`` in-place.
|
||||
See https://github.com/unslothai/unsloth/pull/1200 (Tesla T4 float32).
|
||||
"""
|
||||
if not hasattr(module, "modules_to_save"):
|
||||
|
|
@ -168,8 +167,7 @@ def _offload_frozen_module_for_training(
|
|||
|
||||
new_dtype = module.modules_to_save.default.weight.dtype
|
||||
if new_dtype == torch.float16:
|
||||
# See https://github.com/unslothai/unsloth/pull/1200
|
||||
# Tesla T4 must use float32 and not float16
|
||||
# Tesla T4 must use float32 not float16. See unslothai/unsloth#1200
|
||||
new_dtype = torch.float32
|
||||
|
||||
module.modules_to_save.default.to(device = device_type, dtype = new_dtype, non_blocking = True)
|
||||
|
|
@ -333,7 +331,6 @@ def _fast_prepare_inputs_for_generation(
|
|||
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -350,33 +347,12 @@ def LlamaAttention_fast_forward_inference(
|
|||
attention_mask = None,
|
||||
rotary_seq_len = None,
|
||||
):
|
||||
"""
|
||||
"""Fast inference using KV cache.
|
||||
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.
|
||||
[Q, q] @ [K, k].T splits into 4 chunks; the mask wipes Qk^T so only the new
|
||||
row [qK^T, qk^T] needs computing (the rest is the prior attention). Hence we
|
||||
pass one row of Q but must remember K and V (the KV cache).
|
||||
"""
|
||||
Xn = hidden_states
|
||||
bsz, _, hd = hidden_states.size()
|
||||
|
|
@ -772,19 +748,7 @@ def LlamaDecoderLayer_fast_forward(
|
|||
*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
|
||||
"""
|
||||
"""Fast decoder-layer forward; hidden_states is `(batch, seq_len, embed_dim)`."""
|
||||
if use_cache and hasattr(self, "_flag_for_generation"):
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
|
||||
|
|
@ -936,7 +900,6 @@ def LlamaModel_fast_forward(
|
|||
if position_ids.shape[0] != batch_size:
|
||||
position_ids = position_ids.repeat((batch_size, 1))
|
||||
|
||||
# Embed positions
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
|
|
@ -973,8 +936,7 @@ def LlamaModel_fast_forward(
|
|||
if inputs_requires_grad:
|
||||
inputs_embeds.requires_grad_(True)
|
||||
|
||||
# Fix up attention mask by setting elements to 0
|
||||
# Specifically for DPO
|
||||
# Zero out attention mask elements, specifically for DPO
|
||||
if (
|
||||
getattr(self, "_has_no_labels", False) is True
|
||||
and (attention_mask is not None)
|
||||
|
|
@ -1031,7 +993,6 @@ def LlamaModel_fast_forward(
|
|||
# )
|
||||
# use_cache = False
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = () if use_cache else None
|
||||
|
|
@ -1042,7 +1003,6 @@ def LlamaModel_fast_forward(
|
|||
else:
|
||||
boundaries = None
|
||||
|
||||
# Check checkpointing method
|
||||
gradient_checkpointing = False
|
||||
|
||||
if self.gradient_checkpointing and self.training and not use_cache:
|
||||
|
|
@ -1138,7 +1098,6 @@ def LlamaModel_fast_forward(
|
|||
else:
|
||||
position_embeddings = None
|
||||
|
||||
# Go through every layer!
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
|
@ -1350,7 +1309,7 @@ def _LlamaModel_fast_forward_inference(
|
|||
return LlamaModel_fast_forward_inference_custom
|
||||
|
||||
|
||||
# For ensuring backwards compatibility, we create LlamaModel_fast_forward_inference that is consumed by other models
|
||||
# Backwards-compat alias consumed by other models
|
||||
LlamaModel_fast_forward_inference = _LlamaModel_fast_forward_inference()
|
||||
|
||||
|
||||
|
|
@ -2463,10 +2422,8 @@ class FastLlamaModel:
|
|||
from .loader_utils import check_and_disable_bitsandbytes_loading
|
||||
from unsloth_zoo.utils import get_quant_type
|
||||
|
||||
# Extract load_in_8bit from kwargs if provided
|
||||
load_in_8bit = kwargs.get("load_in_8bit", False)
|
||||
|
||||
# Check and disable bitsandbytes loading if model has non-bitsandbytes quantization
|
||||
load_in_4bit, load_in_8bit, _ckpt_quant_method = check_and_disable_bitsandbytes_loading(
|
||||
model_config, load_in_4bit = load_in_4bit, load_in_8bit = load_in_8bit
|
||||
)
|
||||
|
|
@ -2982,12 +2939,10 @@ class FastLlamaModel:
|
|||
modules_to_save = list(modules_to_save)
|
||||
old_target_modules += modules_to_save
|
||||
|
||||
# Combine all
|
||||
new_target_modules = list(target_modules) + list(
|
||||
modules_to_save if modules_to_save is not None else []
|
||||
)
|
||||
|
||||
# Now check!
|
||||
new_target_modules = set(new_target_modules)
|
||||
check_all = check_all and (len(set(old_target_modules) ^ new_target_modules) == 0)
|
||||
|
||||
|
|
@ -2997,10 +2952,8 @@ class FastLlamaModel:
|
|||
)
|
||||
|
||||
if check_all:
|
||||
# Simply pass through!
|
||||
logger.warning("Unsloth: Already have LoRA adapters! We shall skip this step.")
|
||||
|
||||
# Offload!
|
||||
# [TODO] First offload lm_head and embed_tokens to CPU (should be disk!!)
|
||||
if "embed_tokens" in new_target_modules:
|
||||
print("Unsloth: Training embed_tokens in mixed precision to save VRAM")
|
||||
|
|
|
|||
|
|
@ -104,10 +104,8 @@ from ._utils import (
|
|||
set_task_config_attr,
|
||||
)
|
||||
|
||||
# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers
|
||||
# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working.
|
||||
# Fallback list mirrors zoo for users who upgrade unsloth without upgrading
|
||||
# unsloth_zoo (so this module never fails at import).
|
||||
# Re-export FORCE_FLOAT32 from unsloth_zoo (single source of truth); fallback list
|
||||
# below keeps import working when unsloth_zoo is older than unsloth.
|
||||
try:
|
||||
from unsloth_zoo import FORCE_FLOAT32 # noqa: F401
|
||||
except ImportError:
|
||||
|
|
@ -332,8 +330,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
load_in_8bit = True
|
||||
load_in_4bit = False
|
||||
|
||||
# Login to allow private models
|
||||
token = hf_login(token)
|
||||
token = hf_login(token) # Login to allow private models
|
||||
# Align dtype with bnb_4bit_compute_dtype if provided and dtype is unset.
|
||||
if dtype is None and quantization_config is not None:
|
||||
bnb_compute_dtype = None
|
||||
|
|
@ -430,7 +427,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
fast_inference = False
|
||||
break
|
||||
|
||||
# Check if 4bit is allowed specifically for AMD
|
||||
# AMD is unstable with 4bit bitsandbytes
|
||||
if not ALLOW_BITSANDBYTES and not use_exact_model_name:
|
||||
if load_in_4bit or load_in_8bit or model_name.lower().endswith("-bnb-4bit"):
|
||||
print(
|
||||
|
|
@ -467,13 +464,12 @@ class FastLanguageModel(FastLlamaModel):
|
|||
if load_in_fp8 != False and new_model_name != old_model_name:
|
||||
load_in_fp8 = False
|
||||
|
||||
# Check if pre-quantized models are allowed
|
||||
# AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 64)
|
||||
# AMD Instinct GPUs need blocksize 128 on bitsandbytes < 0.49.2 (our pre-quants use 64)
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
# -BF16 means 16bit only
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
|
|
@ -484,7 +480,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
from modelscope import snapshot_download
|
||||
model_name = snapshot_download(model_name)
|
||||
|
||||
# First check if it's a normal model via AutoConfig
|
||||
from huggingface_hub.utils import (
|
||||
disable_progress_bars,
|
||||
enable_progress_bars,
|
||||
|
|
@ -549,7 +544,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
# Old transformers versions check
|
||||
both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32
|
||||
|
||||
# Error out if both LoRA and normal model config exists.
|
||||
if both_exist:
|
||||
raise RuntimeError(
|
||||
"Unsloth: Your repo has a LoRA adapter and a base model.\n"
|
||||
|
|
@ -569,7 +563,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
|
||||
# New transformers need to check manually.
|
||||
if SUPPORTS_LLAMA32 and is_model and is_peft:
|
||||
# Check if folder exists locally
|
||||
if os.path.isdir(model_name):
|
||||
exist_adapter_config = os.path.exists(
|
||||
os.path.join(model_name, "adapter_config.json")
|
||||
|
|
@ -592,7 +585,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
f'Try `pip install --upgrade "transformers>=4.43.2"`\n'
|
||||
f"to obtain the latest transformers build, then restart this session."
|
||||
)
|
||||
# Create a combined error message showing both failures
|
||||
combined_error = (
|
||||
"Unsloth: Failed to load model. Both AutoConfig and PeftConfig loading failed.\n\n"
|
||||
f"AutoConfig error: {autoconfig_error}\n\n"
|
||||
|
|
@ -600,9 +592,8 @@ class FastLanguageModel(FastLlamaModel):
|
|||
)
|
||||
raise RuntimeError(combined_error)
|
||||
|
||||
# Get base model for PEFT:
|
||||
# Get base model for PEFT
|
||||
if is_peft:
|
||||
# Check base model again for PEFT
|
||||
model_name = peft_config.base_model_name_or_path
|
||||
if not use_exact_model_name:
|
||||
model_name = get_model_name(
|
||||
|
|
@ -612,13 +603,12 @@ class FastLanguageModel(FastLlamaModel):
|
|||
token = token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
# Check if pre-quantized models are allowed
|
||||
# AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 64)
|
||||
# AMD Instinct GPUs need blocksize 128 on bitsandbytes < 0.49.2 (our pre-quants use 64)
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
# -BF16 means 16bit only
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
|
|
@ -750,7 +740,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
# Apply gradient checkpointing with smart heuristics
|
||||
use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
|
||||
use_gradient_checkpointing, max_seq_length, dtype
|
||||
)
|
||||
|
|
@ -809,7 +798,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
if resize_model_vocab is not None:
|
||||
model.resize_token_embeddings(resize_model_vocab)
|
||||
|
||||
# In case the model supports tagging, add the unsloth tag.
|
||||
if hasattr(model, "add_model_tags"):
|
||||
model.add_model_tags(
|
||||
[
|
||||
|
|
@ -852,7 +840,6 @@ class FastLanguageModel(FastLlamaModel):
|
|||
|
||||
if is_peft:
|
||||
# From https://github.com/huggingface/peft/issues/184
|
||||
# Now add PEFT adapters
|
||||
model = PeftModel.from_pretrained(
|
||||
model,
|
||||
old_model_name,
|
||||
|
|
@ -861,11 +848,9 @@ class FastLanguageModel(FastLlamaModel):
|
|||
is_trainable = True,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
# Patch it as well!
|
||||
model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing)
|
||||
|
||||
# Patch Tiled MLP
|
||||
# to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
|
||||
# Tiled MLP: set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"
|
||||
patch_tiled_mlp_choice = os.environ.get(
|
||||
"UNSLOTH_TILED_MLP", "arctic" if unsloth_tiled_mlp else "0"
|
||||
)
|
||||
|
|
@ -944,7 +929,6 @@ class FastModel(FastBaseModel):
|
|||
unsloth_force_compile = False,
|
||||
offload_embedding = False,
|
||||
float32_mixed_precision = None, # Forces float32 mixed precision
|
||||
# Add the missing vLLM/inference parameters
|
||||
fast_inference = False, # uses vLLM
|
||||
gpu_memory_utilization = 0.5,
|
||||
float8_kv_cache = False,
|
||||
|
|
@ -976,8 +960,7 @@ class FastModel(FastBaseModel):
|
|||
load_in_8bit = True
|
||||
load_in_4bit = False
|
||||
|
||||
# Login to allow private models
|
||||
token = hf_login(token)
|
||||
token = hf_login(token) # Login to allow private models
|
||||
if whisper_language is not None:
|
||||
assert type(whisper_language) is str
|
||||
if whisper_task is not None:
|
||||
|
|
@ -1045,7 +1028,7 @@ class FastModel(FastBaseModel):
|
|||
if is_dist:
|
||||
device_map = distributed_device_map
|
||||
|
||||
# Check if 4bit is allowed specifically for AMD
|
||||
# AMD is unstable with 4bit bitsandbytes
|
||||
if not ALLOW_BITSANDBYTES and not use_exact_model_name:
|
||||
if load_in_4bit or load_in_8bit or model_name.lower().endswith("-bnb-4bit"):
|
||||
print(
|
||||
|
|
@ -1095,25 +1078,22 @@ class FastModel(FastBaseModel):
|
|||
if load_in_fp8 != False and new_model_name != old_model_name:
|
||||
load_in_fp8 = False
|
||||
|
||||
# Check if pre-quantized models are allowed
|
||||
# AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 64)
|
||||
# AMD Instinct GPUs need blocksize 128 on bitsandbytes < 0.49.2 (our pre-quants use 64)
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
# -BF16 means 16bit only
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
load_in_fp8 = False
|
||||
load_in_16bit = True
|
||||
|
||||
# Check modelscope
|
||||
if USE_MODELSCOPE and not os.path.exists(model_name):
|
||||
from modelscope import snapshot_download
|
||||
model_name = snapshot_download(model_name)
|
||||
|
||||
# First check if it's a normal model via AutoConfig
|
||||
from huggingface_hub.utils import (
|
||||
disable_progress_bars,
|
||||
enable_progress_bars,
|
||||
|
|
@ -1211,7 +1191,6 @@ class FastModel(FastBaseModel):
|
|||
is_peft = False
|
||||
# Old transformers versions check
|
||||
both_exist = (is_model and is_peft) and not SUPPORTS_LLAMA32
|
||||
# Error out if both LoRA and normal model config exists.
|
||||
if both_exist:
|
||||
raise RuntimeError(
|
||||
"Unsloth: Your repo has a LoRA adapter and a base model.\n"
|
||||
|
|
@ -1389,7 +1368,6 @@ class FastModel(FastBaseModel):
|
|||
|
||||
# New transformers need to check manually.
|
||||
if SUPPORTS_LLAMA32 and is_model and is_peft:
|
||||
# Check if folder exists locally
|
||||
if os.path.isdir(model_name):
|
||||
exist_adapter_config = os.path.exists(
|
||||
os.path.join(model_name, "adapter_config.json")
|
||||
|
|
@ -1412,7 +1390,6 @@ class FastModel(FastBaseModel):
|
|||
f'Try `pip install --upgrade "transformers>=4.43.2"`\n'
|
||||
f"to obtain the latest transformers build, then restart this session."
|
||||
)
|
||||
# Create a combined error message showing both failures
|
||||
combined_error = (
|
||||
"Unsloth: Failed to load model. Both AutoConfig and PeftConfig loading failed.\n\n"
|
||||
f"AutoConfig error: {autoconfig_error}\n\n"
|
||||
|
|
@ -1420,19 +1397,17 @@ class FastModel(FastBaseModel):
|
|||
)
|
||||
raise RuntimeError(combined_error)
|
||||
|
||||
# Get base model for PEFT:
|
||||
# Get base model for PEFT
|
||||
if is_peft:
|
||||
# Check base model again for PEFT
|
||||
model_name = peft_config.base_model_name_or_path
|
||||
if not use_exact_model_name:
|
||||
model_name = get_model_name(model_name, load_in_4bit)
|
||||
# Check if pre-quantized models are allowed
|
||||
# AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 64)
|
||||
# AMD Instinct GPUs need blocksize 128 on bitsandbytes < 0.49.2 (our pre-quants use 64)
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
# -BF16 means 16bit only
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
|
|
@ -1459,14 +1434,13 @@ class FastModel(FastBaseModel):
|
|||
redirector = contextlib.redirect_stdout(open(os.devnull, "w"))
|
||||
|
||||
model_types = ["siglip"] + model_types
|
||||
# Set forced float32 env flag
|
||||
os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
|
||||
do_forced_float32 = False
|
||||
for model_type_arch in model_types:
|
||||
if model_type_arch != "siglip":
|
||||
break
|
||||
for disable_name in FORCE_FLOAT32:
|
||||
# add comma to model_types_all matching in case of exact match for end
|
||||
# model_types_all has a trailing comma so suffixes match exactly
|
||||
if (
|
||||
disable_name.lower() == model_type_arch.lower().replace("-", "").replace("_", "")
|
||||
or disable_name.lower() in model_types_all
|
||||
|
|
@ -1474,7 +1448,6 @@ class FastModel(FastBaseModel):
|
|||
os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
|
||||
dtype = torch.bfloat16 # Change to bfloat16 loading
|
||||
break
|
||||
# Apply gradient checkpointing with smart heuristics
|
||||
use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
|
||||
use_gradient_checkpointing, max_seq_length, dtype
|
||||
)
|
||||
|
|
@ -1538,7 +1511,6 @@ class FastModel(FastBaseModel):
|
|||
for _cfg_key, _cfg_val in task_config_attrs.items():
|
||||
set_task_config_attr(model_config, _cfg_key, _cfg_val)
|
||||
|
||||
# Check if VLM
|
||||
architectures = getattr(model_config, "architectures", None)
|
||||
if architectures is None:
|
||||
architectures = []
|
||||
|
|
@ -1631,7 +1603,6 @@ class FastModel(FastBaseModel):
|
|||
if resize_model_vocab is not None:
|
||||
model.resize_token_embeddings(resize_model_vocab)
|
||||
|
||||
# In case the model supports tagging, add the unsloth tag.
|
||||
if hasattr(model, "add_model_tags"):
|
||||
model.add_model_tags(
|
||||
[
|
||||
|
|
@ -1674,7 +1645,6 @@ class FastModel(FastBaseModel):
|
|||
|
||||
if is_peft:
|
||||
# From https://github.com/huggingface/peft/issues/184
|
||||
# Now add PEFT adapters
|
||||
|
||||
# Gemma4 ClippableLinear wraps nn.Linear -- PEFT can't inject LoRA
|
||||
# on it directly. Monkey-patch PEFT to target the inner .linear
|
||||
|
|
@ -1741,18 +1711,15 @@ class FastModel(FastBaseModel):
|
|||
if _clippable_linear_cls is not None:
|
||||
_LoraModel._create_and_replace = _original_car
|
||||
|
||||
# Patch it as well!
|
||||
model = FastBaseModel.post_patch_model(
|
||||
model, use_gradient_checkpointing, trust_remote_code = trust_remote_code
|
||||
)
|
||||
|
||||
# Apply QAT if specified
|
||||
if qat_scheme is not None:
|
||||
print("Unsloth: Applying QAT to mitigate quantization degradation")
|
||||
model = FastModel._prepare_for_qat(model, qat_scheme)
|
||||
|
||||
# Patch Tiled MLP
|
||||
# to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
|
||||
# Tiled MLP: set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"
|
||||
patch_tiled_mlp_choice = os.environ.get(
|
||||
"UNSLOTH_TILED_MLP", "arctic" if unsloth_tiled_mlp else "0"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ BAD_MAPPINGS = {
|
|||
|
||||
|
||||
def _get_torchao_fp8_config(fp8_mode):
|
||||
# Lazy import so a broken optional vLLM install doesn't break `import unsloth`.
|
||||
# Lazy import: a broken optional vLLM install must not break `import unsloth`
|
||||
from unsloth_zoo.vllm_utils import _get_torchao_fp8_config as _impl
|
||||
return _impl(fp8_mode)
|
||||
|
||||
|
|
@ -118,10 +118,10 @@ def __get_model_name(
|
|||
if load_in_fp8 != False:
|
||||
if load_in_fp8 == True and (os.environ.get("UNSLOTH_HAS_FBGEMM", "0") == "1"):
|
||||
if lower_model_name in FLOAT_TO_FP8_ROW_MAPPER:
|
||||
# Faster row scaling only works if FBGEMM works!
|
||||
# Faster row scaling needs FBGEMM
|
||||
return FLOAT_TO_FP8_ROW_MAPPER[lower_model_name]
|
||||
elif lower_model_name in FLOAT_TO_FP8_BLOCK_MAPPER:
|
||||
# Otherwise we use the slower blockwise type
|
||||
# Fall back to slower blockwise scaling
|
||||
return FLOAT_TO_FP8_BLOCK_MAPPER[lower_model_name]
|
||||
else:
|
||||
if lower_model_name in FLOAT_TO_FP8_BLOCK_MAPPER:
|
||||
|
|
@ -238,7 +238,7 @@ def get_model_name(
|
|||
new_model_name = BAD_MAPPINGS[new_model_name.lower()]
|
||||
|
||||
if new_model_name is None and model_name.count("/") == 1 and model_name[0].isalnum():
|
||||
# Try checking if a new Unsloth version allows it!
|
||||
# Maybe a newer Unsloth version supports it
|
||||
NEW_INT_TO_FLOAT_MAPPER, NEW_FLOAT_TO_INT_MAPPER, NEW_MAP_TO_UNSLOTH_16bit = (
|
||||
_get_new_mapper()
|
||||
)
|
||||
|
|
@ -361,28 +361,16 @@ def check_and_disable_bitsandbytes_loading(
|
|||
verbose = True,
|
||||
):
|
||||
"""
|
||||
Check if we should disable bitsandbytes loading (load_in_4bit/load_in_8bit)
|
||||
because the model already has a non-bitsandbytes quantization config.
|
||||
If so, disable BOTH 4bit and 8bit loading and print a warning message.
|
||||
|
||||
Args:
|
||||
model_config: The AutoConfig object from the model
|
||||
load_in_4bit: Whether load_in_4bit is currently enabled
|
||||
load_in_8bit: Whether load_in_8bit is currently enabled
|
||||
verbose: Whether to print warning messages
|
||||
|
||||
Returns:
|
||||
tuple: (load_in_4bit, load_in_8bit, quant_method)
|
||||
load_in_4bit/load_in_8bit will be False if they were disabled
|
||||
quant_method is the detected quantization method or None
|
||||
Disable bnb 4bit/8bit loading if the model already has a non-bnb quant config,
|
||||
to avoid config conflicts. Returns (load_in_4bit, load_in_8bit, quant_method),
|
||||
where the flags are False if disabled and quant_method is the detected method or None.
|
||||
"""
|
||||
quant_method = get_quant_type(model_config)
|
||||
|
||||
if quant_method is None or quant_method == "bitsandbytes":
|
||||
return load_in_4bit, load_in_8bit, quant_method
|
||||
|
||||
# Model has a non-bitsandbytes quantization config (e.g., compressed-tensors, gptq, awq)
|
||||
# We should disable BOTH bitsandbytes loading to avoid config conflicts
|
||||
# Non-bnb quant config (compressed-tensors/gptq/awq): disable bnb to avoid config conflicts
|
||||
if load_in_4bit or load_in_8bit:
|
||||
if verbose:
|
||||
print(
|
||||
|
|
@ -413,7 +401,6 @@ def _get_fp8_mode_and_check_settings(
|
|||
else:
|
||||
fp8_mode = load_in_fp8
|
||||
|
||||
# Check user settings
|
||||
if fp8_mode not in ["row", "block"]:
|
||||
raise ValueError(f"Unsloth: `load_in_fp8` can only be 'row' or 'block', got '{fp8_mode}'")
|
||||
if full_finetuning:
|
||||
|
|
@ -423,7 +410,7 @@ def _get_fp8_mode_and_check_settings(
|
|||
"Unsloth: `load_in_fp8` is not compatible with `load_in_4bit`, `load_in_8bit` or `load_in_16bit`",
|
||||
)
|
||||
|
||||
# Check if this is Hopper or above
|
||||
# Require Hopper or above
|
||||
if not (
|
||||
torch.cuda.is_available()
|
||||
and torch.version.cuda
|
||||
|
|
@ -433,7 +420,6 @@ def _get_fp8_mode_and_check_settings(
|
|||
"Unsloth: On the fly `load_in_fp8` requires H100 GPUs or after. Try `unsloth/Qwen3-8B` instead."
|
||||
)
|
||||
|
||||
# Check if torch >= 2.9.0
|
||||
if Version(torch.__version__) < Version("2.9.0"):
|
||||
raise ValueError(
|
||||
"Unsloth: On the fly `load_in_fp8` requires torch 2.9.0+. Try `unsloth/Qwen3-8B` instead."
|
||||
|
|
@ -455,14 +441,13 @@ def _get_fp8_mode_and_check_settings(
|
|||
if Version(torchao.__version__) < Version("0.15.0"):
|
||||
raise ValueError(error_message)
|
||||
|
||||
# If fbgemm_gpu_genai is installed and old, disable FBGEMM and use Triton instead
|
||||
# Old fbgemm_gpu_genai: disable FBGEMM, use Triton instead
|
||||
if (
|
||||
importlib.util.find_spec("fbgemm_gpu") is not None
|
||||
and importlib.util.find_spec("fbgemm_gpu.experimental") is not None
|
||||
):
|
||||
import fbgemm_gpu.experimental.gen_ai
|
||||
if Version(fbgemm_gpu.__version__) < Version("1.4.1"):
|
||||
# Old FBGEMM version - disable and use Triton kernels instead
|
||||
os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
|
||||
from unsloth_zoo.log import logger
|
||||
logger.info(
|
||||
|
|
|
|||
|
|
@ -1315,7 +1315,6 @@ __INT_TO_FLOAT_MAPPER = \
|
|||
"google/functiongemma-270m-it",
|
||||
"unsloth/functiongemma-270m-it-unsloth-bnb-4bit",
|
||||
),
|
||||
# Ministral 3 models
|
||||
"unsloth/Ministral-3-3B-Instruct-2512-unsloth-bnb-4bit" : {
|
||||
"8" : (
|
||||
"mistralai/Ministral-3-3B-Instruct-2512",
|
||||
|
|
@ -1455,7 +1454,6 @@ for key, values in __INT_TO_FLOAT_MAPPER.items():
|
|||
_add_with_lower(MAP_TO_UNSLOTH_16bit, row, values[0])
|
||||
pass
|
||||
|
||||
# Get lowercased
|
||||
lowered_key = key.lower()
|
||||
INT_TO_FLOAT_MAPPER[lowered_key] = values[0].lower()
|
||||
|
||||
|
|
|
|||
|
|
@ -65,7 +65,6 @@ def MistralAttention_fast_forward(
|
|||
*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
|
||||
|
|
@ -98,7 +97,7 @@ def MistralAttention_fast_forward(
|
|||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
rope_position_ids = position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
# Useful for LongRoPE
|
||||
# rope_position_ids enables LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
|
|
@ -106,7 +105,6 @@ def MistralAttention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
sw_cfg = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw_cfg is None or sw_cfg == "null") else sw_cfg
|
||||
window_size = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
|
|
@ -176,23 +174,18 @@ def MistralForCausalLM_fast_forward(
|
|||
[q_len] * bsz
|
||||
).make_local_attention(window_size = sliding_window)
|
||||
|
||||
# If attention_mask exists, it will be handled in the attention forward
|
||||
|
||||
elif self.training:
|
||||
# LlamaModel_fast_forward's DPO embed-masking block needs the 2D
|
||||
# attention_mask; it nulls the mask before attention anyway, so
|
||||
# leaving it 2D is safe and avoids a 4D conversion that crashes DPO.
|
||||
# Keep 2D attention_mask: DPO embed-masking nulls it before attention,
|
||||
# and a 4D conversion would crash DPO.
|
||||
pass
|
||||
|
||||
else:
|
||||
# Not using xformers - need to create attention masks
|
||||
if (
|
||||
sliding_window is None
|
||||
or sliding_window == "null"
|
||||
or sliding_window <= 0
|
||||
or q_len <= sliding_window
|
||||
):
|
||||
# Fully causal mask
|
||||
causal_mask_values = torch.triu(
|
||||
torch.full((q_len, q_len), -torch.inf, device = input_ids.device),
|
||||
diagonal = 1,
|
||||
|
|
@ -209,7 +202,6 @@ def MistralForCausalLM_fast_forward(
|
|||
causal_bool_mask & window_bool_mask, 0.0, -torch.inf
|
||||
)
|
||||
|
||||
# Combine with existing attention_mask if present
|
||||
if attention_mask is None:
|
||||
attention_mask = causal_mask_values[None, None, :, :].expand(bsz, 1, q_len, q_len)
|
||||
else:
|
||||
|
|
@ -236,7 +228,6 @@ def MistralForCausalLM_fast_forward(
|
|||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
self.model._has_no_labels = labels is None
|
||||
|
||||
if past_key_values is not None:
|
||||
|
|
@ -268,14 +259,12 @@ def MistralForCausalLM_fast_forward(
|
|||
lm_head = self.lm_head.weight
|
||||
lm_head_device = lm_head.device
|
||||
|
||||
# Move items to same device as lm_head
|
||||
# Move to lm_head device
|
||||
hidden_states = hidden_states.to(lm_head_device)
|
||||
if labels is not None:
|
||||
labels = labels.to(lm_head_device)
|
||||
|
||||
# Merge legacy / new spellings before branching so the decode-time
|
||||
# last-token slice fires on the normal path too. Skip int max() if
|
||||
# either is a tensor (HF selective-decode form).
|
||||
# Merge legacy/new spellings; skip int max() if either is a tensor (HF selective-decode form)
|
||||
if isinstance(num_logits_to_keep, torch.Tensor) or isinstance(logits_to_keep, torch.Tensor):
|
||||
num_logits_to_keep = 0
|
||||
else:
|
||||
|
|
@ -300,9 +289,8 @@ def MistralForCausalLM_fast_forward(
|
|||
logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype))
|
||||
else:
|
||||
RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
|
||||
# < 1024 Normal Unsloth uses less VRAM!
|
||||
# Small batches: fused CE loss uses less VRAM
|
||||
if bsz * q_len <= 1024 and not RETURN_LOGITS:
|
||||
# Use unsloth_fused_ce_loss which actually calculates the best chunk size to reduce VRAM usage
|
||||
RETURN_LOGITS = False
|
||||
|
||||
if not RETURN_LOGITS and labels is not None:
|
||||
|
|
@ -410,7 +398,7 @@ class FastMistralModel(FastLlamaModel):
|
|||
scaled_rope_module = LlamaLinearScalingRotaryEmbedding,
|
||||
attention_module = MistralAttention,
|
||||
)
|
||||
# Just for Mistral Nemo models!
|
||||
# Mistral Nemo only
|
||||
if function is not None and init_name is not None:
|
||||
function = patch_mistral_nemo_attention(function)
|
||||
# if True:#init_name is not None:
|
||||
|
|
@ -425,9 +413,8 @@ class FastMistralModel(FastLlamaModel):
|
|||
PeftModelForCausalLM.forward = PeftModel_fast_forward
|
||||
fix_prepare_inputs_for_generation(MistralForCausalLM)
|
||||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Retain old rotary embeddings: static KV Cache (4.38.0) made training much slower.
|
||||
# https://github.com/unslothai/unsloth/issues/168
|
||||
# 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.mistral.modeling_mistral
|
||||
|
|
|
|||
|
|
@ -56,9 +56,7 @@ class FastQwen2Model(FastLlamaModel):
|
|||
PeftModelForCausalLM.forward = PeftModel_fast_forward
|
||||
fix_prepare_inputs_for_generation(Qwen2ForCausalLM)
|
||||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Retain old rotary embeddings: static KV cache (4.38.0+) slowed training. Solves issue #168
|
||||
# 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.qwen2.modeling_qwen2
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ try:
|
|||
)
|
||||
except:
|
||||
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.50.3"):
|
||||
raise ImportError(
|
||||
f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3 and Qwen3Moe.\n"
|
||||
f"The minimum required version is 4.50.3.\n"
|
||||
|
|
@ -75,7 +75,6 @@ def Qwen3Attention_fast_forward(
|
|||
*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
|
||||
|
|
@ -132,7 +131,6 @@ def Qwen3Attention_fast_forward(
|
|||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
backend = SDPA if attention_mask is not None else select_attention_backend(use_varlen)
|
||||
attention_config = AttentionConfig(
|
||||
|
|
@ -201,7 +199,6 @@ def Qwen3Attention_fast_forward_inference(
|
|||
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:
|
||||
|
|
@ -263,8 +260,7 @@ def Qwen3Attention_fast_forward_inference(
|
|||
# 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
|
||||
# extend 2 steps ahead before the short KV cache fills, else error
|
||||
self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
|
||||
# Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
|
||||
|
|
@ -342,7 +338,6 @@ def Qwen3Attention_fast_forward_inference(
|
|||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
|
||||
# Attention
|
||||
if bsz == 1:
|
||||
Qn *= (
|
||||
self.scalar
|
||||
|
|
@ -402,7 +397,7 @@ class FastQwen3Model(FastLlamaModel):
|
|||
return
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained( # TODO: Change after release
|
||||
def from_pretrained(
|
||||
model_name = "Qwen/Qwen3-7B",
|
||||
max_seq_length = 4096,
|
||||
dtype = None,
|
||||
|
|
|
|||
|
|
@ -67,11 +67,10 @@ def Qwen3MoeSparseMoeBlock_fast_forward(
|
|||
routing_weights = torch_nn_functional_softmax(router_logits, dim = -1, dtype = torch.float32)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim = -1)
|
||||
routing_weights /= routing_weights.sum(dim = -1, keepdim = True)
|
||||
# cast back to the input dtype
|
||||
routing_weights = routing_weights.to(X.dtype)
|
||||
final_X = torch.zeros((bsz * seq_len, hd), dtype = torch.float32, device = X.device)
|
||||
|
||||
# One-hot the selected experts into a mask to index which expert is used
|
||||
# One-hot the selected experts into a mask indexing which expert is used
|
||||
expert_mask = torch.nn.functional.one_hot(
|
||||
selected_experts, num_classes = self.num_experts
|
||||
).permute(2, 1, 0)
|
||||
|
|
@ -80,7 +79,6 @@ def Qwen3MoeSparseMoeBlock_fast_forward(
|
|||
expert_layer = self.experts[expert_idx]
|
||||
idx, top_x = torch.where(expert_mask[expert_idx])
|
||||
|
||||
# Index hidden states for this expert and scale by routing_weights
|
||||
current_state = X[None, top_x].reshape(-1, hd)
|
||||
current_X = (
|
||||
expert_layer(current_state) * routing_weights[top_x, idx, None]
|
||||
|
|
|
|||
|
|
@ -72,13 +72,13 @@ except Exception:
|
|||
except Exception:
|
||||
trl_version = Version("0.0.0")
|
||||
|
||||
# Get PyTorch version for feature detection
|
||||
# PyTorch version for feature detection
|
||||
try:
|
||||
torch_version = Version(torch.__version__.split("+")[0].split("a")[0].split("b")[0])
|
||||
except Exception:
|
||||
torch_version = Version("0.0.0")
|
||||
|
||||
# Get transformers version for feature detection
|
||||
# transformers version for feature detection
|
||||
try:
|
||||
from transformers import __version__ as _transformers_version_raw
|
||||
transformers_version = Version(_transformers_version_raw)
|
||||
|
|
@ -186,10 +186,8 @@ def PatchRL(FastLanguageModel):
|
|||
|
||||
@contextmanager
|
||||
def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
|
||||
# why: snapshot before TRL's unwrap context manager, which calls
|
||||
# gradient_checkpointing_disable() before yielding; preserve the actual
|
||||
# mode value (e.g. "unsloth") rather than collapsing it to a bool, so
|
||||
# the finally restore matches the caller's configured GC mode.
|
||||
# Snapshot the GC mode before TRL's unwrap CM calls gradient_checkpointing_disable();
|
||||
# keep the real value (e.g. "unsloth") not a bool so the finally restore matches.
|
||||
use_gradient_checkpointing = next(
|
||||
(
|
||||
v
|
||||
|
|
@ -199,7 +197,6 @@ def PatchRL(FastLanguageModel):
|
|||
False,
|
||||
)
|
||||
with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
|
||||
# Put the model in inference mode.
|
||||
FastLanguageModel.for_inference(model)
|
||||
|
||||
# We must use .clone for Unsloth since we force inference_mode
|
||||
|
|
@ -217,7 +214,6 @@ def PatchRL(FastLanguageModel):
|
|||
try:
|
||||
yield unwrapped_model
|
||||
finally:
|
||||
# Restore generate and return
|
||||
unwrapped_model.generate = original_generate
|
||||
FastLanguageModel.for_training(
|
||||
model,
|
||||
|
|
@ -229,24 +225,8 @@ def PatchRL(FastLanguageModel):
|
|||
|
||||
@torch.no_grad()
|
||||
def unsloth_prediction_step(self, model, inputs, prediction_loss_only, ignore_keys):
|
||||
"""
|
||||
Perform an evaluation step on `model` using `inputs`.
|
||||
Subclass and override to inject custom behavior.
|
||||
Args:
|
||||
model (`nn.Module`):
|
||||
The model to evaluate.
|
||||
inputs (`Dict[str, Union[torch.Tensor, Any]]`):
|
||||
The inputs and targets of the model.
|
||||
The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
|
||||
argument `labels`. Check your model's documentation for all accepted arguments.
|
||||
prediction_loss_only (`bool`):
|
||||
Whether or not to return the loss only.
|
||||
ignore_keys (`List[str]`, *optional*):
|
||||
A list of keys in the output of your model (if it is a dictionary) that should be ignored when
|
||||
gathering predictions.
|
||||
Return:
|
||||
Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
|
||||
logits and labels (each being optional).
|
||||
"""Evaluation step on `model` using `inputs`.
|
||||
Returns (loss, logits, labels), each optional.
|
||||
"""
|
||||
has_labels = (
|
||||
False
|
||||
|
|
@ -2251,11 +2231,8 @@ def PatchFastRL(algorithm = None, FastLanguageModel = None):
|
|||
# pristine upstream class, not the compiled Unsloth* wrappers.
|
||||
if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
|
||||
return
|
||||
# Install the disable_gradient_checkpointing noop BEFORE
|
||||
# patch_trl_rl_trainers, which imports extra trl.* submodules; any module
|
||||
# imported after the sys.modules walk would keep the original broken
|
||||
# binding. Installing first ensures the canonical symbol is replaced before
|
||||
# those submodules bind it.
|
||||
# Must run before patch_trl_rl_trainers: it imports more trl.* submodules, and any
|
||||
# imported after the sys.modules walk would keep the original broken binding.
|
||||
patch_trl_disable_gradient_checkpointing()
|
||||
patch_trl_rl_trainers()
|
||||
patch_trl_openenv()
|
||||
|
|
|
|||
|
|
@ -420,7 +420,6 @@ def sft_trainer_prepare_dataset(function_name, function):
|
|||
flags = re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
if matched:
|
||||
# Use fast version!
|
||||
function = inspect.getsource(fast_sft_prepare_dataset)
|
||||
function = function.split("\n")
|
||||
function = "\n".join(" " * 4 + x for x in function)
|
||||
|
|
@ -598,7 +597,6 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
if function_name != "_prepare_inputs":
|
||||
return function
|
||||
|
||||
# Add mixed precision training
|
||||
function = function.replace(
|
||||
"with torch.inference_mode():",
|
||||
"with torch.inference_mode(), "
|
||||
|
|
@ -728,7 +726,6 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
|
|||
# Left pad prompt before calculation old and ref hidden states
|
||||
line_to_replace = 'batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size'
|
||||
|
||||
# The new multi-line string that will replace the line above
|
||||
replacement_lines = """
|
||||
max_left_pad = None
|
||||
batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size
|
||||
|
|
@ -759,7 +756,6 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
|
|||
if self.args.gradient_accumulation_steps % generate_every != 0 or (
|
||||
self.use_vllm
|
||||
):"""
|
||||
# Use re.sub() to perform the replacement
|
||||
function, num_replacements = pattern_to_find.subn(replacement_text, function)
|
||||
|
||||
pattern_to_find = re.compile(
|
||||
|
|
@ -1455,12 +1451,9 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entrop
|
|||
|
||||
|
||||
def _unsloth_get_final_logit_softcapping(config):
|
||||
"""Return final_logit_softcapping for a model config, falling back to the
|
||||
nested text sub-config for composite models. Handles both:
|
||||
- Gemma-4-style configs where the attribute lives on ``config.text_config``
|
||||
- T5Gemma-style composite configs where the text sub-config is only
|
||||
reachable via ``config.get_text_config()``
|
||||
Returns 0 if unset, matching the previous behaviour.
|
||||
"""Return final_logit_softcapping for a config, falling back to the nested text sub-config for
|
||||
composite models (Gemma-4 ``config.text_config`` or T5Gemma ``config.get_text_config()``).
|
||||
Returns 0 if unset, matching previous behaviour.
|
||||
"""
|
||||
softcap = getattr(config, "final_logit_softcapping", None)
|
||||
if softcap is None:
|
||||
|
|
|
|||
|
|
@ -152,10 +152,7 @@ def _save_pretrained_gguf(
|
|||
maximum_memory_usage = 0.85,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Saves the SentenceTransformer model to GGUF format by saving the inner transformer model,
|
||||
converting it, and placing the resulting GGUF files in the save directory.
|
||||
"""
|
||||
"""Save the SentenceTransformer to GGUF: convert the inner transformer and place the GGUF files in save_directory."""
|
||||
# 1. Save standard SentenceTransformer structure (configs, modules.json, etc.)
|
||||
self.save_pretrained(save_directory)
|
||||
|
||||
|
|
@ -296,52 +293,20 @@ def _push_to_hub_gguf(
|
|||
tags = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Converts the SentenceTransformer model to GGUF format and pushes to the Hugging Face Hub.
|
||||
"""Convert the SentenceTransformer to GGUF and push it to the Hugging Face Hub, returning the full repo ID.
|
||||
|
||||
This method:
|
||||
1. Saves the model locally to a temporary directory in GGUF format.
|
||||
2. Uploads the GGUF files, config, Ollama Modelfile, and README to the Hub.
|
||||
3. Cleans up the temporary directory.
|
||||
|
||||
Args:
|
||||
repo_id (str): The Hugging Face Hub repo ID (e.g., "username/model-name").
|
||||
tokenizer: The tokenizer to save. Defaults to `self.tokenizer`.
|
||||
quantization_method (str or list): GGUF quantization method(s). Can be a string or list of strings.
|
||||
Choose from the following options:
|
||||
* "not_quantized" : Recommended. Fast conversion. Slow inference, big files.
|
||||
* "fast_quantized" : Recommended. Fast conversion. OK inference, OK file size.
|
||||
* "quantized" : Recommended. Slow conversion. Fast inference, small files.
|
||||
* "f32" : Not recommended. Retains 100% accuracy, but super slow and memory hungry.
|
||||
* "f16" : Fastest conversion + retains 100% accuracy. Slow and memory hungry.
|
||||
* "q8_0" : Fast conversion. High resource use, but generally acceptable.
|
||||
* "q4_k_m" : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q4_K
|
||||
* "q5_k_m" : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q5_K
|
||||
* "q2_k" : Uses Q4_K for the attention.vw and feed_forward.w2 tensors, Q2_K for the other tensors.
|
||||
* "q3_k_l" : Uses Q5_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
|
||||
* "q3_k_m" : Uses Q4_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
|
||||
* "q3_k_s" : Uses Q3_K for all tensors
|
||||
* "q4_0" : Original quant method, 4-bit.
|
||||
* "q4_1" : Higher accuracy than q4_0 but not as high as q5_0. However has quicker inference than q5 models.
|
||||
* "q4_k_s" : Uses Q4_K for all tensors
|
||||
* "q5_0" : Higher accuracy, higher resource usage and slower inference.
|
||||
* "q5_1" : Even higher accuracy, resource usage and slower inference.
|
||||
* "q5_k_s" : Uses Q5_K for all tensors
|
||||
* "q6_k" : Uses Q8_K for all tensors
|
||||
first_conversion (str, optional): The initial conversion format before quantization.
|
||||
token (str, optional): Hugging Face token. Uses cached token if not provided.
|
||||
private (bool, optional): Whether the repo should be private.
|
||||
commit_message (str): Commit message for the upload.
|
||||
commit_description (str): Commit description for the upload.
|
||||
max_shard_size (str): Maximum shard size for saving.
|
||||
temporary_location (str): Temp directory for intermediate files.
|
||||
maximum_memory_usage (float): Max fraction of memory to use.
|
||||
create_pr (bool): Whether to create a pull request instead of pushing directly.
|
||||
revision (str, optional): Branch/revision to push to.
|
||||
tags (list, optional): Additional tags for the repo.
|
||||
|
||||
Returns:
|
||||
str: The full repo ID on Hugging Face Hub.
|
||||
quantization_method (str or list) selects the GGUF method(s):
|
||||
* "not_quantized" : Fast conversion, slow inference, big files.
|
||||
* "fast_quantized" : Fast conversion, OK inference, OK file size.
|
||||
* "quantized" : Slow conversion, fast inference, small files.
|
||||
* "f32" / "f16" : Full accuracy, slow and memory hungry.
|
||||
* "q8_0" : Fast conversion, high resource use.
|
||||
* "q4_k_m" / "q5_k_m" : Q6_K for half the attention.wv/feed_forward.w2 tensors, else Q4_K/Q5_K.
|
||||
* "q2_k" : Q4_K for attention.vw/feed_forward.w2, Q2_K elsewhere.
|
||||
* "q3_k_l" / "q3_k_m" : Q5_K/Q4_K for attention.wv/wo/feed_forward.w2, else Q3_K.
|
||||
* "q3_k_s" / "q4_k_s" / "q5_k_s" : Q3_K/Q4_K/Q5_K for all tensors.
|
||||
* "q4_0" / "q4_1" / "q5_0" / "q5_1" : 4/5-bit, increasing accuracy and cost.
|
||||
* "q6_k" : Q8_K for all tensors.
|
||||
"""
|
||||
if token is None:
|
||||
token = get_token()
|
||||
|
|
@ -486,7 +451,6 @@ This sentence-transformers model was finetuned and converted to GGUF format usin
|
|||
revision = revision,
|
||||
)
|
||||
|
||||
# Add tags
|
||||
all_tags = ["gguf", "llama-cpp", "unsloth", "sentence-transformers"]
|
||||
if is_vlm:
|
||||
all_tags.append("vision-language-model")
|
||||
|
|
@ -507,10 +471,9 @@ This sentence-transformers model was finetuned and converted to GGUF format usin
|
|||
class FastSentenceTransformer(FastModel):
|
||||
@staticmethod
|
||||
def _save_base_config_for_processor_resume(config, output_path):
|
||||
"""sentence-transformers >= 5.4 reloads Transformer modules via
|
||||
AutoProcessor, which falls back to AutoConfig for tokenizer-only
|
||||
roots -- so PEFT adapter checkpoints still need base config.json
|
||||
next to adapter_config.json."""
|
||||
"""Write base config.json next to adapter_config.json so PEFT adapter
|
||||
checkpoints reload: sentence-transformers >= 5.4 reloads Transformer
|
||||
modules via AutoProcessor, which falls back to AutoConfig."""
|
||||
if config is None or not getattr(config, "model_type", None):
|
||||
return
|
||||
if hasattr(config, "save_pretrained"):
|
||||
|
|
@ -926,7 +889,6 @@ class FastSentenceTransformer(FastModel):
|
|||
with open(readme_path, "r", encoding = "utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
# add unsloth tag to frontmatter
|
||||
if "---\ntags:\n" in content:
|
||||
content = content.replace("---\ntags:\n", "---\ntags:\n- unsloth\n")
|
||||
else:
|
||||
|
|
@ -1078,7 +1040,6 @@ class FastSentenceTransformer(FastModel):
|
|||
}
|
||||
transformer_module.model_forward_params |= preinit_model_forward_params
|
||||
|
||||
# determine max_seq_length if not provided
|
||||
if max_seq_length is None:
|
||||
if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"):
|
||||
max_seq_length = model.config.max_position_embeddings
|
||||
|
|
@ -1137,10 +1098,7 @@ class FastSentenceTransformer(FastModel):
|
|||
trust_remote_code = False,
|
||||
) -> tuple[OrderedDict, bool]:
|
||||
"""Load modules from modules.json, else fall back to hard-coded modules.
|
||||
|
||||
Returns:
|
||||
tuple[OrderedDict, bool]: (modules, no_modules_json)
|
||||
"""
|
||||
Returns ``(modules, no_modules_json)``."""
|
||||
from sentence_transformers.util import import_from_string, load_dir_path
|
||||
from sentence_transformers.models import Pooling, Normalize
|
||||
|
||||
|
|
@ -1225,11 +1183,8 @@ class FastSentenceTransformer(FastModel):
|
|||
max_seq_length = None,
|
||||
):
|
||||
"""Estimate the minimum training steps for torch.compile to pay off
|
||||
(with a 1.2x safety margin), from empirical benchmarks.
|
||||
|
||||
Optional batch_size / grad_accum / max_seq_length give a coarse,
|
||||
conservative pre-run adjustment with no runtime measurements.
|
||||
"""
|
||||
(1.2x safety margin) from empirical benchmarks. Optional batch_size /
|
||||
grad_accum / max_seq_length give a coarse pre-run adjustment."""
|
||||
if hasattr(model, "__getitem__"):
|
||||
try:
|
||||
inner = model[0].auto_model
|
||||
|
|
@ -1474,14 +1429,12 @@ class FastSentenceTransformer(FastModel):
|
|||
print("Unsloth: Device does not support bfloat16. Using float16 instead.")
|
||||
dtype = torch.float16
|
||||
|
||||
# Determine device
|
||||
st_device = device_map
|
||||
if isinstance(st_device, dict) or (
|
||||
isinstance(st_device, str) and st_device in ["auto", "sequential"]
|
||||
):
|
||||
st_device = "cuda"
|
||||
|
||||
# Build model_kwargs for SentenceTransformer
|
||||
model_kwargs = {"torch_dtype": dtype}
|
||||
|
||||
encoder_attn_impl = resolve_encoder_attention_implementation(
|
||||
|
|
@ -1494,7 +1447,6 @@ class FastSentenceTransformer(FastModel):
|
|||
if encoder_attn_impl is not None:
|
||||
model_kwargs["attn_implementation"] = encoder_attn_impl
|
||||
|
||||
# Print optimization status
|
||||
sdpa_str = " + SDPA" if supports_sdpa else ""
|
||||
if load_in_4bit:
|
||||
print(
|
||||
|
|
@ -1505,7 +1457,6 @@ class FastSentenceTransformer(FastModel):
|
|||
f"Unsloth: Using fast encoder path for {model_type} (torch.compile{sdpa_str})"
|
||||
)
|
||||
|
||||
# Handle 4-bit quantization via BitsAndBytesConfig
|
||||
if load_in_4bit:
|
||||
from transformers import BitsAndBytesConfig
|
||||
|
||||
|
|
@ -1519,12 +1470,12 @@ class FastSentenceTransformer(FastModel):
|
|||
# When using quantization, device must be handled by accelerate
|
||||
st_device = None
|
||||
|
||||
# Handle gradient checkpointing - warn user it conflicts with torch.compile
|
||||
# Gradient checkpointing conflicts with torch.compile
|
||||
_use_gc = use_gradient_checkpointing
|
||||
if _use_gc and _use_gc != False:
|
||||
print("Unsloth Warning: Gradient checkpointing is incompatible with torch.compile.")
|
||||
print("Disabling torch.compile to enable gradient checkpointing.")
|
||||
compile_mode = None # Disable compilation
|
||||
compile_mode = None
|
||||
|
||||
is_mpnet = "mpnet" == model_type.lower()
|
||||
|
||||
|
|
@ -1553,7 +1504,6 @@ class FastSentenceTransformer(FastModel):
|
|||
st_model[0], getattr(st_model[0].auto_model, "config", None)
|
||||
)
|
||||
|
||||
# Add save methods
|
||||
def _save_pretrained_merged(self, save_directory, **save_kwargs):
|
||||
self.save_pretrained(save_directory)
|
||||
tokenizer = save_kwargs.pop("tokenizer", self.tokenizer)
|
||||
|
|
@ -1612,7 +1562,6 @@ class FastSentenceTransformer(FastModel):
|
|||
print("Unsloth Warning: 4-bit quantization adds ~2.3x overhead for encoder models.")
|
||||
print("Consider using load_in_16bit=True for better performance.")
|
||||
|
||||
# check if the model supports add_pooling_layer
|
||||
if "add_pooling_layer" not in kwargs:
|
||||
supported = FastSentenceTransformer._has_add_pooling_layer(
|
||||
config, kwargs.get("auto_model", AutoModel)
|
||||
|
|
@ -1749,7 +1698,6 @@ class FastSentenceTransformer(FastModel):
|
|||
save_directory, tokenizer = tokenizer, **kwargs
|
||||
)
|
||||
|
||||
# add Unsloth branding to the generated README
|
||||
try:
|
||||
FastSentenceTransformer._add_unsloth_branding(save_directory)
|
||||
except Exception as e:
|
||||
|
|
@ -1840,7 +1788,6 @@ class FastSentenceTransformer(FastModel):
|
|||
transformer_module = model[0]
|
||||
inner_model = transformer_module.auto_model
|
||||
|
||||
# Check if model is quantized (4-bit/8-bit)
|
||||
is_quantized = (
|
||||
getattr(inner_model, "is_quantized", False)
|
||||
or getattr(inner_model.config, "quantization_config", None) is not None
|
||||
|
|
@ -1863,7 +1810,6 @@ class FastSentenceTransformer(FastModel):
|
|||
elif model_type == "mpnet":
|
||||
FastSentenceTransformer._patch_mpnet_v5()
|
||||
|
||||
# Prepare for k-bit training if quantized
|
||||
if is_quantized:
|
||||
from ._utils import prepare_model_for_kbit_training
|
||||
_gc_for_kbit = (
|
||||
|
|
@ -1878,7 +1824,6 @@ class FastSentenceTransformer(FastModel):
|
|||
gc_enabled = bool(_gc_for_kbit)
|
||||
except ValueError as e:
|
||||
if "does not support gradient checkpointing" in str(e):
|
||||
# Model doesn't support gradient checkpointing, disable it
|
||||
print(
|
||||
f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping."
|
||||
)
|
||||
|
|
@ -1905,7 +1850,6 @@ class FastSentenceTransformer(FastModel):
|
|||
f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping."
|
||||
)
|
||||
|
||||
# Create LoRA config
|
||||
lora_config = LoraConfig(
|
||||
r = r,
|
||||
lora_alpha = lora_alpha,
|
||||
|
|
@ -1918,7 +1862,6 @@ class FastSentenceTransformer(FastModel):
|
|||
# Apply PEFT directly (not through FastModel)
|
||||
peft_model = peft_get_peft_model(inner_model, lora_config)
|
||||
|
||||
# Apply QAT if specified
|
||||
qat_scheme = kwargs.get("qat_scheme", None)
|
||||
if qat_scheme is not None:
|
||||
from ._utils import _prepare_model_for_qat
|
||||
|
|
@ -1951,7 +1894,6 @@ class FastSentenceTransformer(FastModel):
|
|||
model._compile_threshold = FastSentenceTransformer._estimate_compile_threshold(
|
||||
model
|
||||
)
|
||||
# Flag to indicate compile has not been applied yet
|
||||
model._compile_pending = True
|
||||
print(
|
||||
f"Unsloth: torch.compile will be applied automatically if max_steps > {model._compile_threshold}"
|
||||
|
|
@ -2019,12 +1961,8 @@ class FastSentenceTransformer(FastModel):
|
|||
|
||||
|
||||
def _patch_sentence_transformer_trainer():
|
||||
"""
|
||||
Patch SentenceTransformerTrainer to automatically apply torch.compile
|
||||
when training steps exceed the breakeven threshold.
|
||||
|
||||
This is called automatically when this module is imported.
|
||||
"""
|
||||
"""Patch SentenceTransformerTrainer to auto-apply torch.compile when training
|
||||
steps exceed the breakeven threshold. Called on module import."""
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformerTrainer
|
||||
except ImportError:
|
||||
|
|
@ -2043,7 +1981,6 @@ def _patch_sentence_transformer_trainer():
|
|||
model = kwargs.get("model") or (args[0] if args else None)
|
||||
training_args = kwargs.get("args") or (args[1] if len(args) > 1 else None)
|
||||
|
||||
# Check if model has pending compile
|
||||
if (
|
||||
model is not None
|
||||
and training_args is not None
|
||||
|
|
@ -2085,7 +2022,6 @@ def _patch_sentence_transformer_trainer():
|
|||
)
|
||||
model._compile_pending = False
|
||||
|
||||
# Call original __init__
|
||||
_original_init(self, *args, **kwargs)
|
||||
|
||||
# Disable mixed precision when FORCE_FLOAT32 is active (matches rl.py behavior)
|
||||
|
|
|
|||
|
|
@ -348,7 +348,6 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
|
|||
|
||||
kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id)
|
||||
|
||||
# Get pixel values for VLMs
|
||||
try:
|
||||
kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype)
|
||||
except:
|
||||
|
|
@ -358,16 +357,14 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
|
|||
except:
|
||||
pass
|
||||
|
||||
# Mixed precision autocast
|
||||
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
|
||||
autocaster = torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = torch.float16)
|
||||
dtype = torch.float16
|
||||
else:
|
||||
autocaster = torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype)
|
||||
# Prepare LoRA
|
||||
# state_dict = convert_lora_modules(self, dtype = dtype)
|
||||
|
||||
# Set compile dynamic shapes
|
||||
# compile dynamic shapes
|
||||
torch._dynamo.mark_static(input_ids, 0)
|
||||
torch._dynamo.mark_dynamic(input_ids, 1)
|
||||
if "attention_mask" in kwargs:
|
||||
|
|
@ -377,8 +374,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
|
|||
torch._dynamo.mark_static(kwargs["token_type_ids"], 0)
|
||||
torch._dynamo.mark_dynamic(kwargs["token_type_ids"], 1)
|
||||
|
||||
# Fix generation_config
|
||||
# Use hybrid if sliding window seen, otherwise try static
|
||||
# use hybrid cache if sliding window seen, otherwise try static
|
||||
cache_implementation = getattr(self.config, "cache_implementation", None)
|
||||
if getattr(self, "_supports_static_cache", getattr(self, "_can_compile_fullgraph", True)):
|
||||
if os.environ.get("UNSLOTH_DISABLE_STATIC_GENERATION", "0") == "0":
|
||||
|
|
@ -386,7 +382,6 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
|
|||
elif Version(transformers_version) < Version("4.56.0.dev0"):
|
||||
cache_implementation = None
|
||||
else:
|
||||
# Should work in latest transformers!
|
||||
cache_implementation = "static"
|
||||
else:
|
||||
cache_implementation = None
|
||||
|
|
@ -450,18 +445,16 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
|
|||
|
||||
|
||||
def _construct_vlm_processor_fallback(tokenizer_name, model_type, token, trust_remote_code):
|
||||
"""Construct a VLM processor manually when AutoProcessor.from_pretrained fails.
|
||||
"""Build a VLM processor manually when AutoProcessor.from_pretrained fails.
|
||||
|
||||
Some VLMs (e.g., LFM2.5-VL) have tokenizer_class entries that AutoTokenizer
|
||||
cannot resolve. This function loads the image processor and tokenizer separately,
|
||||
sets required special token attributes, and constructs the processor.
|
||||
Some VLMs (e.g. LFM2.5-VL) have tokenizer_class entries AutoTokenizer cannot
|
||||
resolve; load the image processor and tokenizer separately and assemble them.
|
||||
"""
|
||||
try:
|
||||
from transformers import AutoImageProcessor, PreTrainedTokenizerFast, AutoConfig
|
||||
from transformers.models.auto.processing_auto import PROCESSOR_MAPPING_NAMES
|
||||
import json
|
||||
|
||||
# Load image processor
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
tokenizer_name,
|
||||
token = token,
|
||||
|
|
@ -481,7 +474,6 @@ def _construct_vlm_processor_fallback(tokenizer_name, model_type, token, trust_r
|
|||
config_path = hf_hub_download(tokenizer_name, "tokenizer_config.json", token = token)
|
||||
with open(config_path, "r", encoding = "utf-8") as f:
|
||||
tok_config = json.load(f)
|
||||
# Set model-specific special tokens and their IDs
|
||||
for key in (
|
||||
"image_token",
|
||||
"image_start_token",
|
||||
|
|
@ -501,8 +493,7 @@ def _construct_vlm_processor_fallback(tokenizer_name, model_type, token, trust_r
|
|||
# Find the processor class - try model_type first, then top-level config model_type
|
||||
proc_class_name = PROCESSOR_MAPPING_NAMES.get(model_type)
|
||||
if proc_class_name is None:
|
||||
# model_type might be a sub-model type (e.g. "lfm2" instead of "lfm2_vl").
|
||||
# Try the top-level config.model_type which often has the processor mapping.
|
||||
# model_type may be a sub-type (e.g. "lfm2" vs "lfm2_vl"); top-level config often maps
|
||||
try:
|
||||
config = AutoConfig.from_pretrained(
|
||||
tokenizer_name,
|
||||
|
|
@ -609,8 +600,8 @@ class FastBaseModel:
|
|||
if os.environ.get("UNSLOTH_MODEL_NAME", "") == "":
|
||||
os.environ["UNSLOTH_MODEL_NAME"] = model_name.lower()
|
||||
|
||||
# Resolve text-only before the is_vlm / vLLM checks so is_vlm stays consistent;
|
||||
# skip the vision tower only for families with their own text decoder (Gemma 3). #5816
|
||||
# Resolve text-only before is_vlm/vLLM checks; skip vision tower only for
|
||||
# families with their own text decoder (Gemma 3). #5816
|
||||
if text_only and auto_config is None:
|
||||
auto_config = AutoConfig.from_pretrained(
|
||||
model_name,
|
||||
|
|
@ -731,7 +722,7 @@ class FastBaseModel:
|
|||
if old_hf_transfer != "0":
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
||||
|
||||
# For debugging - we use a download counter to see if environments are not breaking or if HF is down
|
||||
# download counter: detects broken environments or HF outages
|
||||
get_statistics(kwargs.get("local_files_only", False))
|
||||
|
||||
if dtype is None:
|
||||
|
|
@ -775,7 +766,6 @@ class FastBaseModel:
|
|||
bnb_compute_dtype = eval(_bnb_compute_dtype)
|
||||
correct_dtype = bnb_compute_dtype
|
||||
custom_datatype = _custom_datatype
|
||||
# Execute code as well
|
||||
if len(execute_code.strip()) != 0:
|
||||
exec(execute_code)
|
||||
else:
|
||||
|
|
@ -796,8 +786,7 @@ class FastBaseModel:
|
|||
supports_sdpa = supports_sdpa,
|
||||
)
|
||||
|
||||
# Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with
|
||||
# FP8 weights. We just need to update it here for sanity.
|
||||
# FP8 models were already redirected to a BF16 sibling; sync model_name here
|
||||
auto_config.model_name = model_name
|
||||
kwargs["attn_implementation"] = attn_impl
|
||||
|
||||
|
|
@ -824,9 +813,8 @@ class FastBaseModel:
|
|||
"Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!"
|
||||
)
|
||||
_skip_modules = SKIP_QUANTIZATION_MODULES.copy()
|
||||
# Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers.
|
||||
# Mamba fused kernels pass out_proj.weight directly to F.linear,
|
||||
# which fails with quantized Params4bit. Skip out_proj from quantization.
|
||||
# Nemotron-H Mamba fused kernels pass out_proj.weight to F.linear,
|
||||
# which fails on quantized Params4bit; skip out_proj from quantization.
|
||||
if any(mt == "nemotron_h" for mt in (model_types or [])):
|
||||
_skip_modules.append("out_proj")
|
||||
|
||||
|
|
@ -910,7 +898,6 @@ class FastBaseModel:
|
|||
quantizer = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_method]
|
||||
quantizer_kwargs = {}
|
||||
if quant_method == "compressed-tensors":
|
||||
# Ignore these
|
||||
pass
|
||||
else:
|
||||
# We cannot dequantize since gpt-oss-20b MXFP4 will now be gpt-oss-20b-BF16
|
||||
|
|
@ -951,8 +938,7 @@ class FastBaseModel:
|
|||
if not fast_inference:
|
||||
# Prevent load_in_fp8 from being forwarded into HF internal model loading
|
||||
load_in_fp8 = kwargs.pop("load_in_fp8", None)
|
||||
# Transformers 5.x @strict config classes reject unexpected kwargs.
|
||||
# Move config-level attributes onto the config object directly.
|
||||
# Transformers 5.x @strict config classes reject unexpected kwargs; set them on config
|
||||
_num_labels = kwargs.pop("num_labels", None)
|
||||
if _num_labels is not None:
|
||||
set_task_config_attr(model_config, "num_labels", _num_labels)
|
||||
|
|
@ -1067,10 +1053,9 @@ class FastBaseModel:
|
|||
if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs:
|
||||
load_vllm_kwargs[allowed_arg] = kwargs[allowed_arg]
|
||||
|
||||
# Load vLLM first
|
||||
llm = load_vllm(**load_vllm_kwargs)
|
||||
|
||||
# Convert to HF format
|
||||
# convert to HF format
|
||||
_, quant_state_dict = get_vllm_state_dict(
|
||||
llm,
|
||||
config = model_config,
|
||||
|
|
@ -1091,10 +1076,9 @@ class FastBaseModel:
|
|||
|
||||
raise_handler.remove()
|
||||
|
||||
# Return old flag
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
|
||||
|
||||
# Check float32 norm weights
|
||||
# Force float32 norm weights when requested
|
||||
if os.environ.get("UNSLOTH_HIGH_PRECISION_LAYERNORM", "0") == "1":
|
||||
for jj, (name, module) in enumerate(model.named_modules()):
|
||||
if (
|
||||
|
|
@ -1103,12 +1087,10 @@ class FastBaseModel:
|
|||
or "layer_norm" in name
|
||||
) and hasattr(module, "weight"):
|
||||
module._pre_set_compute_dtype = torch.float32
|
||||
# Edit data-types
|
||||
if custom_datatype is not None:
|
||||
with torch.no_grad():
|
||||
for jj, (name, module) in enumerate(model.named_modules()):
|
||||
exec(custom_datatype)
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
if DEVICE_TYPE in ("cuda", "hip"):
|
||||
|
|
@ -1116,7 +1098,6 @@ class FastBaseModel:
|
|||
elif DEVICE_TYPE == "xpu":
|
||||
torch.xpu.empty_cache()
|
||||
|
||||
# Counteract saved tokenizers
|
||||
tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
|
||||
|
||||
# Fix _Unsloth_Patched_ prefix in local config files from old saves (issue #4085)
|
||||
|
|
@ -1171,10 +1152,8 @@ class FastBaseModel:
|
|||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
||||
# If processor loading failed (e.g., tokenizer class not found),
|
||||
# or if AutoProcessor silently degraded to a text-only tokenizer
|
||||
# instead of returning a full VLM processor (issue #4085),
|
||||
# try constructing the processor manually from separate components.
|
||||
# If processor loading failed or AutoProcessor degraded to a text-only
|
||||
# tokenizer instead of a full VLM processor (#4085), build it manually.
|
||||
_processor_is_degraded = (
|
||||
is_vlm and tokenizer is not None and not hasattr(tokenizer, "image_processor")
|
||||
)
|
||||
|
|
@ -1193,8 +1172,7 @@ class FastBaseModel:
|
|||
f"Unsloth: Warning - VLM processor fallback returned None for model_type={model_type_arch}",
|
||||
file = sys.stderr,
|
||||
)
|
||||
# Backwards compat: if processor has no chat_template (e.g. old saves without
|
||||
# chat_template.jinja) but the inner tokenizer does, copy it to the processor.
|
||||
# Backwards compat: copy chat_template from inner tokenizer when processor lacks one
|
||||
if (
|
||||
hasattr(tokenizer, "tokenizer")
|
||||
and getattr(tokenizer, "chat_template", None) is None
|
||||
|
|
@ -1204,9 +1182,7 @@ class FastBaseModel:
|
|||
|
||||
if hasattr(tokenizer, "tokenizer"):
|
||||
__tokenizer = tokenizer.tokenizer
|
||||
# Add padding side as well
|
||||
__tokenizer.padding_side = "left"
|
||||
# Check bos, eos, pad tokens
|
||||
if hasattr(__tokenizer, "bos_token"):
|
||||
tokenizer.bos_token = __tokenizer.bos_token
|
||||
tokenizer.bos_token_id = __tokenizer.bos_token_id
|
||||
|
|
@ -1216,7 +1192,6 @@ class FastBaseModel:
|
|||
if hasattr(__tokenizer, "pad_token"):
|
||||
tokenizer.pad_token = __tokenizer.pad_token
|
||||
tokenizer.pad_token_id = __tokenizer.pad_token_id
|
||||
# Fix other stuff like BnB compute data types
|
||||
model, tokenizer = patch_model_and_tokenizer(
|
||||
model,
|
||||
tokenizer,
|
||||
|
|
@ -1229,8 +1204,7 @@ class FastBaseModel:
|
|||
try:
|
||||
model, tokenizer = patch_tokenizer(model, tokenizer)
|
||||
except Exception as _patch_err:
|
||||
# Some VLM processors (e.g., ERNIE VL) may fail during tokenizer patching.
|
||||
# Try loading tokenizer separately via AutoTokenizer as fallback.
|
||||
# Some VLM processors (e.g. ERNIE VL) fail patching; retry via AutoTokenizer
|
||||
try:
|
||||
from transformers import AutoTokenizer as _AutoTokenizer
|
||||
|
||||
|
|
@ -1287,13 +1261,12 @@ class FastBaseModel:
|
|||
apply_accepts_loss_kwargs_fix(model)
|
||||
patch_gradient_accumulation_fix(Trainer)
|
||||
|
||||
# Save tokenizer for inference purposes
|
||||
tokenizer.padding_side = "left" # Force inference
|
||||
if hasattr(tokenizer, "tokenizer"):
|
||||
tokenizer.tokenizer.padding_side = "left" # Force inference
|
||||
# Audio feature extractors must stay right padded: left (a text setting,
|
||||
# forwarded by from_pretrained) shifts Whisper mels and desyncs Gemma 4
|
||||
# audio token counts (crash on transformers < 5.10).
|
||||
# Audio feature extractors must stay right padded: left padding (a text
|
||||
# setting forwarded by from_pretrained) shifts Whisper mels and desyncs
|
||||
# Gemma 4 audio token counts (crash on transformers < 5.10).
|
||||
feature_extractor = getattr(tokenizer, "feature_extractor", None)
|
||||
if (
|
||||
feature_extractor is not None
|
||||
|
|
@ -1308,14 +1281,12 @@ class FastBaseModel:
|
|||
m.is_loaded_in_8bit = True if not full_finetuning else False
|
||||
m = m.model
|
||||
m.max_seq_length = max_seq_length
|
||||
# Save to modules as well
|
||||
for module in model.modules():
|
||||
module.max_seq_length = max_seq_length
|
||||
m._saved_temp_tokenizer = tokenizer
|
||||
# Also set is_loaded_in_8bit to disable incorrect DDP
|
||||
m.is_loaded_in_8bit = True if not full_finetuning else False
|
||||
|
||||
# Patch generate
|
||||
if os.environ.get("UNSLOTH_DISABLE_FAST_GENERATION", "0") == "0" and hasattr(
|
||||
model, "generate"
|
||||
):
|
||||
|
|
@ -1324,7 +1295,6 @@ class FastBaseModel:
|
|||
unsloth_base_fast_generate.__doc__ = model._old_generate.__doc__
|
||||
model.generate = types.MethodType(unsloth_base_fast_generate, model)
|
||||
model._unsloth_trust_remote_code = trust_remote_code
|
||||
# Post patches
|
||||
model = FastBaseModel.post_patch_model(
|
||||
model,
|
||||
use_gradient_checkpointing = use_gradient_checkpointing,
|
||||
|
|
@ -1333,7 +1303,6 @@ class FastBaseModel:
|
|||
tokenizer = tokenizer,
|
||||
float32_mixed_precision = float32_mixed_precision,
|
||||
)
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
if DEVICE_TYPE in ("cuda", "hip"):
|
||||
|
|
@ -1425,7 +1394,7 @@ class FastBaseModel:
|
|||
and hasattr(model.vllm_engine.llm_engine, "vllm_config")
|
||||
and getattr(model.vllm_engine.llm_engine.vllm_config, "lora_config", None) is None
|
||||
):
|
||||
# If vLLM is being used but lora is not enabled, throw an error
|
||||
# vLLM in use but LoRA not enabled
|
||||
# Ref https://github.com/vllm-project/vllm/blob/51ba839555a5d122eadd91e9c16463ac288f5fa1/vllm/v1/engine/processor.py#L148-L151
|
||||
raise RuntimeError("Unsloth: LoRA is not enabled for this model!")
|
||||
if finetune_vision_layers:
|
||||
|
|
@ -1444,7 +1413,6 @@ class FastBaseModel:
|
|||
"Unsloth: LoRA finetuning for Llama 3.2 aka mllama models is not supported with fast_inference!"
|
||||
)
|
||||
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
if DEVICE_TYPE in ("cuda", "hip"):
|
||||
|
|
@ -1481,8 +1449,8 @@ class FastBaseModel:
|
|||
model,
|
||||
use_gradient_checkpointing = use_gradient_checkpointing,
|
||||
)
|
||||
# Gemma4 ClippableLinear wraps nn.Linear -- PEFT can't inject LoRA on it directly.
|
||||
# Monkey-patch PEFT to target the inner .linear child instead.
|
||||
# Gemma4 ClippableLinear wraps nn.Linear; PEFT can't inject LoRA directly,
|
||||
# so patch it to target the inner .linear child instead.
|
||||
_clippable_linear_cls = None
|
||||
try:
|
||||
from transformers.models.gemma4.modeling_gemma4 import (
|
||||
|
|
@ -1549,10 +1517,8 @@ class FastBaseModel:
|
|||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
model.max_seq_length = max_seq_length
|
||||
# Save to modules as well
|
||||
for module in model.modules():
|
||||
module.max_seq_length = max_seq_length
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
if DEVICE_TYPE in ("cuda", "hip"):
|
||||
|
|
@ -1562,7 +1528,6 @@ class FastBaseModel:
|
|||
patch_saving_functions(model, vision = True)
|
||||
patch_peft_fast_inference(model)
|
||||
|
||||
# Add for_inference and for_training
|
||||
model.for_training = functools.partial(FastBaseModel.for_training, model)
|
||||
model.for_inference = functools.partial(FastBaseModel.for_inference, model)
|
||||
m = model
|
||||
|
|
@ -1621,11 +1586,9 @@ class FastBaseModel:
|
|||
patch_modules_to_save = True,
|
||||
)
|
||||
|
||||
# Gemma3N audio conformer processes variable-length audio tensors
|
||||
# that cause stride mismatches in AOT autograd compiled backward
|
||||
# when non-reentrant checkpointing is used. The notebook or TRL
|
||||
# may override gradient_checkpointing_kwargs with use_reentrant=False
|
||||
# after this point, so we intercept gradient_checkpointing_enable
|
||||
# Gemma3N audio conformer's variable-length tensors cause stride mismatches
|
||||
# in AOT autograd compiled backward under non-reentrant checkpointing. TRL/notebook
|
||||
# may later set use_reentrant=False, so intercept gradient_checkpointing_enable
|
||||
# to always force use_reentrant=True for Gemma3N.
|
||||
_model_type = getattr(getattr(model, "config", None), "model_type", "") or ""
|
||||
if "gemma3n" in _model_type.lower() or "gemma4" in _model_type.lower():
|
||||
|
|
@ -1663,14 +1626,12 @@ class FastBaseModel:
|
|||
# Also set is_loaded_in_8bit to disable incorrect DDP
|
||||
m.is_loaded_in_8bit = True if not full_finetuning else False
|
||||
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
if DEVICE_TYPE in ("cuda", "hip"):
|
||||
torch.cuda.empty_cache()
|
||||
elif DEVICE_TYPE == "xpu":
|
||||
torch.xpu.empty_cache()
|
||||
# Add for_inference and for_training
|
||||
model.for_training = functools.partial(FastBaseModel.for_training, model)
|
||||
model.for_inference = functools.partial(FastBaseModel.for_inference, model)
|
||||
m = model
|
||||
|
|
@ -1745,8 +1706,7 @@ class FastBaseModel:
|
|||
embeddings = model.get_output_embeddings()
|
||||
if hasattr(embeddings, "training"):
|
||||
embeddings.training = False
|
||||
# Restore use_cache values that prepare_model_for_training disabled
|
||||
# for gradient checkpointing (older unsloth_zoo has no restore helper)
|
||||
# Restore use_cache that prepare_model_for_training disabled for gradient checkpointing
|
||||
try:
|
||||
from unsloth_zoo.training_utils import restore_use_cache
|
||||
restore_use_cache(model)
|
||||
|
|
@ -1811,8 +1771,7 @@ class FastBaseModel:
|
|||
embeddings = model.get_output_embeddings()
|
||||
if hasattr(embeddings, "training"):
|
||||
embeddings.training = True
|
||||
# Re-disable use_cache if prepare_model_for_training had disabled it
|
||||
# and for_inference restored it (record only exists after a disable)
|
||||
# Re-disable use_cache if for_inference restored it (record exists only after a disable)
|
||||
if (
|
||||
use_gradient_checkpointing
|
||||
and getattr(model, "_unsloth_use_cache_originals", None) is not None
|
||||
|
|
@ -1877,21 +1836,17 @@ def check_dataset_for_missing_videos(
|
|||
checked = None,
|
||||
):
|
||||
"""
|
||||
Validate that local video paths referenced in a dataset exist, catching
|
||||
missing files before training (torchvision otherwise returns an empty
|
||||
tensor and the model silently receives no video signal).
|
||||
Validate local video paths in a dataset exist, catching missing files before
|
||||
training (torchvision otherwise silently yields an empty tensor). Returns the
|
||||
list of missing paths (empty when all exist).
|
||||
|
||||
Args:
|
||||
dataset: Map-style Dataset, list of dicts, or iterable of examples
|
||||
(not a streaming IterableDataset - iterating consumes it).
|
||||
column: Chat-messages column, default "messages"; "conversations",
|
||||
"prompt" and "completion" are also scanned.
|
||||
raise_error: True (default) raises FileNotFoundError listing missing
|
||||
files; False warns and returns them.
|
||||
dataset: Map-style Dataset / list / iterable (not a streaming
|
||||
IterableDataset - iterating consumes it).
|
||||
column: Chat-messages column ("messages"); "conversations", "prompt"
|
||||
and "completion" are also scanned.
|
||||
raise_error: True raises FileNotFoundError on missing files; False warns.
|
||||
checked: Optional set of known-good paths for cross-call dedup.
|
||||
|
||||
Returns:
|
||||
List[str]: Missing file paths (empty when all exist).
|
||||
"""
|
||||
try:
|
||||
from datasets import IterableDataset as _IterableDataset
|
||||
|
|
@ -1908,8 +1863,8 @@ def check_dataset_for_missing_videos(
|
|||
pass
|
||||
|
||||
missing = []
|
||||
# Report each missing path once; only confirmed-existing paths enter
|
||||
# `checked`, so retries after an error re-check previously missing files.
|
||||
# Report each missing path once; only existing paths enter `checked`, so
|
||||
# retries after an error re-check previously missing files.
|
||||
seen_missing = set()
|
||||
if checked is None:
|
||||
checked = set()
|
||||
|
|
|
|||
|
|
@ -541,7 +541,7 @@ PARAMETER min_p 0.1
|
|||
OLLAMA_TEMPLATES["gemma_chatml"] = gemma_chatml_ollama
|
||||
|
||||
# =========================================== Gemma 2
|
||||
# Same as Gemma 1, but with sliding window attention!
|
||||
# Gemma 1 plus sliding window attention
|
||||
# https://ollama.com/library/gemma2/blobs/6522ca797f47
|
||||
gemma2_ollama = gemma_ollama + "PARAMETER num_ctx 4096\n"
|
||||
OLLAMA_TEMPLATES["gemma2"] = gemma2_ollama
|
||||
|
|
@ -2219,7 +2219,6 @@ for key, values in OLLAMA_TEMPLATE_TO_MODEL_MAPPER.items():
|
|||
for value in values:
|
||||
MODEL_TO_OLLAMA_TEMPLATE_MAPPER[value] = key
|
||||
|
||||
# Get lowercased
|
||||
lowered_key = key.lower()
|
||||
for value in values:
|
||||
MODEL_TO_OLLAMA_TEMPLATE_MAPPER[value.lower()] = lowered_key
|
||||
|
|
|
|||
|
|
@ -46,29 +46,19 @@ def _require_bnb():
|
|||
|
||||
|
||||
class QGaLoreAdamW8bit(Optimizer2State):
|
||||
"""AdamW optimizer with 8-bit states, GaLore low-rank gradient projection,
|
||||
and optional INT8 weight quantization.
|
||||
"""AdamW with 8-bit states, GaLore low-rank gradient projection, and optional
|
||||
INT8 weight quantization. Three memory-saving techniques:
|
||||
|
||||
This optimizer combines three memory-saving techniques:
|
||||
1. **8-bit optimizer states** (bitsandbytes): Adam moments in 8-bit (~4x less).
|
||||
2. **GaLore low-rank projection**: gradients projected to a low-rank subspace
|
||||
for the step, then back; the projection matrix can be INT4-quantized.
|
||||
3. **INT8 weight quantization**: weights stored in INT8 with stochastic
|
||||
rounding (~2x less) for eligible layers.
|
||||
|
||||
1. **8-bit optimizer states** (via bitsandbytes) — Adam's first and second
|
||||
moments are stored in 8-bit, reducing optimizer state memory by ~4×.
|
||||
|
||||
2. **GaLore low-rank gradient projection** — gradients are projected into a
|
||||
low-rank subspace before the optimizer step, then projected back. The
|
||||
projection matrix itself can be quantized to INT4.
|
||||
|
||||
3. **INT8 weight quantization** — model weights are stored in INT8 during
|
||||
training with stochastic rounding, reducing weight memory by ~2× for
|
||||
eligible layers.
|
||||
|
||||
Param group keys consumed by GaLore projection:
|
||||
``rank``, ``update_proj_gap``, ``scale``, ``proj_type``,
|
||||
``quant`` (projection quantization), ``quant_group_size``,
|
||||
``quant_n_bit``, ``cos_threshold``, ``gamma_proj``, ``queue_size``
|
||||
|
||||
Param group keys for weight quantization:
|
||||
``weight_quant``, ``stochastic_round``, ``weight_group_size``
|
||||
Param group keys: GaLore uses ``rank``, ``update_proj_gap``, ``scale``,
|
||||
``proj_type``, ``quant``, ``quant_group_size``, ``quant_n_bit``,
|
||||
``cos_threshold``, ``gamma_proj``, ``queue_size``; weight quantization uses
|
||||
``weight_quant``, ``stochastic_round``, ``weight_group_size``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -101,17 +91,10 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
|
||||
@torch.no_grad()
|
||||
def step(self, closure = None):
|
||||
"""Perform a single optimization step.
|
||||
|
||||
For each parameter that has a ``rank`` key in its param group, the
|
||||
following sequence is executed:
|
||||
|
||||
1. If ``weight_quant`` is set, dequantize the INT8 weight to float.
|
||||
2. Project the gradient to low-rank via the cached ``GaLoreProjector``.
|
||||
3. Perform the 8-bit Adam update in the low-rank space.
|
||||
4. Project the update back to full rank and add to saved weight.
|
||||
5. If ``weight_quant`` is set, re-quantize the weight to INT8.
|
||||
"""
|
||||
"""Single optimization step. For each ``rank``-group parameter: (1)
|
||||
dequantize INT8 weight if ``weight_quant``; (2) project gradient to
|
||||
low-rank; (3) 8-bit Adam update in low-rank space; (4) project back and
|
||||
add to the saved weight; (5) re-quantize to INT8 if ``weight_quant``."""
|
||||
loss = None
|
||||
if closure is not None:
|
||||
with torch.enable_grad():
|
||||
|
|
@ -133,7 +116,6 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
|
||||
has_weight_quant = self._has_weight_quant(p, group)
|
||||
|
||||
# --- Dequantize weight if INT8 ---
|
||||
if has_weight_quant:
|
||||
if p._q_scales is not None:
|
||||
float_weight = _dequantize(
|
||||
|
|
@ -145,7 +127,6 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
p.data = float_weight
|
||||
# else: first step, weights are still float — skip dequantize
|
||||
|
||||
# --- GaLore projection ---
|
||||
if "rank" in group:
|
||||
if "projector" not in state:
|
||||
state["projector"] = GaLoreProjector(
|
||||
|
|
@ -161,8 +142,7 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
queue_size = group.get("queue_size", 5),
|
||||
)
|
||||
|
||||
# Temporarily disable weight decay for GaLore params
|
||||
# (we apply it manually after project-back)
|
||||
# Disable weight decay here; reapplied manually after project-back.
|
||||
if "weight_decay" in group and group["weight_decay"] > 0:
|
||||
group["_wd_saved"] = group["weight_decay"]
|
||||
group["weight_decay"] = 0
|
||||
|
|
@ -174,16 +154,14 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
p.data = torch.zeros_like(grad, dtype = p.data.dtype, device = p.data.device)
|
||||
p.grad = grad
|
||||
|
||||
# --- 8-bit Adam update ---
|
||||
if "state1" not in state:
|
||||
self.init_state(group, p, gindex, pindex)
|
||||
|
||||
self.prefetch_state(p)
|
||||
self.update_step(group, p, gindex, pindex)
|
||||
|
||||
# --- GaLore project-back ---
|
||||
if "rank" in group:
|
||||
# p.data now holds the weight update in low-rank space
|
||||
# p.data holds the update in low-rank space; project back and add.
|
||||
p.data = p._saved_data.add_(state["projector"].project_back(p.data))
|
||||
|
||||
# Re-apply decoupled weight decay using pre-update weights
|
||||
|
|
@ -197,7 +175,6 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
|
||||
del p._saved_data
|
||||
|
||||
# --- Re-quantize weight to INT8 ---
|
||||
if has_weight_quant:
|
||||
float_data = p.data
|
||||
stochastic = group.get("stochastic_round", True)
|
||||
|
|
@ -208,9 +185,8 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
p._q_scales = scales
|
||||
p._q_zeros = zeros
|
||||
p._q_shape = shape
|
||||
# Scalar placeholder to free float memory; the forward
|
||||
# pre-hook (install_weight_quant_hooks) dequantizes before
|
||||
# the next forward pass.
|
||||
# Scalar placeholder frees float memory; install_weight_quant_hooks
|
||||
# forward pre-hook dequantizes before the next forward pass.
|
||||
p.data = torch.empty(1, dtype = p.data.dtype, device = p.data.device)
|
||||
|
||||
state["step"] += 1
|
||||
|
|
@ -235,13 +211,11 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
group_size: int = 128,
|
||||
stochastic: bool = True,
|
||||
) -> None:
|
||||
"""Tag parameters for INT8 weight quantization.
|
||||
"""Tag eligible parameters with INT8 quantization metadata for ``step()``.
|
||||
|
||||
This marks eligible weights with quantization metadata so that
|
||||
the optimizer knows to quantize/dequantize them during ``step()``.
|
||||
**Weights are NOT converted to uint8 here** — they remain in float
|
||||
so that the first forward/backward pass runs correctly. The actual
|
||||
quantization happens at the end of the first ``step()`` call.
|
||||
**Weights are NOT converted to uint8 here** — they stay float so the first
|
||||
forward/backward runs correctly; actual quantization happens at the end of
|
||||
the first ``step()``.
|
||||
"""
|
||||
weight_quant_params = set()
|
||||
for group in param_groups:
|
||||
|
|
@ -251,9 +225,8 @@ class QGaLoreAdamW8bit(Optimizer2State):
|
|||
|
||||
for name, p in model.named_parameters():
|
||||
if id(p) in weight_quant_params:
|
||||
# Store metadata without converting weights to uint8; the first
|
||||
# step() quantizes after the update. Dummy scales/zeros keep
|
||||
# _has_weight_quant() True on the first step.
|
||||
# Tag only; first step() quantizes after the update. Dummy
|
||||
# scales/zeros keep _has_weight_quant() True on the first step.
|
||||
p._q_scales = None
|
||||
p._q_zeros = None
|
||||
p._q_shape = p.data.shape
|
||||
|
|
@ -288,7 +261,7 @@ def install_weight_quant_hooks(model: torch.nn.Module) -> list:
|
|||
return handles
|
||||
|
||||
|
||||
# Default linear layer names in transformer blocks that should use GaLore.
|
||||
# Default transformer layers that use GaLore.
|
||||
_DEFAULT_GALORE_TARGETS = {
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
|
|
@ -318,33 +291,12 @@ def make_q_galore_param_groups(
|
|||
queue_size: int = 5,
|
||||
target_modules: Optional[List[str]] = None,
|
||||
) -> list:
|
||||
"""Build param groups suitable for :class:`QGaLoreAdamW8bit`.
|
||||
"""Build param groups for :class:`QGaLoreAdamW8bit`, returning
|
||||
``[galore_group, non_galore_group]``.
|
||||
|
||||
Parameters matching ``target_modules`` (or the default set of attention
|
||||
and MLP projection names) are placed in the GaLore group. All other
|
||||
trainable parameters go into the non-GaLore group.
|
||||
|
||||
Args:
|
||||
model: The model whose parameters to partition.
|
||||
lr: Learning rate for all parameter groups.
|
||||
weight_decay: Weight decay coefficient.
|
||||
rank: GaLore projection rank.
|
||||
update_proj_gap: Steps between SVD recomputations.
|
||||
scale: Scaling factor for project-back.
|
||||
proj_quant: Quantize projection matrices.
|
||||
proj_quant_group_size: Group size for projection quantization.
|
||||
proj_quant_n_bit: Bit-width for projection quantization.
|
||||
weight_quant: Enable INT8 weight quantization for GaLore params.
|
||||
stochastic_round: Use stochastic rounding for weight quantization.
|
||||
weight_group_size: Group size for weight quantization.
|
||||
cos_threshold: Cosine similarity threshold for adaptive scheduling.
|
||||
gamma_proj: Multiplier for update_proj_gap when subspace is stable.
|
||||
queue_size: Rolling window size for stability tracking.
|
||||
target_modules: Module name substrings to match for GaLore. If None,
|
||||
uses the default set of attention/MLP projection names.
|
||||
|
||||
Returns:
|
||||
List of two param group dicts: ``[galore_group, non_galore_group]``.
|
||||
Parameters matching ``target_modules`` (or the default attention/MLP
|
||||
projection names) go in the GaLore group; all other trainable params go in
|
||||
the non-GaLore group.
|
||||
"""
|
||||
targets = set(target_modules) if target_modules is not None else _DEFAULT_GALORE_TARGETS
|
||||
|
||||
|
|
@ -355,8 +307,7 @@ def make_q_galore_param_groups(
|
|||
if not param.requires_grad:
|
||||
continue
|
||||
|
||||
# Match target module names; exclude 1-D params (biases, norms) since
|
||||
# GaLoreProjector.project requires 2-D gradients.
|
||||
# Exclude 1-D params (biases, norms): GaLoreProjector.project needs 2-D grads.
|
||||
name_parts = name.split(".")
|
||||
is_galore = param.dim() >= 2 and any(t in name_parts for t in targets)
|
||||
|
||||
|
|
|
|||
|
|
@ -37,18 +37,6 @@ class GaLoreProjector:
|
|||
similarity of consecutive orthogonal vectors exceeds ``cos_threshold``,
|
||||
``update_proj_gap`` is multiplied by ``gamma_proj`` to recompute SVD less
|
||||
often for stabilized layers.
|
||||
|
||||
Args:
|
||||
rank: Target rank for the low-rank projection.
|
||||
update_proj_gap: Number of steps between SVD recomputations.
|
||||
scale: Scaling factor applied when projecting back to full rank.
|
||||
proj_type: Projection type. Only ``'std'`` is supported.
|
||||
quant: Whether to quantize the projection matrix.
|
||||
group_size: Group size for projection matrix quantization.
|
||||
n_bit: Bit-width for projection matrix quantization (4 or 8).
|
||||
cos_threshold: Cosine similarity threshold for adaptive scheduling.
|
||||
gamma_proj: Multiplier for ``update_proj_gap`` on stability detection.
|
||||
queue_size: Number of recent cosine similarities to average.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
|
|
@ -90,12 +78,10 @@ class GaLoreProjector:
|
|||
self.scale = scale
|
||||
self.proj_type = proj_type
|
||||
|
||||
# Quantization settings for the projection matrix
|
||||
self.quant = quant
|
||||
self.quant_group_size = group_size
|
||||
self.quant_n_bit = n_bit
|
||||
|
||||
# Adaptive update scheduling state
|
||||
self.cos_threshold = cos_threshold
|
||||
self.gamma_proj = gamma_proj
|
||||
self.queue_size = queue_size
|
||||
|
|
@ -104,7 +90,6 @@ class GaLoreProjector:
|
|||
self.svd_count = 0
|
||||
self._ortho_float_cache = None
|
||||
|
||||
# Projection matrix state
|
||||
self.ortho_matrix = None
|
||||
self.ortho_matrix_scales = None
|
||||
self.ortho_matrix_zeros = None
|
||||
|
|
@ -115,23 +100,15 @@ class GaLoreProjector:
|
|||
# ------------------------------------------------------------------
|
||||
|
||||
def project(self, full_rank_grad: torch.Tensor, step: int) -> torch.Tensor:
|
||||
"""Project a full-rank gradient into the low-rank subspace.
|
||||
"""Project a full-rank (2-D) gradient into the low-rank subspace.
|
||||
|
||||
The SVD is recomputed every ``update_proj_gap`` steps (subject to
|
||||
adaptive scheduling). Between recomputations the cached orthogonal
|
||||
matrix is reused.
|
||||
|
||||
Args:
|
||||
full_rank_grad: The full-rank gradient tensor (2-D).
|
||||
step: The current optimizer step (0-indexed).
|
||||
|
||||
Returns:
|
||||
The low-rank gradient tensor.
|
||||
SVD is recomputed every ``update_proj_gap`` steps (subject to adaptive
|
||||
scheduling); between recomputations the cached orthogonal matrix is reused.
|
||||
"""
|
||||
assert self.proj_type == "std", "Only proj_type='std' is supported."
|
||||
|
||||
if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
|
||||
# "tall" matrix → right projection (grad @ Q^T)
|
||||
# Tall matrix -> right projection (grad @ Q^T)
|
||||
if self.ortho_matrix is None or step % self.update_proj_gap == 0:
|
||||
float_ortho = self._compute_orthogonal(
|
||||
full_rank_grad,
|
||||
|
|
@ -144,7 +121,7 @@ class GaLoreProjector:
|
|||
self._ortho_float_cache = self._load_ortho()
|
||||
low_rank_grad = torch.matmul(full_rank_grad, self._ortho_float_cache.t())
|
||||
else:
|
||||
# "wide" matrix → left projection (Q^T @ grad)
|
||||
# Wide matrix -> left projection (Q^T @ grad)
|
||||
if self.ortho_matrix is None or step % self.update_proj_gap == 0:
|
||||
float_ortho = self._compute_orthogonal(
|
||||
full_rank_grad,
|
||||
|
|
@ -160,14 +137,7 @@ class GaLoreProjector:
|
|||
return low_rank_grad
|
||||
|
||||
def project_back(self, low_rank_grad: torch.Tensor) -> torch.Tensor:
|
||||
"""Project a low-rank update back to full rank.
|
||||
|
||||
Args:
|
||||
low_rank_grad: The low-rank gradient/update tensor.
|
||||
|
||||
Returns:
|
||||
The full-rank update scaled by ``self.scale``.
|
||||
"""
|
||||
"""Project a low-rank update back to full rank, scaled by ``self.scale``."""
|
||||
float_ortho = self._ortho_float_cache
|
||||
self._ortho_float_cache = None
|
||||
if float_ortho is None:
|
||||
|
|
@ -186,16 +156,9 @@ class GaLoreProjector:
|
|||
|
||||
@staticmethod
|
||||
def _compute_orthogonal(weights: torch.Tensor, rank: int, side: str) -> torch.Tensor:
|
||||
"""Compute the top-``rank`` orthogonal matrix via truncated SVD.
|
||||
|
||||
Args:
|
||||
weights: 2-D tensor (typically the gradient).
|
||||
rank: Number of singular vectors to keep.
|
||||
side: ``'left'`` returns U[:, :rank], ``'right'`` returns Vh[:rank, :].
|
||||
|
||||
Returns:
|
||||
Orthogonal matrix of shape ``(rank, N)`` (right) or ``(M, rank)`` (left).
|
||||
"""
|
||||
"""Top-``rank`` orthogonal matrix of 2-D ``weights`` via truncated SVD.
|
||||
``side='left'`` returns U[:, :rank] shape ``(M, rank)``; ``'right'``
|
||||
returns Vh[:rank, :] shape ``(rank, N)``."""
|
||||
original_dtype = weights.dtype
|
||||
original_device = weights.device
|
||||
|
||||
|
|
@ -318,7 +281,7 @@ def _dequantize(
|
|||
w: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor, original_shape: tuple
|
||||
) -> torch.Tensor:
|
||||
"""Dequantize from uint8 back to float."""
|
||||
# Infer group size: scales has shape (n_groups, 1), so n_groups = scales.shape[0]
|
||||
# Infer group size: scales has shape (n_groups, 1)
|
||||
total = w.numel()
|
||||
n_groups = scales.shape[0] if scales.dim() > 1 else scales.numel()
|
||||
group_size = total // n_groups if n_groups > 0 else total
|
||||
|
|
@ -336,12 +299,9 @@ def _quantize_stochastic(
|
|||
) -> tuple:
|
||||
"""Asymmetric min-max quantization with stochastic rounding.
|
||||
|
||||
Instead of deterministic ``round()``, the rounding direction is chosen
|
||||
probabilistically proportional to the fractional part. This gives an
|
||||
unbiased estimator of the original value in expectation.
|
||||
|
||||
Returns:
|
||||
``(quantized_uint8, scales, zeros, original_shape)``
|
||||
Rounding direction is chosen probabilistically by the fractional part,
|
||||
giving an unbiased estimator in expectation.
|
||||
Returns ``(quantized_uint8, scales, zeros, original_shape)``.
|
||||
"""
|
||||
org_shape = w.shape
|
||||
if q_group_size > 0:
|
||||
|
|
|
|||
|
|
@ -32,11 +32,8 @@ def search_models(
|
|||
quant_types: list[QuantType] = None,
|
||||
search_pattern: str = None,
|
||||
) -> list[ModelInfo]:
|
||||
"""
|
||||
Get model info from the registry. See registry.ModelInfo for more fields.
|
||||
|
||||
search_pattern is matched against the full model path (the HF hub model_id).
|
||||
"""
|
||||
"""Query the registry for ModelInfo. search_pattern matches the full HF
|
||||
hub model_id (model_path)."""
|
||||
if not _ARE_MODELS_REGISTERED:
|
||||
register_models()
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,6 @@ class DeepseekR1ModelInfo(ModelInfo):
|
|||
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||
|
||||
|
||||
# Deepseek V3 Model Meta
|
||||
DeepseekV3Meta = ModelMeta(
|
||||
org = "deepseek-ai",
|
||||
base_name = "DeepSeek",
|
||||
|
|
@ -80,7 +79,6 @@ DeepseekR1DistillLlamaMeta = ModelMeta(
|
|||
quant_types = {"8": [QuantType.UNSLOTH, QuantType.GGUF], "70": [QuantType.GGUF]},
|
||||
)
|
||||
|
||||
# Deepseek R1 Distill Qwen Model Meta
|
||||
DeepseekR1DistillQwenMeta = ModelMeta(
|
||||
org = "deepseek-ai",
|
||||
base_name = "DeepSeek-R1-Distill",
|
||||
|
|
@ -164,7 +162,6 @@ def _list_deepseek_r1_distill_models():
|
|||
for model in models:
|
||||
model_id = model.id
|
||||
model_name = model_id.split("/")[-1]
|
||||
# parse out only the version
|
||||
version = model_name.removeprefix("DeepSeek-R1-Distill-")
|
||||
distill_models.append(version)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ class GemmaModelInfo(ModelInfo):
|
|||
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||
|
||||
|
||||
# Gemma3 Base Model Meta
|
||||
GemmaMeta3Base = ModelMeta(
|
||||
org = "google",
|
||||
base_name = "gemma",
|
||||
|
|
@ -23,7 +22,6 @@ GemmaMeta3Base = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Gemma3 Instruct Model Meta
|
||||
GemmaMeta3Instruct = ModelMeta(
|
||||
org = "google",
|
||||
base_name = "gemma",
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ class LlamaVisionModelInfo(ModelInfo):
|
|||
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||
|
||||
|
||||
# Llama 3.1
|
||||
LlamaMeta_3_1 = ModelMeta(
|
||||
org = "meta-llama",
|
||||
base_name = "Llama",
|
||||
|
|
@ -31,7 +30,6 @@ LlamaMeta_3_1 = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Llama 3.2 Base Models
|
||||
LlamaMeta_3_2_Base = ModelMeta(
|
||||
org = "meta-llama",
|
||||
base_name = "Llama",
|
||||
|
|
@ -43,7 +41,6 @@ LlamaMeta_3_2_Base = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Llama 3.2 Instruction Tuned Models
|
||||
LlamaMeta_3_2_Instruct = ModelMeta(
|
||||
org = "meta-llama",
|
||||
base_name = "Llama",
|
||||
|
|
@ -55,7 +52,6 @@ LlamaMeta_3_2_Instruct = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||
)
|
||||
|
||||
# Llama 3.2 Vision
|
||||
LlamaMeta_3_2_Vision = ModelMeta(
|
||||
org = "meta-llama",
|
||||
base_name = "Llama",
|
||||
|
|
|
|||
|
|
@ -11,25 +11,23 @@ class PhiModelInfo(ModelInfo):
|
|||
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||
|
||||
|
||||
# Phi Model Meta
|
||||
PhiMeta4 = ModelMeta(
|
||||
org = "microsoft",
|
||||
base_name = "phi",
|
||||
instruct_tags = [None],
|
||||
model_version = "4",
|
||||
model_sizes = ["1"], # Assuming only one size
|
||||
model_sizes = ["1"],
|
||||
model_info_cls = PhiModelInfo,
|
||||
is_multimodal = False,
|
||||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Phi Instruct Model Meta
|
||||
PhiInstructMeta4 = ModelMeta(
|
||||
org = "microsoft",
|
||||
base_name = "phi",
|
||||
instruct_tags = ["mini-instruct"],
|
||||
model_version = "4",
|
||||
model_sizes = ["1"], # Assuming only one size
|
||||
model_sizes = ["1"],
|
||||
model_info_cls = PhiModelInfo,
|
||||
is_multimodal = False,
|
||||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||
|
|
|
|||
|
|
@ -33,7 +33,6 @@ class QwenQVQPreviewModelInfo(ModelInfo):
|
|||
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||
|
||||
|
||||
# Qwen2.5 Model Meta
|
||||
Qwen_2_5_Meta = ModelMeta(
|
||||
org = "Qwen",
|
||||
base_name = "Qwen",
|
||||
|
|
@ -45,7 +44,6 @@ Qwen_2_5_Meta = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Qwen2.5 VL Model Meta
|
||||
Qwen_2_5_VLMeta = ModelMeta(
|
||||
org = "Qwen",
|
||||
base_name = "Qwen",
|
||||
|
|
@ -57,7 +55,6 @@ Qwen_2_5_VLMeta = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||
)
|
||||
|
||||
# Qwen QwQ Model Meta
|
||||
QwenQwQMeta = ModelMeta(
|
||||
org = "Qwen",
|
||||
base_name = "QwQ",
|
||||
|
|
@ -69,7 +66,6 @@ QwenQwQMeta = ModelMeta(
|
|||
quant_types = [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||
)
|
||||
|
||||
# Qwen QVQ Preview Model Meta
|
||||
QwenQVQPreviewMeta = ModelMeta(
|
||||
org = "Qwen",
|
||||
base_name = "QVQ",
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ class QuantType(Enum):
|
|||
BF16 = "bf16" # only for Deepseek V3
|
||||
|
||||
|
||||
# Tags for Hugging Face model paths
|
||||
BNB_QUANTIZED_TAG = "bnb-4bit"
|
||||
UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG
|
||||
GGUF_TAG = "GGUF"
|
||||
|
|
@ -159,13 +158,12 @@ def _register_models(model_meta: ModelMeta, include_original_model: bool = False
|
|||
|
||||
for size in model_sizes:
|
||||
for instruct_tag in instruct_tags:
|
||||
# Handle quant types per model size
|
||||
# quant types may vary per model size
|
||||
if isinstance(quant_types, dict):
|
||||
_quant_types = quant_types[size]
|
||||
else:
|
||||
_quant_types = quant_types
|
||||
for quant_type in _quant_types:
|
||||
# NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH
|
||||
_org = "unsloth" # quantized versions of the original model
|
||||
register_model(
|
||||
model_info_cls = model_info_cls,
|
||||
|
|
@ -177,7 +175,7 @@ def _register_models(model_meta: ModelMeta, include_original_model: bool = False
|
|||
quant_type = quant_type,
|
||||
is_multimodal = is_multimodal,
|
||||
)
|
||||
# include original model from releasing organization
|
||||
# original model from the releasing organization
|
||||
if include_original_model:
|
||||
register_model(
|
||||
model_info_cls = model_info_cls,
|
||||
|
|
|
|||
212
unsloth/save.py
212
unsloth/save.py
|
|
@ -24,8 +24,7 @@ from unsloth_zoo.llama_cpp import (
|
|||
_download_convert_hf_to_gguf,
|
||||
)
|
||||
|
||||
# H4: Defensive imports -- these were added in unsloth-zoo PR #526
|
||||
# and may not exist on older versions
|
||||
# Added in unsloth-zoo PR #526; may not exist on older versions
|
||||
try:
|
||||
from unsloth_zoo.llama_cpp import LLAMA_CPP_DEFAULT_DIR, IS_WINDOWS
|
||||
except ImportError:
|
||||
|
|
@ -82,14 +81,12 @@ LLAMA_CPP_TARGETS = [
|
|||
"llama-server",
|
||||
]
|
||||
|
||||
# Check environments
|
||||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames
|
||||
IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames
|
||||
KAGGLE_TMP = "/tmp"
|
||||
del keynames
|
||||
|
||||
# Weights
|
||||
LLAMA_WEIGHTS = (
|
||||
"self_attn.q_proj",
|
||||
"self_attn.k_proj",
|
||||
|
|
@ -347,8 +344,7 @@ def _free_cached_model(model):
|
|||
from huggingface_hub import scan_cache_dir
|
||||
cached_repos = list(scan_cache_dir().repos)
|
||||
|
||||
# Go through every cached repo, and delete the one that matches the model we want to save.
|
||||
# Can save 4GB of disk space - useful for Kaggle systems.
|
||||
# Delete the cached repo matching this model; saves ~4GB on Kaggle.
|
||||
for cached_repo in cached_repos:
|
||||
if cached_repo.repo_id == model.config._name_or_path:
|
||||
remove_cache_commit = list(cached_repo.revisions)[0].commit_hash
|
||||
|
|
@ -367,7 +363,7 @@ def _free_cached_model(model):
|
|||
def _merge_lora(layer, name):
|
||||
bias = getattr(layer, "bias", None)
|
||||
if isinstance(layer, (Bnb_Linear4bit, Peft_Linear4bit, Peft_Linear)):
|
||||
# Is LoRA so we need to merge!
|
||||
# LoRA layer: merge adapters into W
|
||||
W, quant_state, A, B, s, bias = get_lora_parameters_bias(layer)
|
||||
if quant_state is not None:
|
||||
dtype = quant_state.dtype if type(quant_state) is not list else quant_state[2]
|
||||
|
|
@ -412,17 +408,12 @@ def _preserve_tokenizer_eos_token(
|
|||
):
|
||||
"""Restore tokenizer_config.json eos_token from the tokenizer passed to save.
|
||||
|
||||
Some merge paths may re-save or mutate tokenizer metadata after the tokenizer
|
||||
is written. Gemma 4 instruct models use `<turn|>` as their chat EOS token;
|
||||
if tokenizer_config.json is reset to the raw base `<eos>` token, runtimes such
|
||||
as vLLM will not stop generation correctly. Keep the serialized metadata in
|
||||
sync with the source tokenizer without failing the save if the config is not
|
||||
present or cannot be edited.
|
||||
Merge paths may mutate tokenizer metadata after writing. E.g. Gemma 4 instruct
|
||||
uses `<turn|>` as chat EOS; if the config is reset to the base `<eos>`, vLLM
|
||||
won't stop generation correctly. Best-effort, never fails the save.
|
||||
|
||||
`filename_prefix` mirrors the same argument on Transformers'
|
||||
`PreTrainedTokenizerBase.save_pretrained`: when provided, the tokenizer
|
||||
config is written as `{filename_prefix}-tokenizer_config.json` instead of
|
||||
`tokenizer_config.json`.
|
||||
`filename_prefix` mirrors Transformers' save_pretrained: when set, writes
|
||||
`{filename_prefix}-tokenizer_config.json` instead of `tokenizer_config.json`.
|
||||
"""
|
||||
if tokenizer is None or save_directory is None:
|
||||
return
|
||||
|
|
@ -567,7 +558,6 @@ def unsloth_save_model(
|
|||
|
||||
assert maximum_memory_usage > 0 and maximum_memory_usage <= 0.95
|
||||
|
||||
# Clean memory up first
|
||||
for _ in range(3):
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
|
@ -585,7 +575,7 @@ def unsloth_save_model(
|
|||
print("Unsloth: Merging 4bit and LoRA weights to 4bit...")
|
||||
print("This might take 5 minutes...")
|
||||
|
||||
# Counteract no LoRA adapters!
|
||||
# Guard against models without LoRA adapters
|
||||
if hasattr(model, "merge_and_unload"):
|
||||
model = model.merge_and_unload()
|
||||
print("Done.")
|
||||
|
|
@ -613,7 +603,6 @@ def unsloth_save_model(
|
|||
elif save_method == "merged_4bit":
|
||||
print("Unsloth: Saving 4bit Bitsandbytes model. Please wait...")
|
||||
|
||||
# Update model tag
|
||||
_ = upload_to_huggingface(
|
||||
model,
|
||||
save_directory,
|
||||
|
|
@ -659,7 +648,6 @@ def unsloth_save_model(
|
|||
tags = tags,
|
||||
)
|
||||
|
||||
# Revert back padding side
|
||||
_tokenizer.padding_side = old_padding_side
|
||||
|
||||
if hasattr(model, "config"):
|
||||
|
|
@ -684,14 +672,12 @@ def unsloth_save_model(
|
|||
else:
|
||||
internal_model = model
|
||||
|
||||
# Cannot be converted properly!
|
||||
# LoRA / merged_4bit / non-layered models: save directly without merging
|
||||
if (
|
||||
(save_method == "merged_4bit")
|
||||
or (save_method == "lora")
|
||||
or (not hasattr(model, "model") or not hasattr(internal_model.model, "layers"))
|
||||
):
|
||||
# Do general saving
|
||||
# Edit save_pretrained_settings
|
||||
# [TODO] _create_repo has errors due to **kwargs getting accepted
|
||||
# commit_description does not seem to work?
|
||||
what_to_delete = (
|
||||
|
|
@ -721,7 +707,6 @@ def unsloth_save_model(
|
|||
]
|
||||
)
|
||||
|
||||
# Update model tag
|
||||
if push_to_hub:
|
||||
_ = upload_to_huggingface(
|
||||
model,
|
||||
|
|
@ -745,7 +730,6 @@ def unsloth_save_model(
|
|||
|
||||
tokenizer.save_pretrained(**tokenizer_save_settings)
|
||||
|
||||
# Revert back padding side
|
||||
_tokenizer.padding_side = old_padding_side
|
||||
|
||||
print(" Done.")
|
||||
|
|
@ -795,7 +779,7 @@ def unsloth_save_model(
|
|||
|
||||
print("Unsloth: Merging 4bit and LoRA weights to 16bit...")
|
||||
|
||||
# Determine max RAM usage minus sharding
|
||||
# Max RAM for saving, minus per-shard headroom
|
||||
max_ram = psutil.virtual_memory().available
|
||||
sharded_ram_usage = 5 * 1024 * 1024 * 1024
|
||||
if type(max_shard_size) is str:
|
||||
|
|
@ -808,7 +792,6 @@ def unsloth_save_model(
|
|||
elif type(max_shard_size) is int:
|
||||
sharded_ram_usage = max_shard_size
|
||||
|
||||
# Switch to our fast saving modules if it's a slow PC!
|
||||
n_cpus = psutil.cpu_count(logical = False)
|
||||
if n_cpus is None:
|
||||
n_cpus = psutil.cpu_count()
|
||||
|
|
@ -834,7 +817,7 @@ def unsloth_save_model(
|
|||
if safe_serialization:
|
||||
max_ram -= sharded_ram_usage
|
||||
else:
|
||||
max_ram -= sharded_ram_usage * 0.25 # Uses much less
|
||||
max_ram -= sharded_ram_usage * 0.25
|
||||
|
||||
max_ram = int(max(0, max_ram) * maximum_memory_usage)
|
||||
print(
|
||||
|
|
@ -847,20 +830,18 @@ def unsloth_save_model(
|
|||
if IS_KAGGLE_ENVIRONMENT:
|
||||
temporary_location = os.path.join(KAGGLE_TMP, temporary_location)
|
||||
|
||||
# Max directory for disk saving
|
||||
if not os.path.exists(temporary_location):
|
||||
os.makedirs(temporary_location)
|
||||
|
||||
# Check if Kaggle or Colab, since only 20GB of Disk space allowed.
|
||||
# Kaggle/Colab only allow ~20GB disk, so free up the downloaded model
|
||||
if IS_KAGGLE_ENVIRONMENT or IS_COLAB_ENVIRONMENT:
|
||||
# We free up 4GB of space
|
||||
logger.warning_once(
|
||||
"Unsloth: Kaggle/Colab has limited disk space. We need to delete the downloaded\n"
|
||||
"model which will save 4-16GB of disk space, allowing you to save on Kaggle/Colab."
|
||||
)
|
||||
_free_cached_model(internal_model)
|
||||
|
||||
# HF also uses a OrderedDict
|
||||
# HF also uses an OrderedDict
|
||||
from collections import OrderedDict
|
||||
|
||||
state_dict = OrderedDict()
|
||||
|
|
@ -872,7 +853,6 @@ def unsloth_save_model(
|
|||
elif torch_dtype == "bfloat16":
|
||||
torch_dtype = torch.bfloat16
|
||||
|
||||
# Check modules to save float32 dtype
|
||||
state_dict["model.embed_tokens.weight"] = internal_model.model.embed_tokens.weight.data.to(
|
||||
torch_dtype
|
||||
)
|
||||
|
|
@ -889,12 +869,10 @@ def unsloth_save_model(
|
|||
name = f"model.layers.{j}.{item}.weight"
|
||||
W, bias = _merge_lora(proj, name)
|
||||
|
||||
# Bias term
|
||||
if bias is not None:
|
||||
state_dict[f"model.layers.{j}.{item}.bias"] = bias
|
||||
|
||||
if (torch.cuda.memory_allocated() + W.nbytes) < max_vram:
|
||||
# Save to GPU memory
|
||||
state_dict[name] = W
|
||||
# [TODO] Saving to RAM seems to leak memory???
|
||||
# elif (max_ram - W.nbytes) > 0:
|
||||
|
|
@ -903,7 +881,6 @@ def unsloth_save_model(
|
|||
# state_dict[name] = W.to("cpu", non_blocking = True, copy = True)
|
||||
# max_ram = max(max_ram - W.nbytes, 0)
|
||||
else:
|
||||
# Save to Disk
|
||||
logger.warning_once("\nWe will save to Disk and not RAM now.")
|
||||
filename = os.path.join(temporary_location, f"{name}.pt")
|
||||
torch.save(
|
||||
|
|
@ -940,7 +917,6 @@ def unsloth_save_model(
|
|||
if type(value) is not torch.Tensor:
|
||||
logger.warning_once(f"Unsloth: {key} is not a Tensor but a {type(value)}.")
|
||||
|
||||
# Edit save_pretrained_settings
|
||||
# [TODO] _create_repo has errors due to **kwargs getting accepted
|
||||
save_pretrained_settings["state_dict"] = state_dict
|
||||
|
||||
|
|
@ -972,7 +948,6 @@ def unsloth_save_model(
|
|||
]
|
||||
)
|
||||
|
||||
# Update model tag
|
||||
if push_to_hub:
|
||||
_ = upload_to_huggingface(
|
||||
model,
|
||||
|
|
@ -986,7 +961,6 @@ def unsloth_save_model(
|
|||
datasets = datasets,
|
||||
)
|
||||
|
||||
# First check if we're pushing to an organization!
|
||||
save_directory = save_pretrained_settings["save_directory"]
|
||||
|
||||
if save_pretrained_settings["push_to_hub"]:
|
||||
|
|
@ -998,14 +972,12 @@ def unsloth_save_model(
|
|||
else:
|
||||
actual_username = username
|
||||
|
||||
# Check if pushing to an organization
|
||||
# Pushing to an organization: upload everything at the end
|
||||
if save_pretrained_settings["push_to_hub"] and (username != actual_username):
|
||||
print(f"Unsloth: Saving to organization with address {new_save_directory}")
|
||||
# We upload everything at the end!
|
||||
tokenizer_save_settings["push_to_hub"] = False
|
||||
tokenizer_save_settings["save_directory"] = new_save_directory
|
||||
|
||||
# Save tokenizer
|
||||
if tokenizer is not None:
|
||||
print("Unsloth: Saving tokenizer...", end = "")
|
||||
|
||||
|
|
@ -1021,14 +993,13 @@ def unsloth_save_model(
|
|||
filename_prefix = tokenizer_save_settings.get("filename_prefix"),
|
||||
)
|
||||
|
||||
# Revert back padding side
|
||||
_tokenizer.padding_side = old_padding_side
|
||||
|
||||
print(" Done.")
|
||||
else:
|
||||
print()
|
||||
|
||||
# Since merged, edit quantization_config
|
||||
# Merged model is no longer quantized: drop quantization_config
|
||||
old_config = model.config
|
||||
new_config = model.config.to_dict()
|
||||
if "quantization_config" in new_config:
|
||||
|
|
@ -1040,20 +1011,16 @@ def unsloth_save_model(
|
|||
original_model.config = new_config
|
||||
model.config = new_config
|
||||
|
||||
# Save!
|
||||
# [TODO] --> is this correct?
|
||||
# save_pretrained_settings["selected_adapters"] = None
|
||||
|
||||
# Check if pushing to an organization
|
||||
if save_pretrained_settings["push_to_hub"] and (username != actual_username):
|
||||
print(f"Unsloth: Saving to organization with address {new_save_directory}")
|
||||
# Pushing to organization: .save_pretrained doesn't work, so save
|
||||
# locally first then upload manually.
|
||||
# Org push: .save_pretrained doesn't work, so save locally then upload
|
||||
save_pretrained_settings["save_directory"] = new_save_directory
|
||||
save_pretrained_settings["push_to_hub"] = False
|
||||
internal_model.save_pretrained(**save_pretrained_settings)
|
||||
|
||||
# Now manually go through each file and upload them manually!
|
||||
filenames = os.listdir(new_save_directory)
|
||||
|
||||
hf_api = HfApi(token = save_pretrained_settings["token"])
|
||||
|
|
@ -1070,7 +1037,7 @@ def unsloth_save_model(
|
|||
else:
|
||||
internal_model.save_pretrained(**save_pretrained_settings)
|
||||
|
||||
# Revert config back
|
||||
# Restore the original config
|
||||
original_model = model
|
||||
while hasattr(original_model, "model"):
|
||||
original_model = original_model.model
|
||||
|
|
@ -1095,8 +1062,6 @@ def unsloth_save_model(
|
|||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Remove temporary location
|
||||
|
||||
shutil.rmtree(temporary_location, ignore_errors = True)
|
||||
|
||||
for _ in range(3):
|
||||
|
|
@ -1122,16 +1087,15 @@ def install_llama_cpp_make_non_blocking():
|
|||
# https://github.com/ggerganov/llama.cpp/issues/7062
|
||||
# Weirdly GPU conversion for GGUF breaks??
|
||||
# env = { **os.environ, "LLAMA_CUDA": "1", }
|
||||
# Force make clean
|
||||
check = os.system("make clean -C llama.cpp")
|
||||
IS_CMAKE = False
|
||||
if check == 0:
|
||||
# Uses old MAKE
|
||||
# Old MAKE build
|
||||
n_jobs = max(int((psutil.cpu_count() or 1) * 1.5), 1)
|
||||
full_command = ["make", "all", "-j" + str(n_jobs), "-C", "llama.cpp"]
|
||||
IS_CMAKE = False
|
||||
else:
|
||||
# Uses new CMAKE
|
||||
# New CMAKE build
|
||||
n_jobs = max(int(psutil.cpu_count() or 1), 1) # Use less CPUs since 1.5x faster
|
||||
check = os.system(
|
||||
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}"
|
||||
|
|
@ -1202,8 +1166,7 @@ def try_execute(commands, force_complete = False):
|
|||
|
||||
|
||||
def install_llama_cpp_old(version = -10):
|
||||
# Download the 10th latest release since the latest might be broken!
|
||||
# FALLBACK mechanism
|
||||
# Download the 10th latest release since the latest might be broken (fallback)
|
||||
releases = subprocess.check_output(
|
||||
["git", "ls-remote", "--tags", "https://github.com/ggerganov/llama.cpp.git"]
|
||||
)
|
||||
|
|
@ -1215,7 +1178,6 @@ def install_llama_cpp_old(version = -10):
|
|||
latest = releases[-1]
|
||||
version = releases[version].split(" ")[0]
|
||||
|
||||
# Check if the llama.cpp exists
|
||||
if os.path.exists("llama.cpp"):
|
||||
print(
|
||||
"**[WARNING]** You have a llama.cpp directory which is broken.\n"
|
||||
|
|
@ -1230,8 +1192,7 @@ def install_llama_cpp_old(version = -10):
|
|||
|
||||
shutil.rmtree("llama.cpp", ignore_errors = True)
|
||||
|
||||
# Clone a specific commit
|
||||
# Also don't use the GPU!
|
||||
# Clone a specific commit; don't use the GPU
|
||||
commands = [
|
||||
"git clone --recursive https://github.com/ggerganov/llama.cpp",
|
||||
f"cd llama.cpp && git reset --hard {version} && git clean -df",
|
||||
|
|
@ -1254,7 +1215,6 @@ def install_llama_cpp_old(version = -10):
|
|||
|
||||
try_execute(commands)
|
||||
|
||||
# Check if successful
|
||||
if not (
|
||||
os.path.exists("llama.cpp/llama-quantize.exe")
|
||||
or os.path.exists("llama.cpp/llama-quantize")
|
||||
|
|
@ -1302,13 +1262,11 @@ def install_llama_cpp_blocking(use_cuda = False):
|
|||
|
||||
|
||||
def get_executable(executables):
|
||||
# Get system locations (System Path).split(system separator)
|
||||
system_directories = os.environ.get("PATH").split(os.pathsep)
|
||||
|
||||
for directory in system_directories:
|
||||
for executable in executables:
|
||||
path = os.path.join(directory, executable)
|
||||
# Check if the executable exists and is executable
|
||||
if os.path.exists(path) and os.access(path, os.X_OK):
|
||||
return path
|
||||
return None
|
||||
|
|
@ -1325,17 +1283,12 @@ def save_to_gguf(
|
|||
is_vlm: bool = False,
|
||||
is_gpt_oss: bool = False,
|
||||
):
|
||||
"""
|
||||
Orchestrates the complete GGUF conversion process.
|
||||
Handles installation, conversion, and quantization.
|
||||
"""
|
||||
# print_output True only if UNSLOTH_ENABLE_LOGGING=1
|
||||
"""Orchestrate GGUF conversion: install, convert, and quantize."""
|
||||
if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
|
||||
print_output = True
|
||||
else:
|
||||
print_output = False
|
||||
|
||||
# Validate model dtype
|
||||
assert model_dtype == "float16" or model_dtype == "bfloat16"
|
||||
model_dtype = "f16" if model_dtype == "float16" else "bf16"
|
||||
|
||||
|
|
@ -1359,11 +1312,10 @@ def save_to_gguf(
|
|||
)
|
||||
model_dtype = "f16"
|
||||
|
||||
# Check first_conversion as well
|
||||
if first_conversion is None:
|
||||
first_conversion = model_dtype
|
||||
|
||||
# Check I quants
|
||||
# Reject I-quants (not yet supported)
|
||||
for quant_method in quantization_method:
|
||||
if quant_method.startswith("iq2"):
|
||||
raise RuntimeError(
|
||||
|
|
@ -1382,7 +1334,6 @@ def save_to_gguf(
|
|||
elif quant_method is None:
|
||||
quant_method = "q8_0"
|
||||
|
||||
# Check if wrong method
|
||||
if quant_method not in ALLOWED_QUANTS.keys():
|
||||
error = f"Unsloth: Quant method = [{quant_method}] not supported. Choose from below:\n"
|
||||
for key, value in ALLOWED_QUANTS.items():
|
||||
|
|
@ -1395,12 +1346,10 @@ def save_to_gguf(
|
|||
# Determine optimal first_conversion
|
||||
if is_gpt_oss:
|
||||
print("Unsloth: GPT-OSS model detected - using special conversion settings")
|
||||
first_conversion = "None" # No quantization for GPT-OSS
|
||||
# Only keep one conversion method since GPT-OSS doesn't quantize
|
||||
first_conversion = "None" # GPT-OSS isn't quantized
|
||||
quantization_method = ["None"]
|
||||
else:
|
||||
if first_conversion is None:
|
||||
# Check if q8_0 is the ONLY quantization method requested
|
||||
if len(quantization_method) == 1 and quantization_method[0] == "q8_0":
|
||||
first_conversion = "None" # Let llama-quantize do the direct conversion
|
||||
else:
|
||||
|
|
@ -1431,7 +1380,6 @@ def save_to_gguf(
|
|||
first_conversion = "f16"
|
||||
|
||||
first_conversion_dtype = "" if first_conversion == "None" else first_conversion
|
||||
# Print conversion info
|
||||
print_info = (
|
||||
f"==((====))== Unsloth: Conversion from HF to GGUF information\n"
|
||||
f" {chr(92)}{chr(92)} /| [0] Installing llama.cpp might take 3 minutes.\n"
|
||||
|
|
@ -1482,9 +1430,7 @@ def save_to_gguf(
|
|||
max_shard_size = "50GB",
|
||||
print_output = print_output,
|
||||
)
|
||||
# update is_vlm switch
|
||||
is_vlm = is_vlm_update
|
||||
# Check conversion success
|
||||
for file in initial_files:
|
||||
if not os.path.exists(file):
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
|
|
@ -1515,7 +1461,6 @@ def save_to_gguf(
|
|||
# Step 4: Additional quantizations using llama-quantize
|
||||
all_saved_locations = initial_files.copy()
|
||||
|
||||
# Get CPU count for quantization
|
||||
n_cpus = psutil.cpu_count()
|
||||
if n_cpus is None:
|
||||
n_cpus = 1
|
||||
|
|
@ -1639,10 +1584,8 @@ def unsloth_save_pretrained_merged(
|
|||
datasets: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Same as .save_pretrained(...) except 4bit weights are auto
|
||||
converted to float16 with as few overhead as possible.
|
||||
|
||||
Choose for `save_method` to be either:
|
||||
Like .save_pretrained(...) but auto-converts 4bit weights to float16.
|
||||
`save_method`:
|
||||
1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp.
|
||||
2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference.
|
||||
3. `lora`: Save LoRA adapters with no merging. Useful for HF inference.
|
||||
|
|
@ -1681,10 +1624,8 @@ def unsloth_push_to_hub_merged(
|
|||
datasets: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Same as .push_to_hub(...) except 4bit weights are auto
|
||||
converted to float16 with as few overhead as possible.
|
||||
|
||||
Choose for `save_method` to be either:
|
||||
Like .push_to_hub(...) but auto-converts 4bit weights to float16.
|
||||
`save_method`:
|
||||
1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp.
|
||||
2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference.
|
||||
3. `lora`: Save LoRA adapters with no merging. Useful for HF inference.
|
||||
|
|
@ -1770,7 +1711,6 @@ def create_huggingface_repo(
|
|||
private = private,
|
||||
)
|
||||
|
||||
# Create model card
|
||||
from huggingface_hub import ModelCard
|
||||
|
||||
content = MODEL_CARD.format(
|
||||
|
|
@ -1823,7 +1763,6 @@ def upload_to_huggingface(
|
|||
private = private,
|
||||
)
|
||||
|
||||
# Create model card
|
||||
from huggingface_hub import ModelCard
|
||||
|
||||
content = MODEL_CARD.format(
|
||||
|
|
@ -1849,7 +1788,6 @@ def upload_to_huggingface(
|
|||
)
|
||||
|
||||
if file_location is not None:
|
||||
# Now upload file
|
||||
hf_api = HfApi(token = token)
|
||||
|
||||
if "/" in file_location:
|
||||
|
|
@ -1901,7 +1839,7 @@ def upload_to_huggingface(
|
|||
|
||||
|
||||
def fix_tokenizer_bos_token(tokenizer):
|
||||
# Check if BOS added already, then warn
|
||||
# Warn + strip if the model auto-adds a BOS and the template also has one
|
||||
fix_bos_token = False
|
||||
chat_template = getattr(tokenizer, "chat_template", None)
|
||||
|
||||
|
|
@ -1950,7 +1888,7 @@ def create_ollama_modelfile(tokenizer, base_model_name, model_location):
|
|||
f"Unsloth: No Ollama template mapping found for model '{base_model_name}'. Skipping Ollama Modelfile"
|
||||
)
|
||||
return None
|
||||
tokenizer._ollama_modelfile = ollama_modelfile # This comes from the unpacking above
|
||||
tokenizer._ollama_modelfile = ollama_modelfile
|
||||
modelfile = ollama_modelfile
|
||||
|
||||
FILE_LOCATION_REPLACER = "⚫@✅#🦥__FILE_LOCATION__⚡@🦥#⛵"
|
||||
|
|
@ -2113,10 +2051,8 @@ def unsloth_save_pretrained_gguf(
|
|||
maximum_memory_usage: float = 0.85,
|
||||
):
|
||||
"""
|
||||
Same as .save_pretrained(...) except 4bit weights are auto
|
||||
converted to float16 then converted to GGUF / llama.cpp format.
|
||||
|
||||
Choose for `quantization_method` to be:
|
||||
Like .save_pretrained(...) but auto-converts 4bit weights to float16, then to
|
||||
GGUF / llama.cpp format. `quantization_method`:
|
||||
"not_quantized" : "Recommended. Fast conversion. Slow inference, big files.",
|
||||
"fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.",
|
||||
"quantized" : "Recommended. Slow conversion. Fast inference, small files.",
|
||||
|
|
@ -2227,7 +2163,6 @@ def unsloth_save_pretrained_gguf(
|
|||
if is_peft_model:
|
||||
print(f'Unsloth: Merging model weights to {"mxfp4" if is_gpt_oss else "16-bit"} format...')
|
||||
try:
|
||||
# Call unsloth_generic_save directly (it's in the same file)
|
||||
unsloth_generic_save(**arguments)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -2290,11 +2225,9 @@ def unsloth_save_pretrained_gguf(
|
|||
# Step 8: Convert to GGUF format
|
||||
print("Unsloth: Converting to GGUF format...")
|
||||
|
||||
# Convert quantization_method to list if string
|
||||
# Use old style quantization_method
|
||||
# Normalize quantization_method (old-style) to a list
|
||||
quantization_methods = []
|
||||
if quantization_method is not None:
|
||||
# Convert quantization_method to list
|
||||
if isinstance(quantization_method, list):
|
||||
pass
|
||||
elif isinstance(quantization_method, str):
|
||||
|
|
@ -2334,8 +2267,8 @@ def unsloth_save_pretrained_gguf(
|
|||
model_directory = save_directory,
|
||||
quantization_method = quantization_methods,
|
||||
first_conversion = first_conversion,
|
||||
is_vlm = is_vlm, # Pass VLM flag
|
||||
is_gpt_oss = is_gpt_oss, # Pass gpt_oss Flag
|
||||
is_vlm = is_vlm,
|
||||
is_gpt_oss = is_gpt_oss,
|
||||
)
|
||||
except Exception as e:
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
|
|
@ -2434,10 +2367,8 @@ def unsloth_push_to_hub_gguf(
|
|||
datasets: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Same as .push_to_hub(...) except 4bit weights are auto
|
||||
converted to float16 then converted to GGUF / llama.cpp format.
|
||||
|
||||
Choose for `quantization_method` to be:
|
||||
Like .push_to_hub(...) but auto-converts 4bit weights to float16, then to
|
||||
GGUF / llama.cpp format. `quantization_method`:
|
||||
"not_quantized" : "Recommended. Fast conversion. Slow inference, big files.",
|
||||
"fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.",
|
||||
"quantized" : "Recommended. Slow conversion. Fast inference, small files.",
|
||||
|
|
@ -2519,14 +2450,12 @@ def unsloth_push_to_hub_gguf(
|
|||
|
||||
api = HfApi(token = token)
|
||||
|
||||
# Get full repo id
|
||||
if "/" not in repo_id:
|
||||
username = api.whoami()["name"]
|
||||
full_repo_id = f"{username}/{repo_id}"
|
||||
else:
|
||||
full_repo_id = repo_id
|
||||
|
||||
# Create repo
|
||||
api.create_repo(
|
||||
repo_id = full_repo_id,
|
||||
repo_type = "model",
|
||||
|
|
@ -2658,7 +2587,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
|
|||
|
||||
print(f"Unsloth: Successfully uploaded GGUF to https://huggingface.co/{full_repo_id}")
|
||||
|
||||
# Add tags
|
||||
if tags is None:
|
||||
tags = []
|
||||
tags.extend(["gguf", "llama-cpp", "unsloth"])
|
||||
|
|
@ -2687,7 +2615,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
|
|||
raise RuntimeError(f"Failed to upload to Hugging Face Hub: {e}")
|
||||
|
||||
finally:
|
||||
# Clean up temporary directory
|
||||
if cleanup_temp:
|
||||
print("Unsloth: Cleaning up temporary files...")
|
||||
for d in [save_directory, f"{save_directory}_gguf"]:
|
||||
|
|
@ -2700,12 +2627,9 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi
|
|||
return full_repo_id
|
||||
|
||||
|
||||
# Corrected function to save LoRA to a custom directory
|
||||
def save_lora_to_custom_dir(model, tokenizer, save_directory):
|
||||
# Create the custom directory if it doesn't exist
|
||||
os.makedirs(save_directory, exist_ok = True)
|
||||
|
||||
# Call the unsloth_save_model function with the custom directory
|
||||
unsloth_save_model(
|
||||
model,
|
||||
tokenizer,
|
||||
|
|
@ -2715,7 +2639,6 @@ def save_lora_to_custom_dir(model, tokenizer, save_directory):
|
|||
)
|
||||
|
||||
|
||||
# Corrected method within the model class to convert LoRA to GGML and push to Hugging Face Hub
|
||||
def unsloth_convert_lora_to_ggml_and_push_to_hub(
|
||||
self,
|
||||
tokenizer,
|
||||
|
|
@ -2830,7 +2753,6 @@ def unsloth_convert_lora_to_ggml_and_save_locally(
|
|||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
# Use the provided save_directory for local saving
|
||||
save_lora_to_custom_dir(self, tokenizer, save_directory)
|
||||
|
||||
model_type = self.config.model_type
|
||||
|
|
@ -2901,10 +2823,9 @@ def save_to_gguf_generic(
|
|||
if not os.path.exists(os.path.join("llama.cpp", "unsloth_convert_hf_to_gguf.py")):
|
||||
install_llama_cpp(just_clone_repo = True)
|
||||
|
||||
# Use old style quantization_method
|
||||
# Normalize quantization_method (old-style) to a list
|
||||
new_quantization_methods = []
|
||||
if quantization_method is not None:
|
||||
# Convert quantization_method to list
|
||||
if isinstance(quantization_method, list):
|
||||
pass
|
||||
elif isinstance(quantization_method, str):
|
||||
|
|
@ -2930,7 +2851,6 @@ def save_to_gguf_generic(
|
|||
new_quantization_methods.append(quant_method.lower())
|
||||
else:
|
||||
new_quantization_methods.append(quantization_type.lower())
|
||||
# Check if wrong method
|
||||
for quant_method in new_quantization_methods:
|
||||
if quant_method not in ALLOWED_QUANTS.keys():
|
||||
error = f"Unsloth: Quant method = [{quant_method}] not supported. Choose from below:\n"
|
||||
|
|
@ -2938,8 +2858,7 @@ def save_to_gguf_generic(
|
|||
error += f"[{key}] => {value}\n"
|
||||
raise RuntimeError(error)
|
||||
|
||||
# Go through all types and save individually - somewhat inefficient
|
||||
# since we save F16 / BF16 multiple times
|
||||
# Save each type individually (inefficient: F16/BF16 saved repeatedly)
|
||||
for quantization_type in new_quantization_methods:
|
||||
metadata = _convert_to_gguf(
|
||||
save_directory,
|
||||
|
|
@ -3123,10 +3042,8 @@ def unsloth_generic_save_pretrained_merged(
|
|||
datasets: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Same as .push_to_hub(...) except 4bit weights are auto
|
||||
converted to float16 with as few overhead as possible.
|
||||
|
||||
Choose for `save_method` to be either:
|
||||
Like .push_to_hub(...) but auto-converts 4bit weights to float16.
|
||||
`save_method`:
|
||||
1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp.
|
||||
2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference.
|
||||
3. `lora`: Save LoRA adapters with no merging. Useful for HF inference.
|
||||
|
|
@ -3165,10 +3082,8 @@ def unsloth_generic_push_to_hub_merged(
|
|||
datasets: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Same as .push_to_hub(...) except 4bit weights are auto
|
||||
converted to float16 with as few overhead as possible.
|
||||
|
||||
Choose for `save_method` to be either:
|
||||
Like .push_to_hub(...) but auto-converts 4bit weights to float16.
|
||||
`save_method`:
|
||||
1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp.
|
||||
2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference.
|
||||
3. `lora`: Save LoRA adapters with no merging. Useful for HF inference.
|
||||
|
|
@ -3198,9 +3113,8 @@ def _unsloth_save_torchao_with_attached_config(
|
|||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Save a QAT-trained model by converting fake-quantized weights to real quantized weights."""
|
||||
# Convert QAT fake-quantized weights to real quantized weights
|
||||
_convert_torchao_model(model)
|
||||
# PEFT models also might come here, so parse it
|
||||
# PEFT models can also reach here, so parse it
|
||||
if isinstance(model, PeftModelForCausalLM):
|
||||
_unsloth_save_torchao_with_given_config(
|
||||
model = model,
|
||||
|
|
@ -3231,12 +3145,11 @@ def _unsloth_save_torchao_with_given_config(
|
|||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Quantizes the model with torchao and saves a torchao quantized checkpoint
|
||||
"""Quantize the model with torchao and save the quantized checkpoint.
|
||||
|
||||
Args
|
||||
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is set to True, e.g. `my_model`
|
||||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or save locally
|
||||
`save_directory`: local path, or hub repo ID when `push_to_hub` is True.
|
||||
`torchao_config` (TorchAOBaseConfig): torchao quant config, full list:
|
||||
https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
"""
|
||||
|
||||
if push_to_hub:
|
||||
|
|
@ -3337,26 +3250,15 @@ def unsloth_save_pretrained_torchao(
|
|||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Saves a torchao quantized model checkpoint.
|
||||
"""Save a torchao quantized model checkpoint. Two exclusive workflows:
|
||||
|
||||
This function handles two mutually exclusive workflows:
|
||||
1. QAT: model trained with `qat_scheme` -> do NOT pass `torchao_config`;
|
||||
fake-quantized weights are converted to real quantized weights and saved.
|
||||
2. PTQ: model NOT trained with `qat_scheme` -> pass a `torchao_config` to quantize.
|
||||
|
||||
1. **QAT (Quantization-Aware Training)**: If the model was trained with `qat_scheme`
|
||||
parameter, do NOT pass `torchao_config`. The function will convert the QAT
|
||||
fake-quantized weights to real quantized weights and save directly.
|
||||
|
||||
2. **PTQ (Post-Training Quantization)**: If you want to apply quantization to a
|
||||
regular model, pass a `torchao_config`. The model must NOT have been trained
|
||||
with `qat_scheme`.
|
||||
|
||||
Args:
|
||||
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is True
|
||||
`tokenizer`: the tokenizer to save alongside the model
|
||||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization.
|
||||
Required for PTQ, must be None for QAT models.
|
||||
Options: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push to huggingface hub or save locally
|
||||
`token`: HuggingFace token for pushing to hub
|
||||
`save_directory`: local path, or hub repo ID when `push_to_hub` is True.
|
||||
`torchao_config` (TorchAOBaseConfig): required for PTQ, must be None for QAT.
|
||||
Options: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
"""
|
||||
if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)):
|
||||
tokenizer = patch_saving_functions(tokenizer)
|
||||
|
|
@ -3409,7 +3311,7 @@ def patch_saving_functions(model, vision = False):
|
|||
import types
|
||||
from typing import Callable, Optional, Union, List
|
||||
|
||||
# And now re add our saving methods!
|
||||
# Re-add our saving methods
|
||||
if model.push_to_hub.__name__ == "unsloth_push_to_hub":
|
||||
original_push_to_hub = model.original_push_to_hub
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -65,7 +65,6 @@ IGNORED_TOKENIZER_NAMES = frozenset(
|
|||
)
|
||||
os.environ["UNSLOTH_IGNORED_TOKENIZER_NAMES"] = "\n".join(IGNORED_TOKENIZER_NAMES)
|
||||
|
||||
# Check environments
|
||||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames
|
||||
IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames
|
||||
|
|
@ -159,7 +158,7 @@ def convert_to_fast_tokenizer(slow_tokenizer, temporary_location = "_unsloth_sen
|
|||
args = re.findall(r"\n[\s]+([^\s]{1,}) \(", docs, flags = re.MULTILINE)
|
||||
args = [x for x in args if not x.endswith("_file")]
|
||||
|
||||
# Also some missing maybe!
|
||||
# Also pull args from the base class in case some are missing
|
||||
docs = PreTrainedTokenizerFast.__doc__
|
||||
docs = docs[docs.find("Args:") :]
|
||||
args2 = re.findall(r"\n[\s]+([^\s]{1,}) \(", docs, flags = re.MULTILINE)
|
||||
|
|
@ -179,17 +178,14 @@ def convert_to_fast_tokenizer(slow_tokenizer, temporary_location = "_unsloth_sen
|
|||
check_vocab = sorted_slow_tokenizer == sorted_fast_tokenizer
|
||||
check_special = slow_tokenizer.all_special_tokens == fast_tokenizer.all_special_tokens
|
||||
|
||||
# Failure so return slow_tokenizer
|
||||
if not check_vocab or not check_special:
|
||||
return slow_tokenizer
|
||||
|
||||
# Now confirm if they match
|
||||
if not assert_same_tokenization(slow_tokenizer, fast_tokenizer):
|
||||
# Maybe remove prepending of __apple?
|
||||
kwargs["tokenizer_object"] = try_fix_tokenizer(slow_tokenizer, prepend = False)
|
||||
fast_tokenizer = FastTokenizer(**kwargs)
|
||||
if not assert_same_tokenization(slow_tokenizer, fast_tokenizer):
|
||||
# Failure :(
|
||||
return slow_tokenizer
|
||||
|
||||
# Also tokenizer.model is missing!
|
||||
|
|
@ -200,7 +196,6 @@ def convert_to_fast_tokenizer(slow_tokenizer, temporary_location = "_unsloth_sen
|
|||
slow_tokenizer.save_pretrained(new_location)
|
||||
fast_tokenizer.save_pretrained(new_location)
|
||||
|
||||
# Now load it!
|
||||
fast_tokenizer = AutoTokenizer.from_pretrained(new_location)
|
||||
if assert_same_tokenization(slow_tokenizer, fast_tokenizer):
|
||||
return fast_tokenizer
|
||||
|
|
@ -277,7 +272,6 @@ def assert_same_tokenization(slow_tokenizer, fast_tokenizer):
|
|||
replacement_char = b"\xc3\xaf\xc2\xbf\xc2\xbd".decode("utf-8")
|
||||
all_special_tokens = [x for x in all_special_tokens if x != replacement_char]
|
||||
|
||||
# Check if chat template is enabled!
|
||||
check_chat_template1 = True
|
||||
check_chat_template2 = True
|
||||
check_chat_template3 = True
|
||||
|
|
@ -370,25 +364,21 @@ def fix_sentencepiece_tokenizer(
|
|||
if not os.path.exists(temporary_location):
|
||||
os.makedirs(temporary_location)
|
||||
|
||||
# Check if tokenizer.model exists
|
||||
if not os.path.isfile(f"{temporary_location}/tokenizer.model"):
|
||||
return new_tokenizer
|
||||
|
||||
# First save the old tokenizer
|
||||
old_tokenizer.save_pretrained(temporary_location)
|
||||
|
||||
tokenizer_file = sentencepiece_model_pb2.ModelProto()
|
||||
tokenizer_file.ParseFromString(open(f"{temporary_location}/tokenizer.model", "rb").read())
|
||||
|
||||
# Now save the new tokenizer
|
||||
new_tokenizer.save_pretrained(temporary_location)
|
||||
|
||||
# Now correct the old tokenizer's .model file
|
||||
# Correct the old tokenizer's .model file
|
||||
for old_token, new_token in token_mapping.items():
|
||||
ids = old_tokenizer([old_token], add_special_tokens = False).input_ids
|
||||
ids = ids[0]
|
||||
if len(ids) != 1:
|
||||
# Skip this token!
|
||||
print(
|
||||
f"Skip mapping {old_token} to {new_token} since {new_token} is already in the tokenizer!"
|
||||
)
|
||||
|
|
@ -402,11 +392,9 @@ def fix_sentencepiece_tokenizer(
|
|||
assert tokenizer_piece.piece == old_token
|
||||
tokenizer_piece.piece = new_token
|
||||
|
||||
# And now write it
|
||||
with open(f"{temporary_location}/tokenizer.model", "wb") as file:
|
||||
file.write(tokenizer_file.SerializeToString())
|
||||
|
||||
# And load it!
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
|
|
@ -418,14 +406,12 @@ def fix_sentencepiece_tokenizer(
|
|||
|
||||
|
||||
def fix_sentencepiece_gguf(saved_location):
|
||||
"""
|
||||
Fix sentencepiece tokenizers that didn't extend the vocab with user-defined
|
||||
tokens. Inspired by llama.cpp's convert_hf_to_gguf.py.
|
||||
"""Fix sentencepiece tokenizers that didn't extend the vocab with user-defined tokens (inspired
|
||||
by llama.cpp's convert_hf_to_gguf.py).
|
||||
|
||||
Also retypes special tokens (e.g. Gemma 3's <start_of_turn>/<end_of_turn>)
|
||||
that exist in the sentencepiece model but are typed NORMAL instead of CONTROL.
|
||||
NORMAL writes token_type=1 to GGUF, breaking llama.cpp chat inference since
|
||||
parse_special only matches CONTROL (type=3).
|
||||
Also retypes special tokens (e.g. Gemma 3's <start_of_turn>/<end_of_turn>) typed NORMAL instead of
|
||||
CONTROL. NORMAL writes token_type=1 to GGUF, breaking llama.cpp chat inference since parse_special
|
||||
only matches CONTROL (type=3).
|
||||
"""
|
||||
from copy import deepcopy
|
||||
import sys
|
||||
|
|
@ -450,7 +436,6 @@ def fix_sentencepiece_gguf(saved_location):
|
|||
UNUSED = 5
|
||||
BYTE = 6
|
||||
|
||||
# Load tokenizer.model
|
||||
tokenizer_file = sentencepiece_model_pb2.ModelProto()
|
||||
if not os.path.isfile(f"{saved_location}/tokenizer.model"):
|
||||
return
|
||||
|
|
@ -484,7 +469,6 @@ def fix_sentencepiece_gguf(saved_location):
|
|||
f"from NORMAL to CONTROL type so llama.cpp / GGUF chat inference works correctly."
|
||||
)
|
||||
|
||||
# Load added_tokens_json
|
||||
if not os.path.isfile(f"{saved_location}/added_tokens.json"):
|
||||
if patched > 0:
|
||||
with open(f"{saved_location}/tokenizer.model", "wb") as file:
|
||||
|
|
@ -560,14 +544,12 @@ def _load_correct_tokenizer(
|
|||
if IS_COLAB_ENVIRONMENT:
|
||||
cache_dir = cache_dir
|
||||
elif IS_KAGGLE_ENVIRONMENT:
|
||||
# /tmp of Kaggle seems has a 80GB limit!
|
||||
# Let's utilize them
|
||||
# /tmp on Kaggle has a ~80GB limit, so use it
|
||||
cache_dir = os.path.join(KAGGLE_TMP, cache_dir)
|
||||
else:
|
||||
cache_dir = None
|
||||
|
||||
# Try loading the slow tokenizer. If it fails, then try Fast only
|
||||
# Mainly to solve Deepseek models with no tokenizer.model file
|
||||
# Try slow tokenizer, fall back to Fast (e.g. Deepseek has no tokenizer.model)
|
||||
slow_tokenizer = None
|
||||
try:
|
||||
slow_tokenizer = AutoTokenizer.from_pretrained(
|
||||
|
|
@ -603,7 +585,7 @@ def _load_correct_tokenizer(
|
|||
|
||||
if not fix_tokenizer or tokenizer_name in IGNORED_TOKENIZER_NAMES:
|
||||
return fast_tokenizer
|
||||
# Ignore Mistral ones - they're a bit weird to handle!
|
||||
# Mistral tokenizers are weird to handle, so skip them
|
||||
elif "mistral" in tokenizer_name.lower():
|
||||
return fast_tokenizer
|
||||
# Ignore Phi-4 ones as well
|
||||
|
|
@ -688,9 +670,9 @@ def _find_end_position(
|
|||
endfor = None,
|
||||
endif = None,
|
||||
):
|
||||
"""Rightmost {% endfor %}/{% endif %} (any dash variant), as a dict
|
||||
with start/end/text/dash_left/dash_right. Tokens inside Jinja comments
|
||||
are ignored. `endfor`/`endif` kwargs kept for back-compat, ignored."""
|
||||
"""Rightmost {% endfor %}/{% endif %} (any dash variant) as a dict with
|
||||
start/end/text/dash_left/dash_right. Tokens inside Jinja comments are ignored.
|
||||
`endfor`/`endif` kwargs kept for back-compat, ignored."""
|
||||
# Space-pad comments so positions still map 1:1 to the original.
|
||||
scrubbed = _RE_JINJA_COMMENT.sub(lambda m: " " * len(m.group(0)), template)
|
||||
endfor_matches = list(_RE_ENDFOR.finditer(scrubbed))
|
||||
|
|
@ -1307,7 +1289,6 @@ def check_tokenizer(
|
|||
# See https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha/discussions/25
|
||||
# Seems like the Fast tokenizer in Rust breaks things!
|
||||
|
||||
# We ignore some of them!
|
||||
if tokenizer.__repr__().split("(", 1)[0] in IGNORED_TOKENIZER_CHECKING:
|
||||
return tokenizer
|
||||
|
||||
|
|
@ -1322,7 +1303,6 @@ def check_tokenizer(
|
|||
bad_indices = list(added_tokens_fast.keys())[j:]
|
||||
bad_tokens = list(added_tokens_fast.values())[j:]
|
||||
if not _reload:
|
||||
# Try removing the token
|
||||
added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()]
|
||||
special_tokens = tokenizer.special_tokens_map
|
||||
import itertools
|
||||
|
|
@ -1337,7 +1317,6 @@ def check_tokenizer(
|
|||
x for x in can_be_removed1 if x in tokenizer._added_tokens_encoder.keys()
|
||||
]
|
||||
|
||||
# Check of extra tokens can in fact we removed!
|
||||
can_be_removed = (len(can_be_removed1) == len(bad_tokens)) and (
|
||||
len(can_be_removed2) == len(bad_tokens)
|
||||
)
|
||||
|
|
@ -1357,14 +1336,12 @@ def check_tokenizer(
|
|||
try_removal.append(token)
|
||||
try_mapper.append(name_token)
|
||||
|
||||
# Recheck!
|
||||
can_be_removed = len(try_removal) == len(bad_tokens)
|
||||
if can_be_removed:
|
||||
remove_generic = True
|
||||
can_be_removed1 = bad_tokens
|
||||
|
||||
if can_be_removed:
|
||||
# Yes it can be fixed!
|
||||
for j, bad_token in enumerate(can_be_removed1):
|
||||
remove_id = tokenizer._added_tokens_encoder[bad_token]
|
||||
del tokenizer._added_tokens_decoder[remove_id]
|
||||
|
|
@ -1374,7 +1351,6 @@ def check_tokenizer(
|
|||
# Remove sep token for example
|
||||
setattr(tokenizer, try_mapper[j], None)
|
||||
setattr(tokenizer, try_mapper[j] + "_id", None)
|
||||
# Confirm 1 more time!
|
||||
if max(tokenizer.added_tokens_decoder.keys()) < max_embedding_size:
|
||||
logger.warning_once(
|
||||
f"Unsloth loaded a broken tokenizer `{model_name}`, but managed to repair it!\n"
|
||||
|
|
@ -1383,7 +1359,6 @@ def check_tokenizer(
|
|||
)
|
||||
return convert_to_fast_tokenizer(tokenizer)
|
||||
|
||||
# :( Failure
|
||||
raise RuntimeError(
|
||||
f"Unsloth tried to load `{model_name}`, but cannot succeed.\n"
|
||||
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"
|
||||
|
|
@ -1397,7 +1372,6 @@ def check_tokenizer(
|
|||
|
||||
# Sometimes slow tokenizer does not work like Deepseek
|
||||
try:
|
||||
# Try slow tokenizer which can fix things!
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
model_max_length = model_max_length,
|
||||
|
|
@ -1420,8 +1394,7 @@ def check_tokenizer(
|
|||
)
|
||||
break
|
||||
except:
|
||||
# Tokenizer has out of bounds issues and we can't
|
||||
# load the slow tokenizer version :(
|
||||
# Out-of-bounds tokenizer and the slow version won't load either
|
||||
logger.warning_once(
|
||||
"Unsloth: Tokenizer is most likely buggy, and Unsloth failed to repair it.\n"
|
||||
"It will still work, but beware of out of bounds memory accesses.\n"
|
||||
|
|
@ -1432,35 +1405,8 @@ def check_tokenizer(
|
|||
|
||||
|
||||
def get_tokenizer_info(tokenizer) -> dict:
|
||||
"""Return a concise diagnostic summary of a tokenizer instance.
|
||||
|
||||
Collects key properties into a JSON-safe dict for logging, debugging, or the
|
||||
Studio UI. Missing attributes fall back to ``None`` rather than raising.
|
||||
|
||||
Example output::
|
||||
|
||||
{
|
||||
"name_or_path": "unsloth/Llama-3.2-1B-Instruct",
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"is_fast": True,
|
||||
"vocab_size": 128000,
|
||||
"added_tokens_count": 256,
|
||||
"model_max_length": 131072,
|
||||
"padding_side": "right",
|
||||
"bos_token": "<|begin_of_text|>",
|
||||
"eos_token": "<|eot_id|>",
|
||||
"pad_token": "<|finetune_right_pad_id|>",
|
||||
"unk_token": None,
|
||||
"has_chat_template": True,
|
||||
"special_tokens_count": 3,
|
||||
}
|
||||
|
||||
Args:
|
||||
tokenizer: Any HuggingFace ``PreTrainedTokenizer(Fast)`` instance.
|
||||
|
||||
Returns:
|
||||
A ``dict`` of tokenizer properties.
|
||||
"""
|
||||
"""Return a concise JSON-safe diagnostic summary of a tokenizer for logging/debugging/Studio UI.
|
||||
Missing attributes fall back to ``None`` rather than raising."""
|
||||
return {
|
||||
"name_or_path": getattr(tokenizer, "name_or_path", None),
|
||||
"tokenizer_class": type(tokenizer).__name__,
|
||||
|
|
@ -1490,26 +1436,12 @@ try:
|
|||
except:
|
||||
|
||||
def neftune_post_forward_hook(module, input, output):
|
||||
"""
|
||||
Implements the NEFTune forward pass for the model using forward hooks. Note this works only for
|
||||
torch.nn.Embedding layers. This method is slightly adapted from the original source code
|
||||
that can be found here: https://github.com/neelsjain/NEFTune
|
||||
|
||||
Simply add it to your model as follows:
|
||||
"""NEFTune forward hook for torch.nn.Embedding layers (adapted from
|
||||
https://github.com/neelsjain/NEFTune). Set `module.neftune_noise_alpha`, then register:
|
||||
```python
|
||||
model = ...
|
||||
model.embed_tokens.neftune_noise_alpha = 0.1
|
||||
model.embed_tokens.register_forward_hook(neftune_post_forward_hook)
|
||||
```
|
||||
|
||||
Args:
|
||||
module (`torch.nn.Module`):
|
||||
The embedding module where the hook is attached. Note that you need to set
|
||||
`module.neftune_noise_alpha` to the desired noise alpha value.
|
||||
input (`torch.Tensor`):
|
||||
The input tensor to the model.
|
||||
output (`torch.Tensor`):
|
||||
The output tensor of the model (i.e. the embeddings).
|
||||
"""
|
||||
if module.training:
|
||||
dims = torch.tensor(output.size(1) * output.size(2))
|
||||
|
|
@ -1519,9 +1451,7 @@ except:
|
|||
|
||||
|
||||
def patch_sft_trainer_tokenizer():
|
||||
"""
|
||||
Patches the trainer with changes
|
||||
"""
|
||||
"""Patches the SFT trainer with Unsloth changes."""
|
||||
try:
|
||||
sft_trainer = eval(f"trl.trainer.sft_trainer.SFTTrainer")
|
||||
except:
|
||||
|
|
|
|||
|
|
@ -57,11 +57,10 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
|
||||
class UnslothVisionDataCollator(_UnslothVisionDataCollatorBase):
|
||||
"""
|
||||
Drop-in zoo collator that validates local video paths on every batch
|
||||
(deduped across batches), applying formatting_func first so formatter-made
|
||||
paths are checked too. Raises FileNotFoundError on missing files instead
|
||||
of silently training on empty video tensors (issue #5085).
|
||||
"""Drop-in zoo collator that validates local video paths per batch (deduped
|
||||
across batches), applying formatting_func first so formatter-made paths are
|
||||
checked too. Raises FileNotFoundError on missing files instead of silently
|
||||
training on empty video tensors (issue #5085).
|
||||
"""
|
||||
|
||||
__slots__ = ("_checked_video_paths",)
|
||||
|
|
@ -372,8 +371,7 @@ class UnslothTrainer(SFTTrainer):
|
|||
return self.optimizer
|
||||
|
||||
|
||||
# From `trl>=0.13.0`, they changed how to pass several params to the trainer
|
||||
# We need to patch to make the transition smooth
|
||||
# trl>=0.13.0 changed how several params are passed to the trainer; patch for it
|
||||
def _resolve_trainer_params(trainer_class, init_fn):
|
||||
"""Resolve the real named parameters for a trainer __init__.
|
||||
|
||||
|
|
@ -473,7 +471,6 @@ def _backwards_compatible_trainer(trainer_class, config_class):
|
|||
else:
|
||||
config = training_args
|
||||
|
||||
# Reconstruct kwargs for Trainer
|
||||
kwargs = trainer_kwargs
|
||||
kwargs["args"] = config
|
||||
original_init(self, *args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -53,12 +53,10 @@ XFORMERS_BLOCK_DIAG_CLS = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFOR
|
|||
|
||||
@dataclass
|
||||
class AttentionConfig:
|
||||
"""
|
||||
Per-layer attention metadata.
|
||||
"""Per-layer attention metadata.
|
||||
|
||||
NOTE(djsaunde): Constructed on every forward pass (not once per layer) since
|
||||
it can be invalid across passes (e.g. switching training/inference). Kept
|
||||
separate from AttentionContext to group params.
|
||||
NOTE(djsaunde): Rebuilt every forward pass (not once per layer) since it can
|
||||
go stale across passes (e.g. switching training/inference).
|
||||
"""
|
||||
|
||||
backend: str
|
||||
|
|
@ -102,13 +100,11 @@ def select_attention_backend(use_varlen: bool = False) -> str:
|
|||
def run_attention(
|
||||
*, config: AttentionConfig, context: AttentionContext, Q: Tensor, K: Tensor, V: Tensor
|
||||
) -> Tensor:
|
||||
"""
|
||||
Run attention using config / context info.
|
||||
"""Run attention using config / context info.
|
||||
|
||||
Backend priority (speed): FlashAttention if installed (varlen for packed
|
||||
inputs with `seq_info`, else dense), then xFormers, then SDPA as fallback.
|
||||
Varlen flash is preferred for packed batches as it avoids padding; xFormers
|
||||
and SDPA handle packing via a block-diagonal mask.
|
||||
Backend priority (speed): FlashAttention (varlen for packed inputs with
|
||||
`seq_info`, else dense), then xFormers, then SDPA. Varlen flash avoids
|
||||
padding for packed batches; xFormers and SDPA pack via a block-diagonal mask.
|
||||
"""
|
||||
|
||||
backend = config.backend
|
||||
|
|
@ -240,7 +236,6 @@ def run_attention(
|
|||
if local_mask.dtype == torch.bool:
|
||||
key_keep = local_mask
|
||||
else:
|
||||
# tokenizer attention_mask is typically int 0/1
|
||||
key_keep = local_mask != 0
|
||||
|
||||
past_len = k_len_local - q_len_local # works for prefill (0) and decode
|
||||
|
|
@ -260,7 +255,7 @@ def run_attention(
|
|||
|
||||
elif local_mask.dim() == 4:
|
||||
if local_mask.dtype != torch.bool:
|
||||
# Use boolean keep masks for better SDPA stability.
|
||||
# boolean keep masks are more stable in SDPA
|
||||
local_mask = local_mask.eq(0)
|
||||
else:
|
||||
raise ValueError(f"Unsupported SDPA attention_mask rank: {local_mask.dim()}")
|
||||
|
|
|
|||
|
|
@ -28,10 +28,8 @@ def get_model_info(
|
|||
model_id: str, properties: list[str] = ["safetensors", "lastModified"]
|
||||
) -> ModelInfo:
|
||||
"""
|
||||
Get the model info for a specific model.
|
||||
|
||||
Get info for a model. Defaults to minimal info; pass None for full info.
|
||||
properties: see https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/model_info
|
||||
Defaults to minimal info; set to None for the full model information.
|
||||
"""
|
||||
global _HFAPI
|
||||
if _HFAPI is None:
|
||||
|
|
@ -53,11 +51,8 @@ def list_models(
|
|||
limit: int = 10,
|
||||
) -> list[ModelInfo]:
|
||||
"""
|
||||
Retrieve model information from the Hugging Face Hub.
|
||||
|
||||
List models from the Hugging Face Hub. If full is True, properties is ignored.
|
||||
properties: see https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/list_models
|
||||
full: retrieve full model info; if True, properties is ignored.
|
||||
sort/author/search: sort order, model author, and search filter.
|
||||
"""
|
||||
global _HFAPI
|
||||
if _HFAPI is None:
|
||||
|
|
|
|||
|
|
@ -36,13 +36,11 @@ except Exception:
|
|||
_XFORMERS_MASK_CACHE_MAXSIZE = 32
|
||||
_XFORMERS_MASK_CACHE: OrderedDict[Tuple[Tuple[int, ...], int], Any] = OrderedDict()
|
||||
|
||||
# Cache per device for get_packed_info_from_kwargs to avoid repeated D2H sync across layers
|
||||
# Per-device caches avoid repeated D2H sync across layers
|
||||
_PACKED_INFO_CACHE: dict = {}
|
||||
|
||||
# Cache per device for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
|
||||
_SDPA_MASK_CACHE: dict = {}
|
||||
|
||||
# Cache per device for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
|
||||
_XFORMERS_BLOCK_MASK_CACHE: dict = {}
|
||||
|
||||
|
||||
|
|
@ -159,7 +157,7 @@ def enable_sample_packing(
|
|||
lengths = example.get(sequence_lengths_key)
|
||||
if isinstance(lengths, Iterable):
|
||||
seq_lengths.extend(int(length) for length in lengths)
|
||||
# Fallback: infer lengths from tokenized inputs when metadata is absent
|
||||
# fallback: infer lengths from tokenized inputs when metadata absent
|
||||
if not seq_lengths:
|
||||
for example in examples:
|
||||
ids = example.get("input_ids")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue