diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py index 343e737710..da5bc022a4 100644 --- a/tests/saving/language_models/test_merge_4bit_validation.py +++ b/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") diff --git a/tests/saving/language_models/test_save_merged_grpo_model.py b/tests/saving/language_models/test_save_merged_grpo_model.py index 67b649305a..1f0b377a36 100644 --- a/tests/saving/language_models/test_save_merged_grpo_model.py +++ b/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) diff --git a/tests/saving/non_peft/test_mistral_non_peft.py b/tests/saving/non_peft/test_mistral_non_peft.py index e03813367d..755e19f306 100644 --- a/tests/saving/non_peft/test_mistral_non_peft.py +++ b/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: diff --git a/tests/saving/non_peft/test_whisper_non_peft.py b/tests/saving/non_peft/test_whisper_non_peft.py index 303d596c85..c72bd38e93 100644 --- a/tests/saving/non_peft/test_whisper_non_peft.py +++ b/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: diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py index 35fdad6ba0..2e760cd415 100644 --- a/tests/saving/test_unsloth_save.py +++ b/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 diff --git a/tests/saving/text_to_speech_models/test_csm.py b/tests/saving/text_to_speech_models/test_csm.py index c1a892a8d3..2230d2728a 100644 --- a/tests/saving/text_to_speech_models/test_csm.py +++ b/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 diff --git a/tests/saving/text_to_speech_models/test_lasa.py b/tests/saving/text_to_speech_models/test_lasa.py index 804ff512f9..9c373e74cf 100644 --- a/tests/saving/text_to_speech_models/test_lasa.py +++ b/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 diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py index bd8bf14979..aa542af76e 100644 --- a/tests/saving/text_to_speech_models/test_orpheus.py +++ b/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 diff --git a/tests/saving/text_to_speech_models/test_whisper.py b/tests/saving/text_to_speech_models/test_whisper.py index 55f6d98ca0..676b174ef2 100644 --- a/tests/saving/text_to_speech_models/test_whisper.py +++ b/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!") diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index fb9a734c6f..70b225ade3 100644 --- a/tests/test_model_registry.py +++ b/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(): diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index 9f2e8cda4e..38487a7766 100644 --- a/tests/test_raw_text.py +++ b/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: diff --git a/tests/utils/aime_eval.py b/tests/utils/aime_eval.py index 131da3e50b..d70dc37056 100644 --- a/tests/utils/aime_eval.py +++ b/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]: diff --git a/tests/utils/cleanup_utils.py b/tests/utils/cleanup_utils.py index af38528196..4995bc4d79 100644 --- a/tests/utils/cleanup_utils.py +++ b/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(): diff --git a/unsloth/kernels/moe/grouped_gemm/interface.py b/unsloth/kernels/moe/grouped_gemm/interface.py index 5588458973..4b1295c5a8 100644 --- a/unsloth/kernels/moe/grouped_gemm/interface.py +++ b/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) diff --git a/unsloth/kernels/moe/grouped_gemm/reference/layers/llama4_moe.py b/unsloth/kernels/moe/grouped_gemm/reference/layers/llama4_moe.py index 6bb0bfb0c3..f7883f815e 100644 --- a/unsloth/kernels/moe/grouped_gemm/reference/layers/llama4_moe.py +++ b/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 diff --git a/unsloth/kernels/moe/tests/common.py b/unsloth/kernels/moe/tests/common.py index bfe6f2094b..9d4eb6808f 100644 --- a/unsloth/kernels/moe/tests/common.py +++ b/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) diff --git a/unsloth/kernels/moe/tests/moe_utils.py b/unsloth/kernels/moe/tests/moe_utils.py index 26ac9fdb5e..ab1cddc4ba 100644 --- a/unsloth/kernels/moe/tests/moe_utils.py +++ b/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): diff --git a/unsloth/kernels/moe/tests/test_grouped_gemm.py b/unsloth/kernels/moe/tests/test_grouped_gemm.py index bd98b6a276..ea00cdd9d2 100644 --- a/unsloth/kernels/moe/tests/test_grouped_gemm.py +++ b/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( diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 19b5fe0574..748de31e22 100644 --- a/unsloth/models/_utils.py +++ b/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( diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 79ac41c43f..13d2f8daee 100644 --- a/unsloth/models/granite.py +++ b/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() diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 30546a048d..50ea804e60 100755 --- a/unsloth/models/rl.py +++ b/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"): diff --git a/unsloth/save.py b/unsloth/save.py index 6e38d1e952..1b9141dc4e 100644 --- a/unsloth/save.py +++ b/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()) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 65abe6801f..3309730c07 100644 --- a/unsloth/trainer.py +++ b/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