Uninitialized handler

unsloth/models/_utils.py

@@ -66,6 +66,7 @@ __all__ = [
     "unsloth_compile_transformers",
     "patch_fast_lora",
     "validate_loftq_config",
+    "RaiseUninitialized",
 ]
 
 import torch
@@ -205,6 +206,28 @@ try:
 except:
     pass
 
+# Errors out on
+# Some weights of Gemma3nForConditionalGeneration were not initialized from the model checkpoint
+from transformers.modeling_utils import logger as transformers_logger
+class _RaiseUninitialized(logging.Handler):
+    def __init__(self):
+        super().__init__()
+    def emit(self, record):
+        if "some weights of" in str(record).lower():
+            raise Exception(
+                f"Unsloth: Critical error since some weights are not initialized.\n"\
+                f"Please try updating Unsloth, transformers and timm via:\n"\
+                f"`pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo transformers timm`\n"\
+                f"".str(record))
+pass
+class RaiseUninitialized:
+    def __init__(self):
+        self.error_handler = _RaiseUninitialized()
+        transformers_logger.addHandler(self.error_handler)
+    def remove(self):
+        transformers_logger.removeHandler(self.error_handler)
+pass
+
 # Patch get_model_param_count to record correct 4bit / 8bit
 from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled
 

unsloth/models/llama.py

@@ -1967,6 +1967,7 @@ class FastLlamaModel:
         # Cannot be None, since HF now checks for the config
         if load_in_4bit: kwargs["quantization_config"] = bnb_config
 
+        raise_handler = RaiseUninitialized()
         if num_labels is not None:
             model = AutoModelForSequenceClassification.from_pretrained(
                 model_name,
@@ -2030,6 +2031,7 @@ class FastLlamaModel:
             model.fast_generate = model.vllm_engine.generate
             model.fast_generate_batches = functools.partial(generate_batches, model.vllm_engine)
         pass
+        raise_handler.remove()
         # Return old flag
         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
 

unsloth/models/rl.py

@@ -458,26 +458,33 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
 
     # Edit GA / bsz and weight_decay
     replacements = {
-        "output_dir"                  : None,
-        "logging_nan_inf_filter"      : False,
-        "per_device_train_batch_size" : 4,
-        "gradient_accumulation_steps" : 2,
-        "weight_decay"                : 0.01,
-        "warmup_ratio"                : 0.1,
-        "seed"                        : 3407,
-        "optim"                       : "adamw_8bit",
-        "learning_rate"               : 5e-05,
-        "per_device_eval_batch_size"  : 4,
-        "eval_accumulation_steps"     : 2,
-        "torch_empty_cache_steps"     : 250,
-        "logging_steps"               : 1,
-        "max_seq_length"              : None,
-        "num_generations"             : 8,
-        "top_k"                       : None,
-        "vllm_mode"                   : "colocate",
-        "generation_kwargs"           : {},
-        "bf16"                        : False,
-        "fp16"                        : False,
+        "output_dir"                    : None,
+        "logging_nan_inf_filter"        : False,
+        "per_device_train_batch_size"   : 4,
+        "gradient_accumulation_steps"   : 2,
+        "weight_decay"                  : 0.01,
+        "warmup_ratio"                  : 0.1,
+        "seed"                          : 3407,
+        "optim"                         : "adamw_8bit",
+        "learning_rate"                 : 5e-05,
+        "per_device_eval_batch_size"    : 4,
+        "eval_accumulation_steps"       : 2,
+        "torch_empty_cache_steps"       : 250,
+        "logging_steps"                 : 1,
+        "max_seq_length"                : None,
+        "num_generations"               : 8,
+        "top_k"                         : None,
+        "vllm_mode"                     : "colocate",
+        "generation_kwargs"             : {},
+        "bf16"                          : False,
+        "fp16"                          : False,
+        "include_tokens_per_second"     : False,
+        "include_num_input_tokens_seen" : False,
+        "auto_find_batch_size"          : True, # Auto /2 batch size
+        "dataloader_persistent_workers" : True, # Keeps dataloader in RAM
+        "dataloader_prefetch_factor"    : 2,
+        "dataloader_pin_memory"         : True,
+        "dataloader_num_workers"        : 0, # Default is 0 means 1
     }
     for k, v in replacements.items():
         x = f"{k}( = [^,\n]{{1,}})?,\n"
@@ -526,7 +533,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         num_proc_check = \
         "if dataset_num_proc is None:\n"\
         "    from multiprocessing import cpu_count\n"\
-        "    dataset_num_proc = cpu_count()\n"
+        "    dataset_num_proc = min(cpu_count()*2, 2)\n"
         extra_args += num_proc_check
     pass
 

unsloth/models/vision.py

@@ -420,6 +420,7 @@ class FastBaseModel:
         torch_dtype = dtype
         if do_forced_float32: torch_dtype = torch.bfloat16
 
+        raise_handler = RaiseUninitialized()
         model = auto_model.from_pretrained(
             model_name,
             device_map              = device_map,
@@ -430,6 +431,7 @@ class FastBaseModel:
             # attn_implementation   = attn_implementation,
             **kwargs,
         )
+        raise_handler.remove()
         # Return old flag
         os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer