Restore kernel files to upstream to keep their explanatory comments

for more information, see https://pre-commit.ci

unsloth/__init__.py

@@ -16,11 +16,8 @@ import os, importlib.util, platform
 
 os.environ["UNSLOTH_IS_PRESENT"] = "1"
 
-# ── Windows console UTF-8 safety ─────────────────────────────────────────────
+# Force stdout/stderr to UTF-8 on Windows: legacy cp1252 consoles crash on
-# Legacy Windows consoles (cp1252) can't encode Unsloth's emoji/box-drawing
+# Unsloth's emoji/box-drawing glyphs. errors="replace" avoids unencodable-glyph crashes.
-# 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.
 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
+    # Transitional barrier: avoid importing unsloth_zoo on GPU hosts until
-    # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this
+    # unsloth_zoo.mlx is import-safe there.
-    # can collapse back to the centralized zoo runtime call below.
     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
+    # Older unsloth-zoo imports fine but lacks mlx.trainer/mlx.loader; give a
-    # mlx.trainer / mlx.loader submodules. Surface a friendly install hint
+    # friendly install hint instead of a raw submodule ImportError.
-    # instead of a raw ImportError on the submodule path.
     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
+    # Load raw_text helpers directly: dataprep/__init__.py imports torch via
-    # imports synthetic.py -> torch and would defeat the torch-free MLX path.
+    # 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"

unsloth/_gpu_init.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:

unsloth/chat_templates.py

@@ -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
+        # type_chat_template lets us swap the system message later
-        # Since it's not supported yet, we will raise an error first!
         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,
 ):
-    """
+    """Convert a dataset to ShareGPT style (1 input + 1 output field).
-    Converts a dataset to ShareGPT style (1 input + 1 output field).
+    `merged_prompt` merges multiple columns into the input;
-    Merge multiple columns into 1 input via `merged_prompt`; use
+    `conversation_extension` packs that many convos into one.
-    `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
     """
     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,
 ):
-    """
+    """Build an Ollama modelfile and HF Jinja template from a custom template.
-    Creates an Ollama modelfile and a HF Jinja template from a custom
+    Use {INPUT} and {OUTPUT} twice each; {SYSTEM} is optional.
-    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.
     """
     # 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,
 
 ):
-    """
+    """Apply a custom chat template to a dataset (builds the Ollama modelfile
-    Creates an Ollama modelfile and a HF Jinja template from a custom
+    and HF Jinja template). Use {INPUT} and {OUTPUT} twice each; {SYSTEM} is optional.
-    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.
     """
     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"):
-    """
+    """Check GGUF vs HF tokenization to catch tokenization bugs."""
-        Carefully checks the output of GGUF's tokenization and HF.
-        Can catch all tokenization bugs.
-    """
     import subprocess
     import re
     messages = [

unsloth/dataprep/raw_text.py

@@ -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 == input_ids for causal LM
-            labels = [list(ids) for ids in input_ids]
             return Dataset.from_dict(
                 {
                     "input_ids": input_ids,
@@ -125,13 +123,7 @@ class RawTextDataLoader:
         stride,
         return_tokenized = True,
     ):
-        """
+        """Chunk text with stride overlap; return tokenized chunks or text."""
-        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
-        """
         # 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))  # tokenizer returned a count
-            tokens = list(range(tokens))
 
         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):
-        """
+        """Compute dataset stats: lengths, encoding issues, repeats, empties."""
-        Check for:
-        - Minimum/maximum sequence lengths
-        - Character encoding issues
-        - Repeated content
-        - Empty chunks
-        """
         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")
 

unsloth/dataprep/synthetic.py

@@ -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)

unsloth/device_type.py

@@ -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.")

unsloth/import_fixes.py

@@ -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)
+    CUDA device printf (e.g. CUTLASS "Arch conditional MMA" on Blackwell) writes
-    writes to fd 1 at the C level, bypassing Python's sys.stdout, so the
+    to fd 1 at the C level, bypassing sys.stdout, so HidePrintMessage can't catch
-    HidePrintMessage filter can't catch it. Redirect fd 1 and 2 at the OS level,
+    it. Redirect fds 1/2 at the OS level, sync CUDA, then restore.
-    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
+    """Patch enable_input_require_grads to tolerate vision models that raise
-    that raise NotImplementedError from get_input_embeddings()."""
+    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
+    rng_state.pth on resume). Via a thread-local flag, hardens only the rng
-    flag, so it forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD)
+    torch.load: 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
+    and refuses torch < 2.6 (CVE-2025-32434); other torch.load calls untouched.
-    torch.load calls are untouched. No-op if transformers is absent or already
+    No-op if transformers is absent or already guards the load (>= 5.0.0rc3)."""
-    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
+    Operates on the raw importlib_version() string. Standard releases use
-    identifiers). Standard releases use +cu124/+rocm6.3/+cpu/+xpu; custom builds
+    +cu124/+rocm6.3/+cpu/+xpu; custom builds use +gitXXXX or other suffixes.
-    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:
+    Mapping formula:
-      torch 1.x  ->  torchvision 0.(x + 1)   (verified: torch 1.7 through 1.13)
+      torch 1.x -> torchvision 0.(x + 1)   (verified: torch 1.7 - 1.13)
-      torch 2.x  ->  torchvision 0.(x + 15)  (verified: torch 2.0 through 2.9)
+      torch 2.x -> torchvision 0.(x + 15)  (verified: torch 2.0 - 2.9)
-
-    Returns (tv_major, tv_minor) or None if the major version is unrecognized.
     """
     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():
-    """
+    """Re-add num_ctas/cluster_dims to triton CompiledKernel for torch.compile.
-    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).
 
-    The scope dict eagerly evaluates:
+    Triton 3.6.0+ dropped the direct `num_ctas`/`cluster_dims` attrs, but torch
-        binary.metadata.num_ctas, *binary.metadata.cluster_dims
+    2.9.x Inductor's make_launcher() still eagerly reads
-    when hasattr(binary, "metadata") is True, but metadata lacks cluster_dims.
+    binary.metadata.num_ctas/*cluster_dims (metadata lacks cluster_dims), crashing
-    This crashes before reaching the new launch path that doesn't need cta_args.
+    before the new launch path. Upstream fix pytorch/pytorch@97bd4db added hasattr
-
+    guards; we instead patch CompiledKernel.__init__ to inject the missing attrs.
-    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.
     """
     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_ to run fp16/bf16 init in fp32.
-    Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
 
-    torch.nn.init.trunc_normal_ can saturate at truncation bounds in fp16/bf16 on
+    trunc_normal_ can saturate at truncation bounds in fp16/bf16 on some
-    some versions/backends. This was observed in TorchTitan investigations where
+    versions/backends (https://github.com/pytorch/torchtitan/pull/2342). Avoid
-    low-precision truncation produced boundary-heavy initialization behavior:
+    it by initializing into a temporary fp32 tensor, then copying back.
-    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.
     """
     try:
         import torch
@@ -1138,16 +1122,11 @@ def patch_trunc_normal_precision_issue():
 
 
 def check_vllm_torch_sm100_compatibility():
-    """
+    """Raise a helpful error for the vLLM + torch < 2.9.0 + SM100 combination.
-    Check for incompatible vLLM + torch < 2.9.0 + SM100 (Blackwell) combination.
 
-    vLLM's distributed module (device_communicators) crashes with std::bad_alloc
+    vLLM's distributed module crashes with std::bad_alloc when imported on SM100
-    when imported on SM100 GPUs (B200/B100) with torch < 2.9.0. This is due to
+    GPUs (B200/B100) with torch < 2.9.0. Runs early (before vLLM import) to give a
-    C++ code in vLLM's NCCL/distributed layer being incompatible with older
+    clear message instead of the cryptic crash.
-    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 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 SM100 (Blackwell).
-    Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100).
 
-    The issue: vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL
+    vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL on SM90+, but
-    optimization on SM90+ GPUs. This fails on SM100 (B200/B100) during CUDA graph
+    it fails on SM100 (B200/B100) during CUDA graph capture (Triton's pipeliner
-    capture because Triton's pipeliner can't handle gdc_wait in complex kernels.
+    can't handle gdc_wait). See https://github.com/vllm-project/vllm/issues/30872
-
-    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
+    OpenEnv's EnvClient uses async reset/step. nest_asyncio makes nested event
-    these work via top-level await, but converted scripts need
+    loop calls work in both notebooks and converted scripts without replacing the
-    asyncio.get_event_loop().run_until_complete() wrappers. Applying nest_asyncio
+    original async methods (which would break existing sync wrappers).
-    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).
     """
     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
+    transformers and datasets detect torchcodec via find_spec, which returns True
-    returns True even when the native libs cannot dlopen. We flip their
+    even when the native libs can't dlopen. We flip their flags and seat a
-    flags and seat a sys.modules sentinel so downstream imports fall through
+    sys.modules sentinel so downstream imports hit their except ImportError paths.
-    their existing except ImportError handlers cleanly.
     """
     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
+    wandb can fail to import on a protobuf mismatch (e.g. wandb < 0.19.11 with
-    (e.g., wandb < 0.19.11 with protobuf >= 6.0). This causes cascading
+    protobuf >= 6.0), cascading through trl -> transformers/accelerate -> wandb.
-    import failures through trl -> transformers/accelerate -> wandb that
+    trl uses two separate is_wandb_available() functions
-    crash unsloth's import chain.
+    (transformers.integrations.integration_utils and accelerate.utils.imports);
-
+    both must be patched.
-    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.
     """
     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
+    ``Concatenate``/``ConversionOps`` MUST be real classes (peft subclasses them
-    subclasses them at module top); the rest only appear in runtime
+    at module top); the rest only appear in runtime calls gated behind
-    ``isinstance`` / construction calls gated behind ``is_transformers_ge_v5``."""
+    ``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
+    Must run BEFORE ``patch_peft_weight_converter_compatibility``, whose bare
-    function's bare ``except (ImportError, AttributeError): return`` would
+    ``except (ImportError, AttributeError): return`` would otherwise silently
-    otherwise silently no-op.
+    no-op. Idempotent and strictly additive (never overwrites real submodules).
-
-    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``).
 
     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
+    Mirrors the FlashAttention fallback: if import fails with a known binary
-    fails with a known binary symbol error, we disable it at startup so model imports do
+    symbol error, disable it at startup so model imports don't hard-fail.
-    not 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
+    AMD's Windows wheel ships one ``libbitsandbytes_rocm<NN>.dll`` whose suffix
-    suffix can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll),
+    can disagree with ``torch.version.hip`` (HIP 7.13 vs rocm72.dll), breaking the
-    breaking the native 4-bit/8-bit paths. Pin the installed suffix before
+    native 4/8-bit paths; pin the installed suffix before bitsandbytes is imported.
-    bitsandbytes is first imported.
 
-    No-op unless ALL of: Windows, a real HIP torch build (env hints like
+    No-op unless ALL of: Windows, a real HIP torch build (env hints don't count),
-    HIP_PATH do not count), a ROCm DLL installed, and no explicit user value.
+    a ROCm DLL installed, and no explicit user value. sitecustomize-seeded values
-    Linux is untouched. Values seeded by Studio's venv sitecustomize.py
+    are redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
-    (marked ``UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize``) 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
-    Make Accelerate's recursive utilities tolerate Unsloth's EmptyLogits
+    sentinel: recursively_apply returns it unchanged (no TypeError), and
-    sentinel. recursively_apply returns the sentinel unchanged instead of
+    find_device skips it while still finding real tensors, falling back to
-    raising TypeError, and find_device skips it while still finding real
+    PartialState().device only for sentinel-only payloads. Both wrappers are
-    tensors, falling back to PartialState().device only for sentinel-only
+    idempotent and propagated to every already-imported accelerate namespace.
-    payloads. Both wrappers are idempotent and are propagated to every
-    already imported accelerate namespace.
     """
     try:
         import accelerate.utils.operations as acc_ops

unsloth/kernels/moe/autotune_cache.py

@@ -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 for MoE kernels so tuning runs only once at training start."""
-Auto-tuning cache system for MoE kernels to ensure 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]:
-    """
+    """Return cached MoE kernel configs (config_fwd, config_bwd_dx, config_bwd_dw),
-    Get cached kernel configurations or run auto-tuning.
+    running auto-tuning if needed. force_autotune ignores existing caches."""
-
-    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)
-    """
     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]:
-    """
+    """'Safe Heuristic' kernel configs: safe on A100 (SM80), ~9x speedup on H100/B200."""
-    Get 'Safe Heuristic' kernel configurations.
-    These are verified to be safe on A100 (SM80) and provide ~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,

unsloth/kernels/moe/benchmark/benchmark_fused_moe.py

@@ -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,

unsloth/kernels/moe/benchmark/utils.py

@@ -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
 

unsloth/kernels/moe/grouped_gemm/interface.py

@@ -33,17 +33,14 @@ ch.setFormatter(formatter)
 logger.addHandler(ch)
 
 
-# Precompute TMA support to avoid graph breaks
+# Precomputed to avoid graph breaks. TMA needs GPU capability >= 9 (Hopper+) and a Triton TMA API.
-# TMA requires both:
-# 1. NVIDIA GPU with capability >= 9 (Hopper+)
-# 2. Triton version with TMA API (make_tensor_descriptor or _experimental_make_tensor_descriptor)
 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).
+    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
-    NOTE: We do NOT use torch.compiler.is_compiling() because it returns True during both
+    execution; we only want to skip kernels during tracing on fake tensors.
-    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:
+    MoE-specific fusions:
-    - `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.
+    - permute_x: fuse the token->grouped-expert-order permute of X (first GEMM).
-        - When `permute_x` is True, `X` is expected to be of shape (num_tokens, K).
+        True: X is (num_tokens, K). False: X is (total_tokens, K), total_tokens =
-        - 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.
+        num_tokens * topk, already sorted so each expert's tokens 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.
+    - permute_y: fuse the grouped-expert-order->token-order permute of the output
-    - `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.
+        (second GEMM).
-    - `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.
+    - 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`.
+    X: (M, K) hidden states; M = num_tokens if permute_x else total_tokens.
-    W: (E, N, K) expert weights, where E is number of experts, N in the intermediate (output) dim, and K is the reduction dim
+    W: (E, N, K) expert weights (E experts, N output dim, K reduction dim).
-    m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
+    m_sizes: tokens per expert = the M of each per-expert 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.
+    gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
-    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`).
+    topk_weights: (total_tokens,) routed-output weights, used only if fuse_mul_post.
-    use_fast_accum: currently unused; trade off faster accumulation dtype in GEMM for less precision.
+    use_tma_load_x: TMA load of activations, incompatible with permute_x.
-    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: TMA load of weights; prefer when TMA is supported (faster).
-    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: TMA store of output, incompatible with permute_y.
-    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
     """
 
     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
+    dX backward kernel. Shapes: dy is (NUM_TOKENS*TOPK, N) reduced over N,
-    grad_output: (M, N)
+    output dX is (NUM_TOKENS*TOPK, K) (per-expert grads are accumulated in a
-    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.
+    post-processing step).
-    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.
+    gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
-    `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.
+    m_sizes: tokens per expert = the M of each per-expert GEMM.
-    - In the forward pass, if we permuted X on load, we need to permute store in the backward pass
+    topk: experts chosen per token.
-    - Shapes
+    permute_x: whether X was permuted on load in the forward (first GEMM); if so we
-        - the forward pass input X shape is [NUM_TOKENS, K], reduce across K, output y is [NUM_TOKENS * TOPK, K]
+        permute on store here.
-        - the backward pass input dy shape is [NUM_TOKENS * TOPK, N], reduce across N, output dX is [NUM_TOKENS * TOPK, K]
+    permute_y: whether output was permuted on store in the forward (second GEMM); if
-    - 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.
+        so we permute on load here. dX is always stored contiguous.
-    `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.
+    fuse_mul_{pre,post}: must be False (inference only).
-    - 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
+    use_tma_load_dy: TMA load of dy, incompatible with permute_y.
-    - 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
+    use_tma_load_w: TMA load of weights; prefer when TMA is supported (faster).
-    `fuse_mul_{pre,post}`: always set to False since this should only be used for inference.
+    use_tma_store: TMA store of dX, incompatible with permute_x.
-    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.
     """
     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
+    # Flipped from the forward: permuting y on store means permuting on load here
-    # If we permuted y in the forward, we need to permute on load in the backward
     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.
+    # Output stays [NUM_TOKENS * TOPK, K] even when permute_x: per-expert grads are reduced in a later step.
-    # This will be done in a post-processing step reduction 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`.
+    dW backward kernel.
-    dY: (M, N)
+
-    topk: number of experts to choose per token.
+    X: (M, K) hidden states; M = num_tokens if permute_x else total_tokens.
-    m_sizes: tokens assigned to each expert which correspond to the size of M in the respective GEMMs in the grouped GEMM.
+    dY: (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.
+    topk: 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.
+    m_sizes: tokens per expert = the M of each per-expert GEMM.
-    - 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]
+    gather_indices: (total_tokens,) token indices per expert; slice by cumsum(m_sizes).
-    - in the backwards pass, we need to permute on load of X while loading dy in contiguous (expert grouped) order
+    permute_x: whether X was permuted on load in the forward (first GEMM); if so we
-    - since we are writing out dW, there is no need to permute on store
+        permute X on load here. dW never needs 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.
+    permute_y: whether output was permuted on store in the forward (second GEMM); if
-    - 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
+        so we permute dy on load here to match X's order.
-    - 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
+    use_tma_load_dy: TMA load of dy, incompatible with permute_y.
-    - since we are writing out dW, there is no need to permute on store
+    use_tma_load_x: TMA load of x, incompatible with permute_x.
-    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_store: TMA store of dW; prefer when TMA is supported (faster).
-    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.
     """
     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:
+    MoE-specific fusions:
-    - `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.
+    - permute_x: fuse the token->grouped-expert-order permute of X (first GEMM).
-        - When `permute_x` is True, `X` is expected to be of shape (num_tokens, K).
+        True: X is (num_tokens, K). False: X is (total_tokens, K), total_tokens =
-        - 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.
+        num_tokens * topk, already sorted by expert.
-    - `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.
+    - permute_y: fuse the grouped-expert-order->token-order permute of the output
-    - `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.
+        (second GEMM).
-
+    - fuse_mul_post: fuse multiply of routed output by topk_weights; only with
-    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`.
+        permute_y, and inference only (not training).
-    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.
 
+    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 (

unsloth/kernels/moe/grouped_gemm/kernels/autotuning.py

@@ -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_X = None,  # Auto-detect if None
-    TMA_LOAD_W = None,  # Auto-detect if not specified
+    TMA_LOAD_W = None,  # Auto-detect if None
-    TMA_STORE = False,  # NOTE: TMA_STORE is disabled for now
+    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_dY = None,  # Auto-detect if None
-    TMA_LOAD_W = None,  # Auto-detect if not specified
+    TMA_LOAD_W = None,  # Auto-detect if None
-    TMA_STORE = False,  # NOTE: TMA_STORE is disabled for now
+    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_dY = None,  # Auto-detect if None
-    TMA_LOAD_X = None,  # Auto-detect if not specified
+    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

unsloth/kernels/moe/grouped_gemm/kernels/backward.py

@@ -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
+    # TMA descriptors for loading sorted tokens. With TMA load we don't permute_x, so shape is
-    # When using TMA load, we don't permute_x, so shape should be [TOTAL_TOKENS, K]
+    # [TOTAL_TOKENS, K]. Single global descriptor with one block shape -- verify this doesn't error
-    # Also, we are defining a single global descriptor with single block shape
+    # when crossing expert boundaries.
-    # Need to check that this does not result in errors 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
 

unsloth/kernels/moe/grouped_gemm/kernels/forward.py

@@ -11,15 +11,10 @@ from .autotuning import (
 )
 
 
-#
+# PERMUTE_X -> permute X to expert order on load; PERMUTE_Y -> permute Y to token
-# PERMUTE_X -> permute tokens so that they are ordered by expert
+# order on store. Same permutation indices either way (load vs store).
-# PERMUTE_Y -> permute output so that they are ordered by token
+# FUSE_MUL -> multiply routed outputs by topk_weights (token order).
-# 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
+# Fusing mul assumes X in expert order while permuting Y -- checked in the interface.
-# 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
 @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
+    # TMA load implies no permute_x, so descriptor shape is [TOTAL_TOKENS, K].
-    # When using TMA load, we don't permute_x, so shape should be [TOTAL_TOKENS, K]
+    # Single global descriptor; may need checking across expert boundaries.
-    # 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
     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)
+                # Only two cases supported: (PERMUTE_X, not PERMUTE_Y) and (not PERMUTE_X, PERMUTE_Y).
-                # Hence, we can make the following simplifying assumptions when loading and storing
+                # Between them the load/store offsets and strides are flipped.
-                # Note the different strides between the two cases: the offsets for loading and storing are flipped and the strides must also be adjusted
                 if PERMUTE_X:
                     load_idx = (
-                        (expert_token_offsets // TOPK) * K
+                        expert_token_offsets // TOPK
-                    )  # Permute on load from token -> expert order, divide by TOPK to index from original number of tokens
+                    ) * K  # token -> expert order; //TOPK indexes the original tokens
-                    store_idx = indices_to_gather[:, None] * N  # Store in contiguous order
+                    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
+                # Hidden states are grouped by expert, so topk weights are always loaded in expert order
-                # In the pre-multiplication case, we multiply permuted hidden states by weights before the first gemm
+                # (pre-mul: before first gemm; post-mul: after second 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
                 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)
 

unsloth/kernels/moe/grouped_gemm/kernels/tuning.py

@@ -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

unsloth/kernels/moe/grouped_gemm/reference/layers/llama4_moe.py

@@ -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 (token->expert order); not in the autograd graph.
-        # Auxiliary structs; not recorded 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 (token->expert order); not in the autograd graph.
-        # Auxiliary structs; not recorded 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

unsloth/kernels/moe/grouped_gemm/reference/layers/qwen3_moe.py

@@ -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).
+        # Token counts + gather indices (token->expert order); aux structs, not in the autograd graph.
-        # Auxiliary structs; not recorded 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).
+        # Token counts + gather indices (token->expert order); aux structs, not in the autograd graph.
-        # Auxiliary structs; not recorded 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,

unsloth/kernels/moe/grouped_gemm/reference/moe_block.py

@@ -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
+        # Tokens per expert + token->expert gather indices.
-        # 1. Compute tokens per expert and indices for gathering tokes from token order to expert order
+        # Auxiliary structs; not recorded in the autograd graph.
-        # NOTE: these are auxiliary data structs which don't need to be recorded in 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
+        # Unpermute from expert order back to token order
-        # 1. Unpermute from expert order 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]
         )

unsloth/kernels/moe/grouped_gemm/reference/moe_ops.py

@@ -6,21 +6,13 @@ import torch.nn.functional as F
 
 
 def permute(X: torch.Tensor, gather_indices: torch.Tensor, topk: int):
-    """
+    """Reorder tokens by expert for grouped gemm.
-    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.
 
-    Helper for grouped gemm where hidden states need be ordered by expert.
+    X: [num_tokens, hidden_dim], gather_indices: [num_tokens * topk].
-    X: [num_tokens, hidden_dim]
+    Returns [total_tokens, hidden_dim] where total_tokens = num_tokens * topk.
-    sorted_token_idx: [num_tokens * topk]
-    topk: int
-
-    Returns:
-        [total_tokens, hidden_dim]
     """
     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,
 ):
-    """
+    """Run activation before topk (pre_act, e.g. llama4/deepseek) or after
-    If post_act is True, then activation function is run AFTER topk
+    (post_act, aligns with triton_bench)."""
-    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)
-    """
     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] (fwd) else [M, N]; W: [E, N, K]; m_sizes: [E].
-    X: [M, K] if forward, else [M, N]
+    Returns Y: [M, N] (fwd) else [M, K]."""
-    W: [E, N, K]
-    m_sizes: [E]
-
-    Returns:
-        Y: [M, N] if forward, 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
+            X_g = X[m_start:m_end]  # [m_size, K]
-            # m_size x K
+            W_g = W[g]  # [N, K]
-            X_g = X[m_start:m_end]
-            # N x K
-            W_g = W[g]
 
             # Y_g = X_g @ W_g.T -> [m_size, N]
             W_g = W_g.T if transpose else W_g

unsloth/kernels/moe/tests/common.py

@@ -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:

unsloth/kernels/moe/tests/moe_utils.py

@@ -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
+        # torch second_gemm is still permuted vs the fused one; compare via
-        # instead the hidden_states_unpermute should match since hidden_states_unpermute for the fused result is the same as second_gemm
+        # 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.
+    """Reference MoE block like Qwen3MoeGroupedGEMMBlock but with triton (not torch-native) grouped
-
+    gemm. NOT for production: saves intermediates and runs extra debug checks. See
-    Like Qwen3MoeGroupedGEMMBlock but with triton (not torch-native) grouped gemm.
+    grouped_gemm/reference/moe_block.py for a cleaner version.
-    NOT for production: it saves intermediate results and runs extra checks for
-    debugging. 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
+        # Token counts per expert + token-order -> expert-order gather indices
-        # gather indices (auxiliary, not recorded in the autograd graph).
+        # (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)

unsloth/kernels/moe/tests/test_grouped_gemm.py

@@ -43,8 +43,7 @@ from .common import (
 SEED = 0
 
 
-# Only certain (permute_x, permute_y, use_W1) combinations are valid; see the
+# Only certain (permute_x, permute_y, use_W1) combos are valid; see module string below for rationale.
-# module string below for the full 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
+        # topk weights are in token order, so unpermute both before multiplying when permute_y is False
-        # the ref output needs to be unpermuted before multiplying by topk weights since topk weights are in token order
         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
+    # For the first GEMM with permute_x the kernel unpermutes dX, so unpermute ref_grad to match
-    # therefore we need to unpermute the ref_grad to compare to the output of the kernel
     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
+        # Show that the topk reduction introduces diffs vs autograd
-        # First calculate X.grad manually by backpropping through unpermuted ref_grad
         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
+# Llama4 configs are shrunk to avoid OOM; the full size (5120, 8192) shows diffs ~1e-2.
-# Important to note that for the full model size (5120, 8192), the tests do result in diffs on the order of 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
+    # permute_y means grouped_gemm unpermuted on store, so unpermute before backprop to align
-    # Therefore we have to unpermute before backpropping to ensure proper alignment
     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:

unsloth/kernels/moe/tests/test_llama4_moe.py

@@ -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)

unsloth/kernels/moe/tests/test_qwen3_moe.py

@@ -113,9 +113,7 @@ def test_qwen3_moe(
     permute_y: bool,
     autotune: bool,
 ):
-    torch.manual_seed(
+    torch.manual_seed(SEED)  # Redundant under pytest -- conftest.py has an autouse fixture
-        SEED
-    )  # Should not be needed when running using pytest -- autouse fixture in conftest.py
     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,

unsloth/models/__init__.py

@@ -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__

unsloth/models/_utils.py

@@ -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.
+    Apply gradient checkpointing with smart heuristics, returning the effective
-
+    setting (may downgrade "unsloth" to True). For seq < 512, "unsloth" offloading
-    For seq < 512, gc="unsloth" offloading overhead isn't worth it; standard gc is faster.
+    overhead isn't worth it, so standard gc is used instead.
-
-    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)
     """
     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.
+    Intercept warnings.warn and replace matching messages with Unsloth ones, via
-    Uses a list of registered (match_text, replacement, category) rules checked in order.
+    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):
-    """
+    """Param count of a quantized model (Params4bit counted as 2x numel)."""
-    Calculate quant model param count based on difference in param class. Returns int for param count.
-    """
     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):
-    """
+    """Total model param count; if trainable_only, count only params requiring grads."""
-    Calculate model's total param count. If trainable_only is True then count only those 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
+    # Needed for configs that use layer_type_validation (e.g. Gemma-2).
-    # for eg Gemma-2, so we must import it to stop errors.
     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,
+# Skip exec-based config patching on transformers 5.x: its @strict/@auto_docstring/
-# and interval() in config classes. exec(inspect.getsource(...)) fails because
+# interval() config symbols aren't in scope for exec(getsource(...)), and v5 configs
-# those symbols are not in scope. Skip the exec-based config patching for 5.x
+# already use rope_parameters (the rope_scaling replacement).
-# since those configs already use rope_parameters (the v5 replacement for
-# rope_scaling).
 _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.
+    # Log basic env stats by downloading a public README.md from HF (checks for broken/down envs).
-    # We simply download a README.md file from HF - all data is made public.
+    # Disable by commenting the below out.
-    # This is simply so we can check if some envs are broken or not.
-    # You can disable this 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.
+    # Log basic env stats by downloading a public README.md from HF (also detects if HF is down).
-    # This is also to check if HuggingFace is down or not!
+    # Disable via UNSLOTH_DISABLE_STATISTICS.
-    # 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
     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.
+    Untie input/output embeddings in-place (so they can be quantized differently).
-    This is useful if we want to quantize the input/ouput embeddings differently.
-    Model is modified in-place.
     """
 
     # 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]:
-    """
+    """Yield FQNs of modules matching filter_fn(module, fqn) -> bool."""
-    Given a model and a filter function (m, fqn) -> bool,
-    yield fully qualified names (FQNs) of modules that match.
-    """
     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.
+    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)
-    On a high level, this means fake quantizing the base (frozen) model during training.
+    to reduce post-training quantization degradation. Combinable with LoRA.
-    Fake quantization refers to simulating quantization numerics in high precision (e.g. bf16).
+    See https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700
-    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
     """
     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
+    For torchao-quantized models, return torch.no_grad() instead of
-    the following error when calling ops like aten.t() in inference mode.
+    torch.inference_mode(), since ops like aten.t() on tensor subclasses hit a
-    This is a bug in PyTorch that affects all tensor subclasses.
+    PyTorch bug ("Cannot set version_counter for inference tensor").
-
+    See https://github.com/pytorch/pytorch/issues/164872
-        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()`.
     """
     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 (or config) is a Mixture of Experts (MoE) model."""
-    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
-    """
     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.
+    Resolve the parameter path for MoE expert weights. Most models use
-
+    ``mlp.experts.*``; Gemma4 uses ``experts.*``. Prefer whichever exists on the
-    Most current Unsloth MoE models expose expert weights under
+    loaded module.
-    ``mlp.experts.*``. Gemma4 stores them directly under ``experts.*``.
-    Prefer the path that exists on the loaded module when possible.
     """
     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.
+    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
-    For MoE models, returns the parameter paths for expert weights
+    target_modules are included:
-    (gate_up_proj, down_proj) that should be targeted by PEFT's
+    - "down_proj" -> "<prefix>.experts.down_proj"
-    target_parameters for LoRA on nn.Parameter. The exact parameter path
+    - "gate_proj"/"up_proj"/"gate_up_proj" -> "<prefix>.experts.gate_up_proj"
-    depends on the model layout, for example ``mlp.experts.*`` or
+    The prefix depends on layout (``mlp.experts.*`` or ``experts.*``).
-    ``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
     """
     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):
-    """
+    """Wrap model.generate to reject vLLM-style usage when fast_inference=False."""
-    Creates a wrapper around model.generate that checks for incorrect
-    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

unsloth/models/cohere.py

@@ -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

unsloth/models/diffusion.py

@@ -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: transformers-only slow path for text-diffusion models (e.g. DiffusionGemma).
-FastDiffusionModel: a 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
+These models use a block-diffusion sampling loop and a novel backbone, so we skip Unsloth's
-Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay
+autoregressive kernel/compile patching and load the unmodified HF model (outputs stay bit-identical to
-bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the
+transformers), keeping only safe conveniences: 4bit/8bit loading, PEFT LoRA, the (model, tokenizer)
-(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land.
+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(

unsloth/models/falcon_h1.py

@@ -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,
 ):
-    """
+    """Fast inference using the 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 splits into 4 chunks; the mask zeroes Qk^T and softmax is row-wise, so
-    [QK^T, Qk^T]
+    softmax(QK^T)V is just the prior step's attention. We therefore only compute
-    [qK^T, qk^T]
+    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
-    Since the attention mask wipes Qk^T, we just get
-    [QK^T,    0]
-    [qK^T, qk^T]
-
-    Since softmax is row-wise, we get
-    softmax([QK^T,    0])
-    softmax([qK^T, qk^T])
-
-    We then multiply by   [V]
-                          [v]
-    softmax([QK^T,    0]) [softmax(QK^T)V] *
-    softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
-
-    But notice * [softmax(QK^T)V] is just the last attention.
-    We just need to compute the last final row.
-
-    This means we can pass in a row of Q, but we need to
-    remember K and V, which are called the KV cache.
     """
     Xn = hidden_states
     bsz, _, hd = hidden_states.size()
@@ -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
+    # Extend 2 steps ahead to avoid errors on short KV cache
-    # or else error
     self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, 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]]]:
-    """
+    """FalconH1 decoder layer: mamba + attention mixer, then SwiGLU MLP, with residuals."""
-    Args:
-        hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-        attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-            `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-        output_attentions (`bool`, *optional*):
-            Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-            returned tensors for more detail.
-        use_cache (`bool`, *optional*):
-            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-            (see `past_key_values`).
-        past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-    """
     if use_cache and hasattr(self, "_flag_for_generation"):
         residual = hidden_states
         hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
@@ -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
 

unsloth/models/gemma.py

@@ -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

unsloth/models/gemma2.py

@@ -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
+        # Pipeline parallelism: move tensors to this layer's device (once per GPU)
-        # note that this movement is once per GPU in PP
         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

unsloth/models/glm4_moe.py

@@ -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:
+Differences from Qwen3 MoE: sigmoid router (not softmax), routed_scaling_factor 1.8, 1 shared expert
-- Router uses sigmoid activation (not softmax)
+processing all tokens, group-based selection before topk, and MLA (Multi-head Latent Attention).
-- 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)
 """
 
 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
+    # Import first to apply the TMA compatibility shim (old + new TMA API names)
-    # for both old and 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 (sigmoid router + 1 shared expert)."""
-    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
-    """
     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. hidden_states [num_tokens, hidden_dim],
-    Optimized expert forward using grouped GEMM.
+    top_k_index/top_k_weights [num_tokens, top_k] -> [num_tokens, hidden_dim]."""
-
-    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
-    """
     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 (2-3x MoE throughput
-    Fast GLM-4.7 Flash (GLM4 MoE Lite) model with grouped GEMM optimization.
+    via grouped GEMM, fused permutation, and optimized RMS LayerNorm / SwiGLU)."""
-
-    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
-    """
 
     @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
+        # Patch MoE forward with grouped GEMM (TMA compat in grouped_gemm/__init__.py)
-        # grouped_gemm/__init__.py)
         if HAS_GROUPED_GEMM:
             Glm4MoeLiteNaiveMoe.forward = Glm4MoeLiteNaiveMoe_fast_forward
             Glm4MoeLiteMoE.forward = Glm4MoeLiteMoE_fast_forward

unsloth/models/granite.py

@@ -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,
+        # Stash config so GraniteModel_fast_forward_inference can reach residual_multiplier. See:
-        # so stash the whole config here to pass it around. 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 for torch.compile failing on the embedding matrix (torch < 2.2)
-        # Workaround randomnly fixes it for torch versions < 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?
+        # Patch all dtypes - BnB defaults to float16 instead of the correct type
-        # BnB default dtype seems to be float16!
         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  # BnB defaults to float16
-                    module.weight.quant_state[2] = correct_dtype  # Cast to correct dtype
                 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):

unsloth/models/llama.py

@@ -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).
+    float16 is promoted to float32 (Tesla T4). ``offload_device`` only supports
-    ``offload_device`` currently only supports "cpu"; None leaves the frozen
+    "cpu"; None leaves the frozen module in place. Modifies ``module`` in-place.
-    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 not float16. See unslothai/unsloth#1200
-        # Tesla T4 must use float32 and not float16
         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.
+    [Q, q] @ [K, k].T splits into 4 chunks; the mask wipes Qk^T so only the new
-    [QK^T, Qk^T]
+    row [qK^T, qk^T] needs computing (the rest is the prior attention). Hence we
-    [qK^T, qk^T]
+    pass one row of Q but must remember K and V (the KV cache).
-
-    Since the attention mask wipes Qk^T, we just get
-    [QK^T,    0]
-    [qK^T, qk^T]
-
-    Since softmax is row-wise, we get
-    softmax([QK^T,    0])
-    softmax([qK^T, qk^T])
-
-    We then multiply by   [V]
-                          [v]
-    softmax([QK^T,    0]) [softmax(QK^T)V] *
-    softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
-
-    But notice * [softmax(QK^T)V] is just the last attention.
-    We just need to compute the last final row.
-
-    This means we can pass in a row of Q, but we need to
-    remember K and V, which are called the KV cache.
     """
     Xn = hidden_states
     bsz, _, hd = hidden_states.size()
@@ -772,19 +748,7 @@ def LlamaDecoderLayer_fast_forward(
     *args,
     **kwargs,
 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-    """
+    """Fast decoder-layer forward; hidden_states is `(batch, seq_len, embed_dim)`."""
-    Args:
-        hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-        attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-            `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-        output_attentions (`bool`, *optional*):
-            Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-            returned tensors for more detail.
-        use_cache (`bool`, *optional*):
-            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-            (see `past_key_values`).
-        past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-    """
     if use_cache and hasattr(self, "_flag_for_generation"):
         residual = hidden_states
         hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
@@ -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
+    # Zero out attention mask elements, specifically for DPO
-    # 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")

unsloth/models/loader.py

@@ -104,10 +104,8 @@ from ._utils import (
     set_task_config_attr,
 )
 
-# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers
+# Re-export FORCE_FLOAT32 from unsloth_zoo (single source of truth); fallback list
-# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working.
+# below keeps import working when unsloth_zoo is older than unsloth.
-# Fallback list mirrors zoo for users who upgrade unsloth without upgrading
-# unsloth_zoo (so this module never fails at import).
 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)  # Login to allow private models
-        token = hf_login(token)
         # 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 64)
-        # AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 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 64)
-            # AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 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
+        # Tiled MLP: set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"
-        # to turn on 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)  # Login to allow private models
-        token = hf_login(token)
         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 64)
-        # AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 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 64)
-            # AMD Instinct GPUs need blocksize = 128 on bitsandbytes < 0.49.2 (our pre-quants use blocksize = 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
+        # Tiled MLP: set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"
-        # to turn on 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"
         )

unsloth/models/loader_utils.py

@@ -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)
+    Disable bnb 4bit/8bit loading if the model already has a non-bnb quant config,
-    because the model already has a non-bitsandbytes quantization config.
+    to avoid config conflicts. Returns (load_in_4bit, load_in_8bit, quant_method),
-    If so, disable BOTH 4bit and 8bit loading and print a warning message.
+    where the flags are False if disabled and quant_method is the detected method or None.
-
-    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
     """
     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)
+    # Non-bnb quant config (compressed-tensors/gptq/awq): disable bnb to avoid config conflicts
-    # We should disable BOTH bitsandbytes loading 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(

unsloth/models/mapper.py

@@ -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()
 

unsloth/models/mistral.py

@@ -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
+            # Keep 2D attention_mask: DPO embed-masking nulls it before attention,
-            # attention_mask; it nulls the mask before attention anyway, so
+            # and a 4D conversion would crash DPO.
-            # leaving it 2D is safe and avoids a 4D conversion that crashes 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
+    # Merge legacy/new spellings; skip int max() if either is a tensor (HF selective-decode form)
-    # last-token slice fires on the normal path too. 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
+        # Retain old rotary embeddings: static KV Cache (4.38.0) made training much slower.
-        # Static KV Cache was introduced in 4.38.0, causing training to be much slower.
+        # https://github.com/unslothai/unsloth/issues/168
-        # Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
         # https://github.com/huggingface/transformers/pull/27931
         # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
         import transformers.models.mistral.modeling_mistral

unsloth/models/qwen2.py

@@ -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
+        # Retain old rotary embeddings: static KV cache (4.38.0+) slowed training. Solves issue #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.
         # 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

unsloth/models/qwen3.py

@@ -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
+    # extend 2 steps ahead before the short KV cache fills, else error
-    # or 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,

unsloth/models/qwen3_moe.py

@@ -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]

unsloth/models/rl.py

@@ -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
+        # Snapshot the GC mode before TRL's unwrap CM calls gradient_checkpointing_disable();
-        # gradient_checkpointing_disable() before yielding; preserve the actual
+        # keep the real value (e.g. "unsloth") not a bool so the finally restore matches.
-        # mode value (e.g. "unsloth") rather than collapsing it to a bool, so
-        # the finally restore matches the caller's configured GC mode.
         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):
-        """
+        """Evaluation step on `model` using `inputs`.
-        Perform an evaluation step on `model` using `inputs`.
+        Returns (loss, logits, labels), each optional.
-        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).
         """
         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
+    # Must run before patch_trl_rl_trainers: it imports more trl.* submodules, and any
-    # patch_trl_rl_trainers, which imports extra trl.* submodules; any module
+    # imported after the sys.modules walk would keep the original broken binding.
-    # 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.
     patch_trl_disable_gradient_checkpointing()
     patch_trl_rl_trainers()
     patch_trl_openenv()

unsloth/models/rl_replacements.py

@@ -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
+    """Return final_logit_softcapping for a config, falling back to the nested text sub-config for
-    nested text sub-config for composite models. Handles both:
+    composite models (Gemma-4 ``config.text_config`` or T5Gemma ``config.get_text_config()``).
-      - Gemma-4-style configs where the attribute lives on ``config.text_config``
+    Returns 0 if unset, matching previous behaviour.
-      - 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.
     """
     softcap = getattr(config, "final_logit_softcapping", None)
     if softcap is None:

unsloth/models/sentence_transformer.py

@@ -152,10 +152,7 @@ def _save_pretrained_gguf(
     maximum_memory_usage = 0.85,
     **kwargs,
 ):
-    """
+    """Save the SentenceTransformer to GGUF: convert the inner transformer and place the GGUF files in save_directory."""
-    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.
-    """
     # 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,
 ):
-    """
+    """Convert the SentenceTransformer to GGUF and push it to the Hugging Face Hub, returning the full repo ID.
-    Converts the SentenceTransformer model to GGUF format and pushes to the Hugging Face Hub.
 
-    This method:
+    quantization_method (str or list) selects the GGUF method(s):
-    1. Saves the model locally to a temporary directory in GGUF format.
+        * "not_quantized"  : Fast conversion, slow inference, big files.
-    2. Uploads the GGUF files, config, Ollama Modelfile, and README to the Hub.
+        * "fast_quantized" : Fast conversion, OK inference, OK file size.
-    3. Cleans up the temporary directory.
+        * "quantized"      : Slow conversion, fast inference, small files.
-
+        * "f32" / "f16"    : Full accuracy, slow and memory hungry.
-    Args:
+        * "q8_0"           : Fast conversion, high resource use.
-        repo_id (str): The Hugging Face Hub repo ID (e.g., "username/model-name").
+        * "q4_k_m" / "q5_k_m" : Q6_K for half the attention.wv/feed_forward.w2 tensors, else Q4_K/Q5_K.
-        tokenizer: The tokenizer to save. Defaults to `self.tokenizer`.
+        * "q2_k"           : Q4_K for attention.vw/feed_forward.w2, Q2_K elsewhere.
-        quantization_method (str or list): GGUF quantization method(s). Can be a string or list of strings.
+        * "q3_k_l" / "q3_k_m" : Q5_K/Q4_K for attention.wv/wo/feed_forward.w2, else Q3_K.
-            Choose from the following options:
+        * "q3_k_s" / "q4_k_s" / "q5_k_s" : Q3_K/Q4_K/Q5_K for all tensors.
-            * "not_quantized"  : Recommended. Fast conversion. Slow inference, big files.
+        * "q4_0" / "q4_1" / "q5_0" / "q5_1" : 4/5-bit, increasing accuracy and cost.
-            * "fast_quantized" : Recommended. Fast conversion. OK inference, OK file size.
+        * "q6_k"           : Q8_K for all tensors.
-            * "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.
     """
     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
+        """Write base config.json next to adapter_config.json so PEFT adapter
-        AutoProcessor, which falls back to AutoConfig for tokenizer-only
+        checkpoints reload: sentence-transformers >= 5.4 reloads Transformer
-        roots -- so PEFT adapter checkpoints still need base config.json
+        modules via AutoProcessor, which falls back to AutoConfig."""
-        next to adapter_config.json."""
         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 ``(modules, no_modules_json)``."""
-        Returns:
-            tuple[OrderedDict, bool]: (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.
+        (1.2x safety margin) from empirical benchmarks. Optional batch_size /
-
+        grad_accum / max_seq_length give a coarse pre-run adjustment."""
-        Optional batch_size / grad_accum / max_seq_length give a coarse,
-        conservative pre-run adjustment with no runtime measurements.
-        """
         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 auto-apply torch.compile when training
-    Patch SentenceTransformerTrainer to automatically apply torch.compile
+    steps exceed the breakeven threshold. Called on module import."""
-    when training steps exceed the breakeven threshold.
-
-    This is called automatically when this module is imported.
-    """
     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)

unsloth/models/vision.py

@@ -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 cache if sliding window seen, otherwise try static
-    # Use hybrid 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
+    Some VLMs (e.g. LFM2.5-VL) have tokenizer_class entries AutoTokenizer cannot
-    cannot resolve. This function loads the image processor and tokenizer separately,
+    resolve; load the image processor and tokenizer separately and assemble them.
-    sets required special token attributes, and constructs the processor.
     """
     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").
+            # model_type may be a sub-type (e.g. "lfm2" vs "lfm2_vl"); top-level config often maps
-            # Try the top-level config.model_type which often has the processor mapping.
             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;
+        # Resolve text-only before is_vlm/vLLM checks; skip vision tower only for
-        # skip the vision tower only for families with their own text decoder (Gemma 3). #5816
+        # 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 models were already redirected to a BF16 sibling; sync model_name here
-        # FP8 weights. We just need to update it here for sanity.
         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.
+        # Nemotron-H Mamba fused kernels pass out_proj.weight to F.linear,
-        # Mamba fused kernels pass out_proj.weight directly to F.linear,
+        # which fails on quantized Params4bit; skip out_proj from quantization.
-        # which fails with 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.
+            # Transformers 5.x @strict config classes reject unexpected kwargs; set them on config
-            # Move config-level attributes onto the config object directly.
             _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),
+        # If processor loading failed or AutoProcessor degraded to a text-only
-        # or if AutoProcessor silently degraded to a text-only tokenizer
+        # tokenizer instead of a full VLM processor (#4085), build it manually.
-        # instead of returning a full VLM processor (issue #4085),
-        # try constructing the processor manually from separate components.
         _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
+        # Backwards compat: copy chat_template from inner tokenizer when processor lacks one
-        # chat_template.jinja) but the inner tokenizer does, copy it to the processor.
         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.
+            # Some VLM processors (e.g. ERNIE VL) fail patching; retry via AutoTokenizer
-            # Try loading tokenizer separately via AutoTokenizer as fallback.
             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,
+        # Audio feature extractors must stay right padded: left padding (a text
-        # forwarded by from_pretrained) shifts Whisper mels and desyncs Gemma 4
+        # setting forwarded by from_pretrained) shifts Whisper mels and desyncs
-        # audio token counts (crash on transformers < 5.10).
+        # 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.
+        # Gemma4 ClippableLinear wraps nn.Linear; PEFT can't inject LoRA directly,
-        # Monkey-patch PEFT to target the inner .linear child instead.
+        # 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
+        # Gemma3N audio conformer's variable-length tensors cause stride mismatches
-        # that cause stride mismatches in AOT autograd compiled backward
+        # in AOT autograd compiled backward under non-reentrant checkpointing. TRL/notebook
-        # when non-reentrant checkpointing is used. The notebook or TRL
+        # may later set use_reentrant=False, so intercept gradient_checkpointing_enable
-        # may override gradient_checkpointing_kwargs with use_reentrant=False
-        # after this point, so we 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
+        # Restore use_cache that prepare_model_for_training disabled for gradient checkpointing
-        # for gradient checkpointing (older unsloth_zoo has no restore helper)
         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
+        # Re-disable use_cache if for_inference restored it (record exists only after a disable)
-        # and for_inference restored it (record only exists 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
+    Validate local video paths in a dataset exist, catching missing files before
-    missing files before training (torchvision otherwise returns an empty
+    training (torchvision otherwise silently yields an empty tensor). Returns the
-    tensor and the model silently receives no video signal).
+    list of missing paths (empty when all exist).
 
     Args:
-        dataset:     Map-style Dataset, list of dicts, or iterable of examples
+        dataset:     Map-style Dataset / list / iterable (not a streaming
-                     (not a streaming IterableDataset - iterating consumes it).
+                     IterableDataset - iterating consumes it).
-        column:      Chat-messages column, default "messages"; "conversations",
+        column:      Chat-messages column ("messages"); "conversations", "prompt"
-                     "prompt" and "completion" are also scanned.
+                     and "completion" are also scanned.
-        raise_error: True (default) raises FileNotFoundError listing missing
+        raise_error: True raises FileNotFoundError on missing files; False warns.
-                     files; False warns and returns them.
         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
+    # Report each missing path once; only existing paths enter `checked`, so
-    # `checked`, so retries after an error re-check previously missing files.
+    # retries after an error re-check previously missing files.
     seen_missing = set()
     if checked is None:
         checked = set()

unsloth/ollama_template_mappers.py

@@ -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

unsloth/optimizers/q_galore_adamw.py

@@ -46,29 +46,19 @@ def _require_bnb():
 
 
 class QGaLoreAdamW8bit(Optimizer2State):
-    """AdamW optimizer with 8-bit states, GaLore low-rank gradient projection,
+    """AdamW with 8-bit states, GaLore low-rank gradient projection, and optional
-    and optional INT8 weight quantization.
+    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
+    Param group keys: GaLore uses ``rank``, ``update_proj_gap``, ``scale``,
-       moments are stored in 8-bit, reducing optimizer state memory by ~4×.
+    ``proj_type``, ``quant``, ``quant_group_size``, ``quant_n_bit``,
-
+    ``cos_threshold``, ``gamma_proj``, ``queue_size``; weight quantization uses
-    2. **GaLore low-rank gradient projection** — gradients are projected into a
+    ``weight_quant``, ``stochastic_round``, ``weight_group_size``.
-       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``
     """
 
     def __init__(
@@ -101,17 +91,10 @@ class QGaLoreAdamW8bit(Optimizer2State):
 
     @torch.no_grad()
     def step(self, closure = None):
-        """Perform a single optimization step.
+        """Single optimization step. For each ``rank``-group parameter: (1)
-
+        dequantize INT8 weight if ``weight_quant``; (2) project gradient to
-        For each parameter that has a ``rank`` key in its param group, the
+        low-rank; (3) 8-bit Adam update in low-rank space; (4) project back and
-        following sequence is executed:
+        add to the saved weight; (5) re-quantize to INT8 if ``weight_quant``."""
-
-        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.
-        """
         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
+                    # Disable weight decay here; reapplied manually after project-back.
-                    # (we apply it 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
+                    # Scalar placeholder frees float memory; install_weight_quant_hooks
-                    # pre-hook (install_weight_quant_hooks) dequantizes before
+                    # forward pre-hook dequantizes before the next forward pass.
-                    # 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
+        **Weights are NOT converted to uint8 here** — they stay float so the first
-        the optimizer knows to quantize/dequantize them during ``step()``.
+        forward/backward runs correctly; actual quantization happens at the end of
-        **Weights are NOT converted to uint8 here** — they remain in float
+        the first ``step()``.
-        so that the first forward/backward pass runs correctly.  The actual
-        quantization happens at the end of the first ``step()`` call.
         """
         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
+                # Tag only; first step() quantizes after the update. Dummy
-                # step() quantizes after the update. Dummy scales/zeros keep
+                # scales/zeros keep _has_weight_quant() True on the first step.
-                # _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
+    Parameters matching ``target_modules`` (or the default attention/MLP
-    and MLP projection names) are placed in the GaLore group.  All other
+    projection names) go in the GaLore group; all other trainable params go in
-    trainable parameters go into the non-GaLore group.
+    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]``.
     """
     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
+        # Exclude 1-D params (biases, norms): GaLoreProjector.project needs 2-D grads.
-        # GaLoreProjector.project requires 2-D gradients.
         name_parts = name.split(".")
         is_galore = param.dim() >= 2 and any(t in name_parts for t in targets)
 

unsloth/optimizers/q_galore_projector.py

@@ -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
+        SVD is recomputed every ``update_proj_gap`` steps (subject to adaptive
-        adaptive scheduling).  Between recomputations the cached orthogonal
+        scheduling); between recomputations the cached orthogonal matrix is reused.
-        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.
         """
         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.
+        """Project a low-rank update back to full rank, scaled by ``self.scale``."""
-
-        Args:
-            low_rank_grad: The low-rank gradient/update tensor.
-
-        Returns:
-            The full-rank update 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.
+        """Top-``rank`` orthogonal matrix of 2-D ``weights`` via truncated SVD.
-
+        ``side='left'`` returns U[:, :rank] shape ``(M, rank)``; ``'right'``
-        Args:
+        returns Vh[:rank, :] shape ``(rank, N)``."""
-            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).
-        """
         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
+    Rounding direction is chosen probabilistically by the fractional part,
-    probabilistically proportional to the fractional part. This gives an
+    giving an unbiased estimator in expectation.
-    unbiased estimator of the original value in expectation.
+    Returns ``(quantized_uint8, scales, zeros, original_shape)``.
-
-    Returns:
-        ``(quantized_uint8, scales, zeros, original_shape)``
     """
     org_shape = w.shape
     if q_group_size > 0:

unsloth/registry/__init__.py

@@ -32,11 +32,8 @@ def search_models(
     quant_types: list[QuantType] = None,
     search_pattern: str = None,
 ) -> list[ModelInfo]:
-    """
+    """Query the registry for ModelInfo. search_pattern matches the full HF
-    Get model info from the registry. See registry.ModelInfo for more fields.
+    hub model_id (model_path)."""
-
-    search_pattern is matched against the full model path (the HF hub model_id).
-    """
     if not _ARE_MODELS_REGISTERED:
         register_models()
 

unsloth/registry/_deepseek.py

@@ -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)
 

unsloth/registry/_gemma.py

@@ -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",

unsloth/registry/_llama.py

@@ -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",

unsloth/registry/_phi.py

@@ -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],

unsloth/registry/_qwen.py

@@ -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",

unsloth/registry/registry.py

@@ -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,

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
+# Added in unsloth-zoo PR #526; may not exist on older versions
-# and 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.
+    # Delete the cached repo matching this model; saves ~4GB on Kaggle.
-    # Can save 4GB of disk space - useful for Kaggle systems.
     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
+    Merge paths may mutate tokenizer metadata after writing. E.g. Gemma 4 instruct
-    is written. Gemma 4 instruct models use `<turn|>` as their chat EOS token;
+    uses `<turn|>` as chat EOS; if the config is reset to the base `<eos>`, vLLM
-    if tokenizer_config.json is reset to the raw base `<eos>` token, runtimes such
+    won't stop generation correctly. Best-effort, never fails the save.
-    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.
 
-    `filename_prefix` mirrors the same argument on Transformers'
+    `filename_prefix` mirrors Transformers' save_pretrained: when set, writes
-    `PreTrainedTokenizerBase.save_pretrained`: when provided, the tokenizer
+    `{filename_prefix}-tokenizer_config.json` instead of `tokenizer_config.json`.
-    config is written as `{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
+        # Org push: .save_pretrained doesn't work, so save locally then upload
-        # locally first then upload manually.
         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!
+    # Download the 10th latest release since the latest might be broken (fallback)
-    # FALLBACK mechanism
     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
+    # Clone a specific commit; don't use the GPU
-    # Also 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,
 ):
-    """
+    """Orchestrate GGUF conversion: install, convert, and quantize."""
-    Orchestrates the complete GGUF conversion process.
-    Handles installation, conversion, and quantization.
-    """
-    # print_output True only if UNSLOTH_ENABLE_LOGGING=1
     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
+        first_conversion = "None"  # GPT-OSS isn't quantized
-        # Only keep one conversion method since GPT-OSS doesn't quantize
         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
+    Like .save_pretrained(...) but auto-converts 4bit weights to float16.
-    converted to float16 with as few overhead as possible.
+    `save_method`:
-
-    Choose for `save_method` to be either:
     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
+    Like .push_to_hub(...) but auto-converts 4bit weights to float16.
-    converted to float16 with as few overhead as possible.
+    `save_method`:
-
-    Choose for `save_method` to be either:
     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
+    Like .save_pretrained(...) but auto-converts 4bit weights to float16, then to
-    converted to float16 then converted to GGUF / llama.cpp format.
+    GGUF / llama.cpp format. `quantization_method`:
-
-    Choose for `quantization_method` to be:
     "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
+    # Normalize quantization_method (old-style) to a list
-    # Use old style quantization_method
     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_vlm = is_vlm,
-            is_gpt_oss = is_gpt_oss,  # Pass gpt_oss Flag
+            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
+    Like .push_to_hub(...) but auto-converts 4bit weights to float16, then to
-    converted to float16 then converted to GGUF / llama.cpp format.
+    GGUF / llama.cpp format. `quantization_method`:
-
-    Choose for `quantization_method` to be:
     "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
+    # Save each type individually (inefficient: F16/BF16 saved repeatedly)
-    # since we save F16 / BF16 multiple times
     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
+    Like .push_to_hub(...) but auto-converts 4bit weights to float16.
-    converted to float16 with as few overhead as possible.
+    `save_method`:
-
-    Choose for `save_method` to be either:
     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
+    Like .push_to_hub(...) but auto-converts 4bit weights to float16.
-    converted to float16 with as few overhead as possible.
+    `save_method`:
-
-    Choose for `save_method` to be either:
     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 path, or hub repo ID when `push_to_hub` is True.
-      `save_directory`: local folder path or huggingface hub ID when `push_to_hub` is set to True, e.g. `my_model`
+    `torchao_config` (TorchAOBaseConfig): torchao quant config, full list:
-      `torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
+        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
     """
 
     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`
+    `save_directory`: local path, or hub repo ID when `push_to_hub` is True.
-       parameter, do NOT pass `torchao_config`. The function will convert the QAT
+    `torchao_config` (TorchAOBaseConfig): required for PTQ, must be None for QAT.
-       fake-quantized weights to real quantized weights and save directly.
+        Options: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
-
-    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
     """
     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:

unsloth/tokenizer_utils.py

@@ -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
-    Fix sentencepiece tokenizers that didn't extend the vocab with user-defined
+    by llama.cpp's convert_hf_to_gguf.py).
-    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>)
+    Also retypes special tokens (e.g. Gemma 3's <start_of_turn>/<end_of_turn>) typed NORMAL instead of
-    that exist in the sentencepiece model but are typed NORMAL instead of CONTROL.
+    CONTROL. NORMAL writes token_type=1 to GGUF, breaking llama.cpp chat inference since parse_special
-    NORMAL writes token_type=1 to GGUF, breaking llama.cpp chat inference since
+    only matches CONTROL (type=3).
-    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!
+        # /tmp on Kaggle has a ~80GB limit, so use it
-        # Let's utilize them
         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
+    # Try slow tokenizer, fall back to Fast (e.g. Deepseek has no tokenizer.model)
-    # Mainly to solve Deepseek models with no tokenizer.model file
     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
+    """Rightmost {% endfor %}/{% endif %} (any dash variant) as a dict with
-    with start/end/text/dash_left/dash_right. Tokens inside Jinja comments
+    start/end/text/dash_left/dash_right. Tokens inside Jinja comments are ignored.
-    are ignored. `endfor`/`endif` kwargs kept for back-compat, 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
+                # Out-of-bounds tokenizer and the slow version won't load either
-                # load the slow tokenizer version :(
                 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.
+    """Return a concise JSON-safe diagnostic summary of a tokenizer for logging/debugging/Studio UI.
-
+    Missing attributes fall back to ``None`` rather than raising."""
-    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 {
         "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):
-        """
+        """NEFTune forward hook for torch.nn.Embedding layers (adapted from
-        Implements the NEFTune forward pass for the model using forward hooks. Note this works only for
+        https://github.com/neelsjain/NEFTune). Set `module.neftune_noise_alpha`, then register:
-        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:
         ```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 SFT trainer with Unsloth changes."""
-    Patches the trainer with changes
-    """
     try:
         sft_trainer = eval(f"trl.trainer.sft_trainer.SFTTrainer")
     except:

unsloth/trainer.py

@@ -57,11 +57,10 @@ logger = logging.getLogger(__name__)
 
 
 class UnslothVisionDataCollator(_UnslothVisionDataCollatorBase):
-    """
+    """Drop-in zoo collator that validates local video paths per batch (deduped
-    Drop-in zoo collator that validates local video paths on every batch
+    across batches), applying formatting_func first so formatter-made paths are
-    (deduped across batches), applying formatting_func first so formatter-made
+    checked too. Raises FileNotFoundError on missing files instead of silently
-    paths are checked too. Raises FileNotFoundError on missing files instead
+    training on empty video tensors (issue #5085).
-    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
+# trl>=0.13.0 changed how several params are passed to the trainer; patch for it
-# We need to patch to make the transition smooth
 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)

unsloth/utils/attention_dispatch.py

@@ -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
+    NOTE(djsaunde): Rebuilt every forward pass (not once per layer) since it can
-        it can be invalid across passes (e.g. switching training/inference). Kept
+    go stale across passes (e.g. switching training/inference).
-        separate from AttentionContext to group params.
     """
 
     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
+    Backend priority (speed): FlashAttention (varlen for packed inputs with
-    inputs with `seq_info`, else dense), then xFormers, then SDPA as fallback.
+    `seq_info`, else dense), then xFormers, then SDPA. Varlen flash avoids
-    Varlen flash is preferred for packed batches as it avoids padding; xFormers
+    padding for packed batches; xFormers and SDPA pack via a block-diagonal mask.
-    and SDPA handle packing 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()}")

unsloth/utils/hf_hub.py

@@ -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:

unsloth/utils/packing.py

@@ -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")