[pre-commit.ci] auto fixes from pre-commit.com hooks

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

.pre-commit-config.yaml

@@ -14,4 +14,4 @@ repos:
         language: python
         types: [python]
         additional_dependencies:
-          - ruff==0.6.9
+          - ruff==0.14.7

tests/saving/language_models/test_merge_4bit_validation.py

@@ -24,9 +24,9 @@ def formatting_prompts_func(examples):
     return {"text": texts}
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 1: Loading Base Model and Initial Training")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 if torch.cuda.is_bf16_supported():
     compute_dtype = torch.bfloat16
@@ -58,9 +58,9 @@ dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
 
 print("✅ Base model loaded successfully!")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 2: First Fine-tuning")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 model = FastLanguageModel.get_peft_model(
     model,
@@ -114,9 +114,9 @@ trainer = SFTTrainer(
 trainer_stats = trainer.train()
 print("✅ First fine-tuning completed!")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 3: Save with Forced 4bit Merge")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 model.save_pretrained_merged(
     save_directory = "./test_4bit_model",
@@ -126,9 +126,9 @@ model.save_pretrained_merged(
 
 print("✅ Model saved with forced 4bit merge!")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 4: Loading 4bit Model and Second Fine-tuning")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 # Clean up first model
 del model
@@ -202,9 +202,9 @@ trainer_4bit = SFTTrainer(
 trainer_4bit.train()
 print("✅ Second fine-tuning on 4bit model completed!")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 5: Testing TypeError on Regular Merge (Should Fail)")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 try:
     model_4bit.save_pretrained_merged(
@@ -219,9 +219,9 @@ except TypeError as e:
     print("✅ Correct TypeError raised for 4bit base model regular merge attempt!")
     print(f"Error message: {str(e)}")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 6: Successful Save with Forced 4bit Method")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 try:
     model_4bit.save_pretrained_merged(
@@ -233,9 +233,9 @@ try:
 except Exception as e:
     assert False, f"Phase 6 failed unexpectedly: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 CLEANUP")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 # Cleanup
 safe_remove_directory("./outputs")

tests/saving/language_models/test_save_merged_grpo_model.py

@@ -40,7 +40,7 @@ def evaluate_merged_model(result_queue, load_in_4bit = False, load_in_8bit = Fal
         gpu_memory_utilization = 0.8,  # Reduce if out of memory
     )
 
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     if load_in_4bit:
         print("🔍 EVALUATION Merged model: 4 bits load")
         model_type = "merged_model_4bits"
@@ -50,7 +50,7 @@ def evaluate_merged_model(result_queue, load_in_4bit = False, load_in_8bit = Fal
     else:
         print("🔍 EVALUATION Merged model: 16 bits load")
         model_type = "merged_model_16bits"
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
 
     evaluate_model_aime(
         model = model,
@@ -374,9 +374,9 @@ def training_run(result_queue):
 
     def compare_model_results(all_results):
         """Generate comprehensive comparison of multiple model results"""
-        print(f"\n{'='*80}")
+        print(f"\n{'=' * 80}")
         print("COMPREHENSIVE MODEL COMPARISON")
-        print(f"{'='*80}")
+        print(f"{'=' * 80}")
 
         # Main table
         print(
@@ -395,9 +395,9 @@ def training_run(result_queue):
 
         # Improvement analysis
         if len(all_results) > 1:
-            print(f"\n{'='*50}")
+            print(f"\n{'=' * 50}")
             print("IMPROVEMENT ANALYSIS")
-            print(f"{'='*50}")
+            print(f"{'=' * 50}")
 
             base_result = all_results[0]
             for result in all_results[1:]:
@@ -504,9 +504,9 @@ def training_run(result_queue):
     from transformers import DataCollatorForSeq2Seq, TrainingArguments
     from unsloth import is_bfloat16_supported
 
-    print(f"\n{'*'*60}")
+    print(f"\n{'*' * 60}")
     print("🎯 STAGE 1: Qlora Fine-Tuning on LIMO")
-    print(f"{'*'*60}")
+    print(f"{'*' * 60}")
 
     model = FastLanguageModel.get_peft_model(
         model,
@@ -629,9 +629,9 @@ def training_run(result_queue):
                 continue
         return scores
 
-    print(f"\n{'*'*60}")
+    print(f"\n{'*' * 60}")
     print("🎯 STAGE 2: GRPO Fine-Tuning on GSM8K")
-    print(f"{'*'*60}")
+    print(f"{'*' * 60}")
 
     # Get max prompt length
     max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
@@ -691,9 +691,9 @@ def training_run(result_queue):
 
     print("✅ GRPO training completed!")
 
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     print("🔍 EVALUATION 3: Final GRPO Model")
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
 
     grpo_results = evaluate_model_aime(
         model = model,
@@ -709,9 +709,9 @@ def training_run(result_queue):
     all_results.append(grpo_results)
     print("✅ Final model evaluation complete!")
 
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     print("💾 SAVING FINAL MODEL")
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
 
     # Save as merged model
     try:
@@ -817,9 +817,9 @@ if __name__ == "__main__":
 
     # AIME-specific comparison function
 
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print("🏆 FINAL TRAINING PIPELINE RESULTS")
-    print(f"{'='*80}")
+    print(f"{'=' * 80}")
 
     # Use the AIME-specific comparison
     compare_aime_results(all_results)

tests/saving/non_peft/test_mistral_non_peft.py

@@ -11,9 +11,9 @@ sys.path.insert(0, str(REPO_ROOT))
 from tests.utils.cleanup_utils import safe_remove_directory
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 1: Loading Base Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name = "unsloth/mistral-7b-v0.3",
@@ -30,9 +30,9 @@ print("✅ Base model loaded successfully!")
 ### Attemtping save merge
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings(record = True) as w:
     warnings.simplefilter("always")
@@ -48,9 +48,9 @@ with warnings.catch_warnings(record = True) as w:
 print("✅ Correct warning detected for non-PeftModel merge attempt!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 3: Using save_pretrained (Should Succeed)")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 try:

tests/saving/non_peft/test_whisper_non_peft.py

@@ -11,9 +11,9 @@ sys.path.insert(0, str(REPO_ROOT))
 from tests.utils.cleanup_utils import safe_remove_directory
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 1: Loading Base Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 model, tokenizer = FastModel.from_pretrained(
     model_name = "unsloth/whisper-large-v3",
@@ -30,9 +30,9 @@ print("✅ Base model loaded successfully!")
 ### Attemtping save merge
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings(record = True) as w:
     warnings.simplefilter("always")
@@ -48,9 +48,9 @@ with warnings.catch_warnings(record = True) as w:
 print("✅ Correct warning detected for non-PeftModel merge attempt!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 PHASE 3: Using save_pretrained (Should Succeed)")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 try:

tests/saving/test_unsloth_save.py

@@ -138,9 +138,9 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
 
     # Check model files
     assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isfile(os.path.join(save_path, "config.json")), (
+        "config.json not found."
+    )
 
     weight_files = [
         f
@@ -151,18 +151,18 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
 
     # Check tokenizer files
     for file in tokenizer_files:
-        assert os.path.isfile(
-            os.path.join(save_path, file)
-        ), f"{file} not found in the save directory."
+        assert os.path.isfile(os.path.join(save_path, file)), (
+            f"{file} not found in the save directory."
+        )
 
     # Check config to see if it is 16bit by checking for quantization config
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" not in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" not in config, (
+        "Quantization config not found in the model config."
+    )
 
     # Store the size of the model files
     total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files)
@@ -191,9 +191,9 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
 
     # Check model files
     assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isfile(os.path.join(save_path, "config.json")), (
+        "config.json not found."
+    )
 
     weight_files = [
         f
@@ -204,9 +204,9 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
 
     # Check tokenizer files
     for file in tokenizer_files:
-        assert os.path.isfile(
-            os.path.join(save_path, file)
-        ), f"{file} not found in the save directory."
+        assert os.path.isfile(os.path.join(save_path, file)), (
+            f"{file} not found in the save directory."
+        )
 
     # Store the size of the model files
     total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files)
@@ -214,18 +214,18 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
 
     print(f"Total size of merged_4bit files: {total_size} bytes")
 
-    assert (
-        total_size < save_file_sizes["merged_16bit"][model.config._name_or_path]
-    ), "Merged 4bit files are larger than merged 16bit files."
+    assert total_size < save_file_sizes["merged_16bit"][model.config._name_or_path], (
+        "Merged 4bit files are larger than merged 16bit files."
+    )
 
     # Check config to see if it is 4bit
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" in config, (
+        "Quantization config not found in the model config."
+    )
 
     # Test loading the model from the saved path
     loaded_model, loaded_tokenizer = FastModel.from_pretrained(
@@ -269,12 +269,12 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
     torchao_save_path = save_path + "-torchao"
 
     # Check model files
-    assert os.path.isdir(
-        torchao_save_path
-    ), f"Directory {torchao_save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(torchao_save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isdir(torchao_save_path), (
+        f"Directory {torchao_save_path} does not exist."
+    )
+    assert os.path.isfile(os.path.join(torchao_save_path, "config.json")), (
+        "config.json not found."
+    )
 
     weight_files = [
         f
@@ -285,9 +285,9 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
 
     # Check tokenizer files
     for file in tokenizer_files:
-        assert os.path.isfile(
-            os.path.join(torchao_save_path, file)
-        ), f"{file} not found in the save directory."
+        assert os.path.isfile(os.path.join(torchao_save_path, file)), (
+            f"{file} not found in the save directory."
+        )
 
     # Store the size of the model files
     total_size = sum(
@@ -295,18 +295,18 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
     )
     save_file_sizes["torchao"][model.config._name_or_path] = total_size
 
-    assert (
-        total_size < save_file_sizes["merged_16bit"][model.config._name_or_path]
-    ), "torchao files are larger than merged 16bit files."
+    assert total_size < save_file_sizes["merged_16bit"][model.config._name_or_path], (
+        "torchao files are larger than merged 16bit files."
+    )
 
     # Check config to see if it is quantized with torchao
     config_path = os.path.join(torchao_save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" in config, (
+        "Quantization config not found in the model config."
+    )
 
     # Test loading the model from the saved path
     # can't set `load_in_4bit` to True because the model is torchao quantized
@@ -351,9 +351,9 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
     torchao_save_path = save_path + "-torchao"
 
     # Verify files exist
-    assert os.path.isdir(
-        torchao_save_path
-    ), f"TorchAO directory {torchao_save_path} does not exist."
+    assert os.path.isdir(torchao_save_path), (
+        f"TorchAO directory {torchao_save_path} does not exist."
+    )
 
     # Load with safe globals
     import torch.serialization

tests/saving/text_to_speech_models/test_csm.py

@@ -20,9 +20,9 @@ require_python_package("soundfile")
 
 import soundfile as sf
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 1: Loading Model and LoRA Adapters")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastModel.from_pretrained(
@@ -62,17 +62,17 @@ model = FastModel.get_peft_model(
 print("✅ Model and LoRA adapters loaded successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 2: Checking Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
 print("✅ Model is an instance of PeftModel!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 3: Checking Config Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 def find_lora_base_model(model_to_inspect):
@@ -86,15 +86,15 @@ def find_lora_base_model(model_to_inspect):
 
 config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
 
-assert (
-    config_model.__class__.__name__ == base_model_class
-), f"Expected config_model class to be {base_model_class}"
+assert config_model.__class__.__name__ == base_model_class, (
+    f"Expected config_model class to be {base_model_class}"
+)
 print("✅ config_model returns correct Base Model class:", str(base_model_class))
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 4: Saving and Merging Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings():
     warnings.simplefilter("error")  # Treat warnings as errors
@@ -104,9 +104,9 @@ with warnings.catch_warnings():
     except Exception as e:
         assert False, f"Model saving/merging failed with exception: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 5: Loading Model for Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, processor = FastModel.from_pretrained(
@@ -124,9 +124,9 @@ processor = AutoProcessor.from_pretrained("unsloth/csm-1b")
 print("✅ Model loaded for inference successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 6: Running Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 from transformers import pipeline

tests/saving/text_to_speech_models/test_lasa.py

@@ -36,9 +36,9 @@ except Exception as e:
 
 codec_model.to("cpu")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 1: Loading Model and LoRA Adapters")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 max_seq_length = 2048
 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -69,17 +69,17 @@ model = FastLanguageModel.get_peft_model(
 print("✅ Model and LoRA adapters loaded successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 2: Checking Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
 print("✅ Model is an instance of PeftModel!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 3: Checking Config Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 def find_lora_base_model(model_to_inspect):
@@ -93,15 +93,15 @@ def find_lora_base_model(model_to_inspect):
 
 config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
 
-assert (
-    config_model.__class__.__name__ == base_model_class
-), f"Expected config_model class to be {base_model_class}"
+assert config_model.__class__.__name__ == base_model_class, (
+    f"Expected config_model class to be {base_model_class}"
+)
 print("✅ config_model returns correct Base Model class:", str(base_model_class))
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 4: Saving and Merging Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings():
     warnings.simplefilter("error")  # Treat warnings as errors
@@ -111,9 +111,9 @@ with warnings.catch_warnings():
     except Exception as e:
         assert False, f"Model saving/merging failed with exception: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 5: Loading Model for Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -130,9 +130,9 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 print("✅ Model loaded for inference successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 6: Running Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 from transformers import pipeline

tests/saving/text_to_speech_models/test_orpheus.py

@@ -24,9 +24,9 @@ from snac import SNAC
 
 snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
 snac_model = snac_model.to("cuda")
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 1: Loading Model and LoRA Adapters")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -64,17 +64,17 @@ model = FastLanguageModel.get_peft_model(
 print("✅ Model and LoRA adapters loaded successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 2: Checking Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
 print("✅ Model is an instance of PeftModel!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 3: Checking Config Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 def find_lora_base_model(model_to_inspect):
@@ -88,15 +88,15 @@ def find_lora_base_model(model_to_inspect):
 
 config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
 
-assert (
-    config_model.__class__.__name__ == base_model_class
-), f"Expected config_model class to be {base_model_class}"
+assert config_model.__class__.__name__ == base_model_class, (
+    f"Expected config_model class to be {base_model_class}"
+)
 print("✅ config_model returns correct Base Model class:", str(base_model_class))
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 4: Saving and Merging Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings():
     warnings.simplefilter("error")  # Treat warnings as errors
@@ -106,9 +106,9 @@ with warnings.catch_warnings():
     except Exception as e:
         assert False, f"Model saving/merging failed with exception: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 5: Loading Model for Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -125,9 +125,9 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 print("✅ Model loaded for inference successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 6: Running Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 # @title Run Inference

tests/saving/text_to_speech_models/test_whisper.py

@@ -22,9 +22,9 @@ require_python_package("soundfile")
 
 import soundfile as sf
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 1: Loading Model and LoRA Adapters")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastModel.from_pretrained(
@@ -62,17 +62,17 @@ model = FastModel.get_peft_model(
 print("✅ Model and LoRA adapters loaded successfully!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 2: Checking Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
 print("✅ Model is an instance of PeftModel!")
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 3: Checking Config Model Class Type")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 def find_lora_base_model(model_to_inspect):
@@ -86,15 +86,15 @@ def find_lora_base_model(model_to_inspect):
 
 config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
 
-assert (
-    config_model.__class__.__name__ == base_model_class
-), f"Expected config_model class to be {base_model_class}"
+assert config_model.__class__.__name__ == base_model_class, (
+    f"Expected config_model class to be {base_model_class}"
+)
 print("✅ config_model returns correct Base Model class:", str(base_model_class))
 
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 4: Saving and Merging Model")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 with warnings.catch_warnings():
     warnings.simplefilter("error")  # Treat warnings as errors
@@ -104,9 +104,9 @@ with warnings.catch_warnings():
     except Exception as e:
         assert False, f"Model saving/merging failed with exception: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 5: Loading Model for Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 model, tokenizer = FastModel.from_pretrained(
@@ -124,9 +124,9 @@ model, tokenizer = FastModel.from_pretrained(
 
 print("✅ Model loaded for inference successfully!")
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 6: Downloading Sample Audio File")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 audio_url = "https://upload.wikimedia.org/wikipedia/commons/5/5b/Speech_12dB_s16.flac"
 audio_file = "Speech_12dB_s16.flac"
@@ -143,9 +143,9 @@ try:
 except Exception as e:
     assert False, f"Failed to download audio file: {e}"
 
-print(f"\n{'='*80}")
+print(f"\n{'=' * 80}")
 print("🔍 SECTION 7: Running Inference")
-print(f"{'='*80}")
+print(f"{'=' * 80}")
 
 
 from transformers import pipeline
@@ -185,9 +185,9 @@ all_phrases_found = all(
     phrase.lower() in transcribed_lower for phrase in expected_phrases
 )
 
-assert (
-    all_phrases_found
-), f"Expected phrases not found in transcription: {transcribed_text['text']}"
+assert all_phrases_found, (
+    f"Expected phrases not found in transcription: {transcribed_text['text']}"
+)
 print("✅ Transcription contains all expected phrases!")
 
 

tests/test_model_registry.py

@@ -71,9 +71,9 @@ def test_model_registration(model_test_param: ModelTestParam):
     registration_method()
     registered_models = MODEL_REGISTRY.keys()
     missing_models = _test_model_uploaded(registered_models)
-    assert (
-        not missing_models
-    ), f"{model_test_param.name} missing following models: {missing_models}"
+    assert not missing_models, (
+        f"{model_test_param.name} missing following models: {missing_models}"
+    )
 
 
 def test_all_model_registration():

tests/test_raw_text.py

@@ -107,30 +107,30 @@ def test_raw_text_loader():
         # Test loading with tokenized output (new efficient mode)
         tokenized_dataset = loader.load_from_file(test_file, return_tokenized = True)
         assert len(tokenized_dataset) > 0, "Should create at least one tokenized chunk"
-        assert (
-            "input_ids" in tokenized_dataset.column_names
-        ), "Dataset should have 'input_ids' column"
-        assert (
-            "attention_mask" in tokenized_dataset.column_names
-        ), "Dataset should have 'attention_mask' column"
+        assert "input_ids" in tokenized_dataset.column_names, (
+            "Dataset should have 'input_ids' column"
+        )
+        assert "attention_mask" in tokenized_dataset.column_names, (
+            "Dataset should have 'attention_mask' column"
+        )
 
         # Verify tokenized data structure
         first_sample = tokenized_dataset[0]
         assert isinstance(first_sample["input_ids"], list), "input_ids should be a list"
-        assert isinstance(
-            first_sample["attention_mask"], list
-        ), "attention_mask should be a list"
-        assert len(first_sample["input_ids"]) == len(
-            first_sample["attention_mask"]
-        ), "input_ids and attention_mask should have same length"
+        assert isinstance(first_sample["attention_mask"], list), (
+            "attention_mask should be a list"
+        )
+        assert len(first_sample["input_ids"]) == len(first_sample["attention_mask"]), (
+            "input_ids and attention_mask should have same length"
+        )
 
         # Verify labels field exists (for causal LM training)
-        assert (
-            "labels" in tokenized_dataset.column_names
-        ), "Dataset should have 'labels' column"
-        assert (
-            first_sample["labels"] == first_sample["input_ids"]
-        ), "labels should match input_ids"
+        assert "labels" in tokenized_dataset.column_names, (
+            "Dataset should have 'labels' column"
+        )
+        assert first_sample["labels"] == first_sample["input_ids"], (
+            "labels should match input_ids"
+        )
 
         # Test constructor validation
         try:

tests/utils/aime_eval.py

@@ -205,10 +205,10 @@ def evaluate_model_aime(
 ):
     """Evaluate model on combined AIME dataset with official configuration"""
 
-    print(f"\n{'='*70}")
+    print(f"\n{'=' * 70}")
     print(f"🧮 AIME EVALUATION - {model_type.upper()} MODEL")
     print(f"Combined Dataset: test2024 + test2025-I + test2025-II")
-    print(f"{'='*70}")
+    print(f"{'=' * 70}")
 
     # Load combined AIME dataset
     try:
@@ -417,9 +417,9 @@ def evaluate_model_aime(
         json.dump({"results": results, "records": records}, f, indent = 4)
 
     # Print comprehensive summary
-    print(f"\n{'='*70}")
+    print(f"\n{'=' * 70}")
     print(f"📊 AIME EVALUATION RESULTS - {model_type.upper()}")
-    print(f"{'='*70}")
+    print(f"{'=' * 70}")
 
     print(f"\n🎯 Overall Performance:")
     print(f"   Total problems:       {total_problems:>6}")
@@ -464,7 +464,7 @@ def evaluate_model_aime(
 
     print(f"\n🎖️  AIME Performance:     {tier} ({accuracy:.1f}%)")
     print(f"\n💾 Detailed results saved to: {filename}")
-    print(f"\n{'='*70}")
+    print(f"\n{'=' * 70}")
 
     return results
 
@@ -472,9 +472,9 @@ def evaluate_model_aime(
 # Comparison functions for multiple model results
 def compare_aime_results(all_results):
     """Generate comprehensive comparison for AIME evaluation results"""
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print("COMPREHENSIVE AIME MODEL COMPARISON")
-    print(f"{'='*80}")
+    print(f"{'=' * 80}")
 
     # Main comparison table
     print(
@@ -493,9 +493,9 @@ def compare_aime_results(all_results):
 
     # Performance improvement analysis
     if len(all_results) > 1:
-        print(f"\n{'='*50}")
+        print(f"\n{'=' * 50}")
         print("IMPROVEMENT ANALYSIS")
-        print(f"{'='*50}")
+        print(f"{'=' * 50}")
 
         base_result = all_results[0]  # Assume first is base model
 
@@ -509,9 +509,9 @@ def compare_aime_results(all_results):
             print(f"  Pass@K improvement:    {pass_k_improvement:+.1f}%")
 
     # Dataset breakdown
-    print(f"\n{'='*50}")
+    print(f"\n{'=' * 50}")
     print("PERFORMANCE BY DATASET")
-    print(f"{'='*50}")
+    print(f"{'=' * 50}")
 
     # Get all unique datasets from the first result
     if all_results and "source_accuracies" in all_results[0]:

tests/utils/cleanup_utils.py

@@ -57,7 +57,7 @@ def clear_memory(variables_to_clear = None, verbose = False, clear_all_caches =
     for i in range(3):
         collected = gc.collect()
         if verbose and collected > 0:
-            print(f"GC pass {i+1}: collected {collected} objects")
+            print(f"GC pass {i + 1}: collected {collected} objects")
 
     # 4. CUDA cleanup
     if torch.cuda.is_available():

unsloth/kernels/moe/grouped_gemm/interface.py

@@ -222,22 +222,22 @@ def grouped_gemm_forward(
     W = W.view(-1, W.shape[-1])
 
     if permute_x or permute_y:
-        assert (
-            gather_indices is not None
-        ), "gather_indices must be provided when permute_x or permute_y is True"
+        assert gather_indices is not None, (
+            "gather_indices must be provided when permute_x or permute_y is True"
+        )
         assert gather_indices.is_contiguous()
         assert gather_indices.device.type == "cuda"
         assert gather_indices.ndim == 1
         total_tokens = gather_indices.shape[0]
         num_tokens = total_tokens // topk
         if permute_x:
-            assert (
-                X.shape[0] == num_tokens
-            ), f"X.shape[0] ({X.shape[0]}) must match num_tokens ({num_tokens})"
+            assert X.shape[0] == num_tokens, (
+                f"X.shape[0] ({X.shape[0]}) must match num_tokens ({num_tokens})"
+            )
         else:
-            assert (
-                X.shape[0] == total_tokens
-            ), f"X.shape[0] ({X.shape[0]}) must match total_tokens ({total_tokens})"
+            assert X.shape[0] == total_tokens, (
+                f"X.shape[0] ({X.shape[0]}) must match total_tokens ({total_tokens})"
+            )
     else:
         total_tokens = X.shape[0]
         num_tokens = total_tokens // topk
@@ -383,12 +383,12 @@ def grouped_gemm_dX(
     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"
+    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"
+    )
     assert dY.is_contiguous()
     assert W.is_contiguous()
     assert m_sizes.is_contiguous()
@@ -433,15 +433,15 @@ def grouped_gemm_dX(
     # N = N_total // num_experts
     assert N_grad == N, f"Grad_output N ({N_grad}) must match weight N ({N})"
 
-    assert (
-        M_total % topk == 0
-    ), f"M_total ({M_total}) must be divisible by topk ({topk})"
+    assert M_total % topk == 0, (
+        f"M_total ({M_total}) must be divisible by topk ({topk})"
+    )
     num_tokens = M_total // topk
 
     total_tokens = gather_indices.shape[0]
-    assert (
-        total_tokens == M_total
-    ), f"Total tokens ({total_tokens}) must match M_total ({M_total})"
+    assert total_tokens == M_total, (
+        f"Total tokens ({total_tokens}) must match M_total ({M_total})"
+    )
 
     # Note that the output shape is [NUM_TOKENS * TOPK, K] even when `permute_x` is True since we need to accumulate gradients across all experts chosen by the token.
     # This will be done in a post-processing step reduction step.
@@ -766,17 +766,17 @@ class GroupedGemm(torch.autograd.Function):
 
         if not autotune:
             if not dW_only:
-                assert (
-                    kernel_config_bwd_dX is not None
-                ), "kernel_config_bwd_dX must be provided if autotune is False"
+                assert kernel_config_bwd_dX is not None, (
+                    "kernel_config_bwd_dX must be provided if autotune is False"
+                )
             if not dX_only:
-                assert (
-                    kernel_config_bwd_dW is not None
-                ), "kernel_config_bwd_dW must be provided if autotune is False"
+                assert kernel_config_bwd_dW is not None, (
+                    "kernel_config_bwd_dW must be provided if autotune is False"
+                )
 
-        assert (
-            not fuse_mul_post
-        ), "fused_mul should only be used for inference, not for training"
+        assert not fuse_mul_post, (
+            "fused_mul should only be used for inference, not for training"
+        )
 
         if not dX_only:
             bwd_dW_config = {}
@@ -872,21 +872,21 @@ def check_valid_config_fwd(
     is_second_gemm = not is_first_gemm
 
     assert not (permute_x and permute_y), "Cannot permute both X and Y"
-    assert not (
-        is_second_gemm and permute_x
-    ), "Cannot permute X for the second grouped GEMM"
-    assert not (
-        is_first_gemm and permute_y
-    ), "Cannot permute Y for the first grouped GEMM"
-    assert not (
-        fuse_mul_post and is_first_gemm
-    ), "Cannot fuse mul for the first grouped GEMM"
-    assert not (
-        use_tma_load_x and permute_x
-    ), "Cannot use TMA load and permute X unless on sm100+ (Blackwell+)"
-    assert not (
-        use_tma_store and permute_y and is_second_gemm
-    ), "Cannot use TMA store and permute Y for the second grouped GEMM unless on sm100+ (Blackwell+)"
+    assert not (is_second_gemm and permute_x), (
+        "Cannot permute X for the second grouped GEMM"
+    )
+    assert not (is_first_gemm and permute_y), (
+        "Cannot permute Y for the first grouped GEMM"
+    )
+    assert not (fuse_mul_post and is_first_gemm), (
+        "Cannot fuse mul for the first grouped GEMM"
+    )
+    assert not (use_tma_load_x and permute_x), (
+        "Cannot use TMA load and permute X unless on sm100+ (Blackwell+)"
+    )
+    assert not (use_tma_store and permute_y and is_second_gemm), (
+        "Cannot use TMA store and permute Y for the second grouped GEMM unless on sm100+ (Blackwell+)"
+    )
 
 
 def check_valid_config_bwd_dW(
@@ -974,9 +974,9 @@ def grouped_gemm(
 
     """
     if not autotune:
-        assert (
-            kernel_config_fwd is not None
-        ), "kernel_config_fwd must be provided if autotune is False"
+        assert kernel_config_fwd is not None, (
+            "kernel_config_fwd must be provided if autotune is False"
+        )
 
         check_valid_config_fwd(
             permute_x,
@@ -1009,14 +1009,14 @@ def grouped_gemm(
             )
 
     if permute_x or permute_y:
-        assert (
-            gather_indices is not None
-        ), "gather_indices is required when either permute_x or permute_y is True"
+        assert gather_indices is not None, (
+            "gather_indices is required when either permute_x or permute_y is True"
+        )
 
     if fuse_mul_post:
-        assert (
-            topk_weights is not None
-        ), "topk_weights is required when fuse_mul_post is True"
+        assert topk_weights is not None, (
+            "topk_weights is required when fuse_mul_post is True"
+        )
 
     X = X.view(-1, X.shape[-1])
     m_sizes = m_sizes.view(-1)

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

@@ -62,9 +62,9 @@ class Llama4GroupedGemmTextMoe(Llama4TextMoe):
 
         # Permute in-place expert weights
         E, K, N = self.num_experts, self.hidden_dim, self.experts.expert_dim
-        assert self.experts.gate_up_proj.shape == torch.Size(
-            [E, K, 2 * N]
-        ), f"{self.experts.gate_up_proj.shape} != {[E, K, 2 * N]}"
+        assert self.experts.gate_up_proj.shape == torch.Size([E, K, 2 * N]), (
+            f"{self.experts.gate_up_proj.shape} != {[E, K, 2 * N]}"
+        )
         permuted_shape = [E, 2 * N, K]
         permuted_stride = [2 * N * K, K, 1]
         if verbose:
@@ -79,9 +79,9 @@ class Llama4GroupedGemmTextMoe(Llama4TextMoe):
                 f"{self.experts.gate_up_proj.shape}:{self.experts.gate_up_proj.stride()}"
             )
 
-        assert self.experts.down_proj.shape == torch.Size(
-            [E, N, K]
-        ), f"{self.experts.down_proj.shape} != {[E, N, K]}"
+        assert self.experts.down_proj.shape == torch.Size([E, N, K]), (
+            f"{self.experts.down_proj.shape} != {[E, N, K]}"
+        )
         permuted_shape = [E, K, N]
         permuted_stride = [K * N, N, 1]
         if verbose:
@@ -110,9 +110,9 @@ class Llama4GroupedGemmTextMoe(Llama4TextMoe):
             if any(n in name for n in self.EXPERT_WEIGHT_NAMES):
                 param_to_copy = param_to_copy.permute(0, 2, 1)
 
-            assert (
-                param.shape == param_to_copy.shape
-            ), f"{param.shape} != {param_to_copy.shape}"
+            assert param.shape == param_to_copy.shape, (
+                f"{param.shape} != {param_to_copy.shape}"
+            )
             param.copy_(param_to_copy)
 
         return self
@@ -269,7 +269,9 @@ class Llama4TritonTextMoe(Llama4GroupedGemmTextMoe):
         verbose = False,
     ):
         super().__init__(config, overlap_router_shared = overlap_router_shared)
-        assert not permute_x, "Llama4 triton grouped gemm does not support permute x due to pre-multiplication of router weights"
+        assert not permute_x, (
+            "Llama4 triton grouped gemm does not support permute x due to pre-multiplication of router weights"
+        )
         self.permute_x = permute_x
         self.permute_y = permute_y
         self.autotune = autotune
@@ -295,9 +297,9 @@ class Llama4TritonTextMoe(Llama4GroupedGemmTextMoe):
             if any(n in name for n in self.EXPERT_WEIGHT_NAMES):
                 param_to_copy = param_to_copy.permute(0, 2, 1)
 
-            assert (
-                param.shape == param_to_copy.shape
-            ), f"{param.shape} != {param_to_copy.shape}"
+            assert param.shape == param_to_copy.shape, (
+                f"{param.shape} != {param_to_copy.shape}"
+            )
             param.copy_(param_to_copy)
 
         return self

unsloth/kernels/moe/tests/common.py

@@ -124,16 +124,16 @@ def assert_close(ref, tri, maxtol = None, rmstol = None, description = "--", ver
     # 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 = }"
+    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"
+    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)
 

unsloth/kernels/moe/tests/moe_utils.py

@@ -142,12 +142,12 @@ def check_gate_up_proj_grad(
         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}"
-        assert (
-            ref_up_proj_grad.shape == test_up_proj_grad.shape
-        ), f"{ref_up_proj_grad.shape} != {test_up_proj_grad.shape}"
+        assert ref_gate_proj_grad.shape == test_gate_proj_grad.shape, (
+            f"{ref_gate_proj_grad.shape} != {test_gate_proj_grad.shape}"
+        )
+        assert 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()
@@ -199,9 +199,9 @@ def check_tensor_allclose(
     diff = (X_ref - X_test).abs().max()
     if verbose:
         print(f"{name} diff: {diff.detach().cpu().item():.6f}")
-    assert torch.allclose(
-        X_ref, X_test, atol = atol, rtol = rtol
-    ), f"{name} diff: {diff.detach().cpu().item():.6f}"
+    assert torch.allclose(X_ref, X_test, atol = atol, rtol = rtol), (
+        f"{name} diff: {diff.detach().cpu().item():.6f}"
+    )
 
 
 def check_expert_grads(
@@ -217,26 +217,26 @@ def check_expert_grads(
     for field in fields_to_check:
         ref_grads = getattr(ref_result, field)
         test_grads = getattr(test_result, field)
-        assert (
-            ref_grads.shape == test_grads.shape
-        ), f"{field}: {ref_grads.shape} != {test_grads.shape}"
+        assert 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]
             diff = (ref_grad - test_grad).abs().max()
-            assert torch.allclose(
-                ref_grad, test_grad, atol = atol, rtol = rtol
-            ), f"{field}[{i}] diff: {diff.detach().cpu().item():.6f}"
+            assert torch.allclose(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}")
-        assert torch.allclose(
-            ref_grads, test_grads, atol = atol, rtol = rtol
-        ), f"{field} diff: {diff.detach().cpu().item():.6f}"
+        assert torch.allclose(ref_grads, test_grads, atol = atol, rtol = rtol), (
+            f"{field} diff: {diff.detach().cpu().item():.6f}"
+        )
 
 
 def check_grads(
@@ -268,9 +268,9 @@ def check_fwd(
     diff = (ref_output - test_output).abs().max()
     if verbose:
         print(f"output diff: {diff.detach().cpu().item():.6f}")
-    assert torch.allclose(
-        ref_output, test_output, atol = atol, rtol = rtol
-    ), f"output diff: {diff.detach().cpu().item():.6f}"
+    assert torch.allclose(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
@@ -304,9 +304,9 @@ def check_grouped_gemm_results(
         if verbose:
             print(f"{field.name} diff: {diff.detach().cpu().item():.6f}")
 
-        assert torch.allclose(
-            ref_value, test_value, atol = atol, rtol = rtol
-        ), f"{field.name} diff: {diff.detach().cpu().item():.6f}"
+        assert torch.allclose(ref_value, test_value, atol = atol, rtol = rtol), (
+            f"{field.name} diff: {diff.detach().cpu().item():.6f}"
+        )
 
 
 def run_forward(model: nn.Module, X: torch.Tensor, is_grouped_gemm: bool = False):

unsloth/kernels/moe/tests/test_grouped_gemm.py

@@ -179,9 +179,9 @@ def _test_grouped_gemm_forward(
 
     Xref = Xperm
 
-    assert (
-        Xperm.shape == (total_tokens, K) if use_W1 else (total_tokens, N)
-    ), f"Xperm.shape: {Xperm.shape}, total_tokens: {total_tokens}, K: {K}"
+    assert Xperm.shape == (total_tokens, K) if use_W1 else (total_tokens, N), (
+        f"Xperm.shape: {Xperm.shape}, total_tokens: {total_tokens}, K: {K}"
+    )
 
     ref_output = torch_grouped_gemm(X = Xref, W = W, m_sizes = expert_token_counts)
 
@@ -267,9 +267,9 @@ def _test_grouped_gemm_forward(
             test_output = unpermute(test_output, gather_indices)
         ref_output = ref_output * topk_weights[:, None]
 
-    assert torch.allclose(
-        ref_output, test_output, atol = atol, rtol = rtol
-    ), f"Grouped gemm forward failed: {(ref_output - test_output).abs().max().item():.6f}"
+    assert torch.allclose(ref_output, test_output, atol = atol, rtol = rtol), (
+        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.
@@ -540,9 +540,9 @@ def _test_grouped_gemm_backward_dX(
 
     output_shape = (total_tokens, 2 * N) if use_W1 else (total_tokens, K)
     ref_output = torch_grouped_gemm(X = Xperm, W = W, m_sizes = expert_token_counts)
-    assert (
-        ref_output.shape == output_shape
-    ), f"ref_output.shape: {ref_output.shape}, output_shape: {output_shape}"
+    assert ref_output.shape == output_shape, (
+        f"ref_output.shape: {ref_output.shape}, output_shape: {output_shape}"
+    )
 
     if permute_y:
         ref_output = unpermute(ref_output, gather_indices)
@@ -620,12 +620,12 @@ def _test_grouped_gemm_backward_dX(
             is_first_gemm = use_W1,
             dX_only = True,
         )
-        assert (
-            test_output.shape == ref_output.shape
-        ), f"test_output.shape: {test_output.shape}, ref_output.shape: {ref_output.shape}"
-        assert torch.allclose(
-            test_output, ref_output, atol = atol, rtol = rtol
-        ), f"Grouped gemm backward_dX forward outputs mismatch: {(test_output - ref_output).abs().max().item():.6f}"
+        assert test_output.shape == ref_output.shape, (
+            f"test_output.shape: {test_output.shape}, ref_output.shape: {ref_output.shape}"
+        )
+        assert torch.allclose(test_output, ref_output, atol = atol, rtol = rtol), (
+            f"Grouped gemm backward_dX forward outputs mismatch: {(test_output - ref_output).abs().max().item():.6f}"
+        )
         test_output.backward(grad_output)
         assert X_.grad is not None
 
@@ -636,19 +636,19 @@ def _test_grouped_gemm_backward_dX(
         if permute_x and use_W1:
             X_grad_unperm = unpermute(Xperm.grad, gather_indices)
             manual_grad_check = X_grad_unperm.view(num_tokens, topk, K).sum(dim = 1)
-            assert (
-                manual_grad_check.shape == X_.grad.shape
-            ), f"manual_grad_check.shape: {manual_grad_check.shape}, X_.grad.shape: {X_.grad.shape}"
-            assert torch.allclose(
-                manual_grad_check, X_.grad, atol = atol, rtol = rtol
-            ), f"Grouped gemm backward_dX forward outputs mismatch: {(manual_grad_check - X_.grad).abs().max().item():.6f}"
+            assert manual_grad_check.shape == X_.grad.shape, (
+                f"manual_grad_check.shape: {manual_grad_check.shape}, X_.grad.shape: {X_.grad.shape}"
+            )
+            assert torch.allclose(manual_grad_check, X_.grad, atol = atol, rtol = rtol), (
+                f"Grouped gemm backward_dX forward outputs mismatch: {(manual_grad_check - X_.grad).abs().max().item():.6f}"
+            )
             manual_diff = (X_.grad - manual_grad_check).abs().max().item()
             autograd_diff = (X_.grad - X.grad).abs().max().item()
             print(f"manual_diff: {manual_diff:.6f}, autograd_diff: {autograd_diff:.6f}")
         else:
-            assert torch.allclose(
-                X_.grad, ref_grad, atol = atol, rtol = rtol
-            ), f"Grouped gemm backward_dX forward outputs mismatch: {(X_.grad - ref_grad).abs().max().item():.6f}"
+            assert torch.allclose(X_.grad, ref_grad, atol = atol, rtol = rtol), (
+                f"Grouped gemm backward_dX forward outputs mismatch: {(X_.grad - ref_grad).abs().max().item():.6f}"
+            )
         return
     else:
         dX_test = grouped_gemm_dX(
@@ -677,14 +677,14 @@ def _test_grouped_gemm_backward_dX(
     if permute_x and use_W1:
         ref_grad = unpermute(ref_grad, gather_indices)
 
-    assert (
-        ref_grad.shape == dX_test.shape
-    ), f"Grouped gemm manual backward_dX outputs mismatch: ref_grad: {ref_grad.shape}, dX_test: {dX_test.shape}"
+    assert ref_grad.shape == dX_test.shape, (
+        f"Grouped gemm manual backward_dX outputs mismatch: ref_grad: {ref_grad.shape}, dX_test: {dX_test.shape}"
+    )
     diff = (ref_grad - dX_test).abs().max().item()
 
-    assert torch.allclose(
-        ref_grad, dX_test, atol = atol, rtol = rtol
-    ), f"Grouped gemm manual backward_dX outputs mismatch: {diff:.6f}"
+    assert torch.allclose(ref_grad, dX_test, atol = atol, rtol = rtol), (
+        f"Grouped gemm manual backward_dX outputs mismatch: {diff:.6f}"
+    )
 
     if permute_x and use_W1:
         # Show that reduction results in diffs
@@ -1020,12 +1020,12 @@ def _test_grouped_gemm_backward_dW(
             is_first_gemm = use_W1,
             dW_only = True,
         )
-        assert (
-            test_output.shape == ref_output.shape
-        ), f"Grouped gemm autograd backward_dW outputs mismatch: {test_output.shape} != {ref_output.shape}"
-        assert torch.allclose(
-            test_output, ref_output, atol = atol, rtol = rtol
-        ), f"Grouped gemm autograd backward_dW forward outputs mismatch: {test_output.shape} != {ref_output.shape}"
+        assert test_output.shape == ref_output.shape, (
+            f"Grouped gemm autograd backward_dW outputs mismatch: {test_output.shape} != {ref_output.shape}"
+        )
+        assert torch.allclose(test_output, ref_output, atol = atol, rtol = rtol), (
+            f"Grouped gemm autograd backward_dW forward outputs mismatch: {test_output.shape} != {ref_output.shape}"
+        )
         test_output.backward(grad_output)
         assert W_test.grad is not None
         dW_test = W_test.grad
@@ -1050,9 +1050,9 @@ def _test_grouped_gemm_backward_dW(
             autotune = autotune,
             debug = debug,
         )
-    assert (
-        W.grad.shape == dW_test.shape
-    ), f"Grouped gemm manual backward_dW outputs mismatch: W.grad: {W.grad.shape}, dW_test: {dW_test.shape}"
+    assert W.grad.shape == dW_test.shape, (
+        f"Grouped gemm manual backward_dW outputs mismatch: W.grad: {W.grad.shape}, dW_test: {dW_test.shape}"
+    )
 
     if debug:
         with torch.no_grad():
@@ -1067,14 +1067,14 @@ def _test_grouped_gemm_backward_dW(
                 print(f"Expert {i} diff: {expert_diff:.6f}")
 
             diff = (W.grad - dW_test).abs().max().item()
-            assert (
-                False
-            ), f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"
+            assert False, (
+                f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"
+            )
     else:
         diff = (W.grad - dW_test).abs().max().item()
-        assert torch.allclose(
-            W.grad, dW_test, atol = atol, rtol = rtol
-        ), f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"
+        assert torch.allclose(W.grad, dW_test, atol = atol, rtol = rtol), (
+            f"Grouped gemm manual backward_dW outputs mismatch: {diff:.6f}"
+        )
 
 
 @pytest.mark.parametrize(

unsloth/models/_utils.py

@@ -760,7 +760,7 @@ model_architectures = [
 for model_name in model_architectures:
     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"  # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
     try:
         exec(f"from {config_filepath} import {config_filename}", globals())
     except:
@@ -2349,9 +2349,9 @@ def _untie_input_output_embeddings(model: torch.nn.Module) -> None:
         raise AttributeError("Couldn't locate output projection (lm_head).")
 
     # (Optional) sanity: shapes should match [vocab, hidden]
-    assert (
-        out_proj.weight.shape == in_emb.weight.shape
-    ), f"Shape mismatch: out_proj {out_proj.weight.shape} vs in_emb {in_emb.weight.shape}"
+    assert out_proj.weight.shape == in_emb.weight.shape, (
+        f"Shape mismatch: out_proj {out_proj.weight.shape} vs in_emb {in_emb.weight.shape}"
+    )
 
     # 3) Only clone if they are actually tied (shared storage)
     if out_proj.weight.data_ptr() == in_emb.weight.data_ptr():
@@ -2366,9 +2366,9 @@ def _untie_input_output_embeddings(model: torch.nn.Module) -> None:
     model.tie_weights = _no_tie.__get__(model, model.__class__)
 
     # 5) Verify no shared storage
-    assert (
-        out_proj.weight.data_ptr() != in_emb.weight.data_ptr()
-    ), "Embeddings still tied!"
+    assert out_proj.weight.data_ptr() != in_emb.weight.data_ptr(), (
+        "Embeddings still tied!"
+    )
 
 
 def _filter_fn_to_fqns(

unsloth/models/granite.py

@@ -283,9 +283,9 @@ def GraniteAttention_fast_forward_inference(
     use_sliding_window = False,
     position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
 ):
-    assert (
-        position_embeddings is not None
-    ), f"Granite model requires position embeddings to be specified"
+    assert position_embeddings is not None, (
+        f"Granite model requires position embeddings to be specified"
+    )
 
     Xn = hidden_states
     bsz, _, hd = hidden_states.size()

unsloth/models/rl.py

@@ -686,8 +686,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
             else:
                 continue
             call_args.append(f"{k} = {k}")
-        arguments = f"\n{' '*8}" + f",\n{' '*8}".join(arguments)
-        call_args = f"\n{' '*12}" + f",\n{' '*12}".join(call_args)
+        arguments = f"\n{' ' * 8}" + f",\n{' ' * 8}".join(arguments)
+        call_args = f"\n{' ' * 12}" + f",\n{' ' * 12}".join(call_args)
         processed.append(
             (
                 arguments,
@@ -701,7 +701,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
 
     # Add tokenizer if not seen
     if "tokenizer" not in parameters and "processing_class" in parameters:
-        arguments += f",\n{' '*8}tokenizer = None"
+        arguments += f",\n{' ' * 8}tokenizer = None"
         call_args = call_args.replace(
             "processing_class = processing_class",
             "processing_class = tokenizer if tokenizer is not None else processing_class",
@@ -1705,8 +1705,8 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
                 sampling_params = re.sub(r"[\,][\s]{0,}\,", ",", sampling_params)
 
                 new_vllm_part = (
-                    f"\n{' '*8}if {args}.use_vllm:\n{sampling_params}"
-                    f"\n{' '*8}else:\n"
+                    f"\n{' ' * 8}if {args}.use_vllm:\n{sampling_params}"
+                    f"\n{' ' * 8}else:\n"
                 )
 
         if trl_version >= Version("0.18.0"):

unsloth/save.py

@@ -601,8 +601,8 @@ def unsloth_save_model(
     max_ram = int(max(0, max_ram) * maximum_memory_usage)
     print(
         f"Unsloth: Will use up to "
-        f"{round(max_ram/1024/1024/1024, 2)} out of "
-        f"{round(psutil.virtual_memory().total/1024/1024/1024, 2)} RAM for saving."
+        f"{round(max_ram / 1024 / 1024 / 1024, 2)} out of "
+        f"{round(psutil.virtual_memory().total / 1024 / 1024 / 1024, 2)} RAM for saving."
     )
 
     # Move temporary_location to /tmp in Kaggle
@@ -992,7 +992,9 @@ def install_llama_cpp_old(version = -10):
         import time
 
         for i in range(30):
-            print(f"**[WARNING]** Deleting llama.cpp directory... {30-i} seconds left.")
+            print(
+                f"**[WARNING]** Deleting llama.cpp directory... {30 - i} seconds left."
+            )
             time.sleep(1)
         import shutil
 
@@ -1009,13 +1011,13 @@ def install_llama_cpp_old(version = -10):
     # Try using MAKE
     commands = [
         "make clean -C llama.cpp",
-        f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
+        f"make all -j{(psutil.cpu_count() or 1) * 2} -C llama.cpp",
     ]
     if try_execute(commands) == "CMAKE":
         # Instead use CMAKE
         commands = [
             f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
-            f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
+            f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1) * 2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
             "cp llama.cpp/build/bin/llama-* llama.cpp",
             "rm -rf llama.cpp/build",
         ]
@@ -1056,13 +1058,13 @@ def install_llama_cpp_blocking(use_cuda = False):
         # https://github.com/ggerganov/llama.cpp/issues/7062
         # Weirdly GPU conversion for GGUF breaks??
         # f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
-        f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
+        f"make all -j{(psutil.cpu_count() or 1) * 2} -C llama.cpp",
     ]
     if try_execute(commands) == "CMAKE":
         # Instead use CMAKE
         commands = [
             f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
-            f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
+            f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1) * 2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
             "cp llama.cpp/build/bin/llama-* llama.cpp",
             "rm -rf llama.cpp/build",
         ]
@@ -1977,7 +1979,7 @@ def unsloth_save_pretrained_gguf(
 
     # Step 4: Save/merge model to 16-bit format
     print(
-        f'Unsloth: Merging model weights to {"mxfp4" if is_gpt_oss else "16-bit"} format...'
+        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)
@@ -2917,9 +2919,9 @@ def _unsloth_save_torchao_with_given_config(
     if push_to_hub:
         assert token is not None, "Unsloth: Please specify a token for uploading!"
 
-    assert (
-        torchao_config is not None
-    ), "Unsloth: Please specify a torchao_config for post-training quantization!"
+    assert torchao_config is not None, (
+        "Unsloth: Please specify a torchao_config for post-training quantization!"
+    )
 
     # first merge the lora weights
     arguments = dict(locals())

unsloth/trainer.py

@@ -383,7 +383,7 @@ def _patch_sft_trainer_auto_packing(trl_module):
                 reason = "vision-language model"
             elif is_unsupported_model:
                 reason = f"unsupported model type(s): {', '.join(model_types)}"
-            message = "Unsloth: Sample packing skipped " f"({reason} detected)."
+            message = f"Unsloth: Sample packing skipped ({reason} detected)."
             print(message)
 
         packing_active = False