Fix #563 (#564)

* Update llama.py * offload * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * continued pretraining trainer * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * is_bfloat16_supported * Update __init__.py * Update README.md * Update llama.py * is_bfloat16_supported * Update __init__.py * Mistral v3 * Phi 3 medium * Update chat_templates.py * Update chat_templates.py * Phi-3 * Update save.py * Update README.md Mistral v3 to Mistral v0.3 * Untrained tokens * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update llama.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update save.py * Update save.py * Update save.py * checkpoint * Update _utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update llama.py * accelerate * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update tokenizer_utils.py * train_dataloader * Update llama.py * Update llama.py * Update llama.py * use_fast_convert --------- Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
2024-05-31 00:41:44 +10:00 · 2024-05-31 00:41:44 +10:00 · 27fa021a7b
commit 27fa021a7b
parent cd1b448786
4 changed files with 50 additions and 33 deletions
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@ -410,36 +410,50 @@ pass


 """
-    Remove warnings about missing kwargs
+    Remove warnings about missing kwargs and patch stuff
 """
-try:
-    from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
-    from inspect import getsource
-    import re
-    BitsAndBytesConfig__init__ = getsource(BitsAndBytesConfig.__init__)
-    BitsAndBytesConfig__init__ = re.sub(
-        r"if[\s]{1,}kwargs\:[\s]{1,}.+?\n",
-        "",
-        BitsAndBytesConfig__init__,
-        flags = re.MULTILINE,
-    )
-    BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.split("\n")
-    length_spaces = len(re.match(r"[\s]{1,}", BitsAndBytesConfig__init__[0]).group(0))
-    BitsAndBytesConfig__init__ = "\n".join(x[length_spaces:] for x in BitsAndBytesConfig__init__)
-    BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace(
-        "__init__",
-        "_BitsAndBytesConfig__init__",
-    )
-    exec(BitsAndBytesConfig__init__, globals())
-    
-    import transformers.utils.quantization_config
-    transformers.utils.quantization_config.BitsAndBytesConfig.__init__ = _BitsAndBytesConfig__init__
-except:
-    logger.warning_once(
-        "Unsloth unsuccessfully patched bitsandbytes. Please file a bug report.\n"\
-        "Luckily, your training run will still work in the meantime!"
-    )
+from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
+from inspect import getsource
+from accelerate.utils.dataclasses import DistributedType
+import re
+BitsAndBytesConfig__init__ = getsource(BitsAndBytesConfig.__init__)
+BitsAndBytesConfig__init__ = re.sub(
+    r"if[\s]{1,}kwargs\:[\s]{1,}.+?\n",
+    "",
+    BitsAndBytesConfig__init__,
+    flags = re.MULTILINE,
+)
+BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.split("\n")
+length_spaces = len(re.match(r"[\s]{1,}", BitsAndBytesConfig__init__[0]).group(0))
+BitsAndBytesConfig__init__ = "\n".join(x[length_spaces:] for x in BitsAndBytesConfig__init__)
+BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace(
+    "__init__",
+    "_BitsAndBytesConfig__init__",
+)
+
+def _prepare_backend(
+    self, cpu: bool = False, sagemaker_dp = False, backend: str = None,
+) -> tuple[str, DistributedType]:
+    return None, DistributedType.NO
 pass
+import accelerate.state
+accelerate.state.PartialState._prepare_backend = _prepare_backend
+
+import accelerate.accelerator
+prepare = inspect.getsource(accelerate.accelerator.Accelerator.prepare)
+prepare = prepare.split("\n")
+spaces = prepare[0].find("def")
+prepare = "\n".join(x[spaces:] for x in prepare)
+x = "for obj in args:"
+s = " "*spaces
+prepare = prepare.replace(x, f'self.state.distributed_type = DistributedType.NO\n{s}{x}', 1)
+exec(prepare, globals())
+accelerate.accelerator.Accelerator.prepare = prepare
+
+exec(BitsAndBytesConfig__init__, globals())
+
+import transformers.utils.quantization_config
+transformers.utils.quantization_config.BitsAndBytesConfig.__init__ = _BitsAndBytesConfig__init__


 # Offloading to disk for modules (lm_head, embed_tokens)
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@ -1277,6 +1277,7 @@ class FastLlamaModel:
            "is_sagemaker_mp_enabled()",
            "False",
        )
+        exec(inner_training_loop, globals())
        Trainer._inner_training_loop = _fast_inner_training_loop

        # Save max_seq_length
@ -1316,6 +1317,7 @@ class FastLlamaModel:

        # Add save modules
        patch_saving_functions(model)
+        Trainer._inner_training_loop = _fast_inner_training_loop

        # Save tokenizer for inference purposes
        tokenizer.padding_side = "left" # Force inference
@ -1336,18 +1338,18 @@ class FastLlamaModel:
        layers = model.model.layers

        # Torch.compile fails on embedding matrix??
-        # Workaround randomnly fixes it for torch versions < 2.2
-        model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight)
+        # Workaround randomnly fixes it for torch versions < 2.
+        model.set_input_embeddings(torch.nn.Embedding.from_pretrained(model.get_input_embeddings().weight))
        model.config.update({"unsloth_version" : __version__})

        # We also do this for the lm_head
        lm_head = torch.nn.Linear(1, 1, bias = None)
        del lm_head.weight
-        lm_head.weight = model.lm_head.weight
+        lm_head.weight = model.get_output_embeddings().weight
        lm_head.in_features  = lm_head.weight.shape[1]
        lm_head.out_features = lm_head.weight.shape[0]
        model.lm_head = lm_head
-
+        
        # Also patch all dtypes - BnB seems to not allocate the correct type?
        # BnB default dtype seems to be float16!
        correct_dtype = lm_head.weight.dtype
--- a/unsloth/save.py
+++ b/unsloth/save.py
@ -975,6 +975,7 @@ def save_to_gguf(
        vocab_type = "bpe"
    pass

+    use_fast_convert = False
    if use_fast_convert:
        command = f"python llama.cpp/convert.py {model_directory} "\
            f"--outfile {final_location} --vocab-type {vocab_type} "\
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@ -579,7 +579,7 @@ def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16):
    actual_bad_tokens = tokenizer.convert_ids_to_tokens(where_untrained)
    # Remove None items in actual_bad_tokens
    actual_bad_tokens = [x for x in actual_bad_tokens if x is not None]
-    
+
    # Check if tokenizer and training datasets have bad tokens
    if_bad_first  = False
    if_bad_second = False