From 5cff582ccff5899ac3f7d8423922aecbda6edea8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 29 May 2024 00:30:38 +1000 Subject: [PATCH] Nightly (#548) * 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 --------- Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> --- README.md | 4 +- unsloth/models/_utils.py | 2 +- unsloth/models/llama.py | 19 ++-- unsloth/models/loader.py | 16 +-- unsloth/save.py | 49 ++++++++- unsloth/tokenizer_utils.py | 198 ++++++++++++++++++++++++++++++------- 6 files changed, 229 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index 1d335101d8..3537b8de38 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and | Unsloth supports | Free Notebooks | Performance | Memory use | |-----------|---------|--------|----------| | **Llama 3 (8B)** | [▶️ Start for free](https://colab.research.google.com/drive/135ced7oHytdxu3N2DNe1Z0kqjyYIkDXp?usp=sharing) | 2x faster | 60% less | -| **Mistral v3 (7B)** | [▶️ Start for free](https://colab.research.google.com/drive/1_yNCks4BTD5zOnjozppphh5GzMFaMKq_?usp=sharing) | 2.2x faster | 73% less | +| **Mistral v0.3 (7B)** | [▶️ Start for free](https://colab.research.google.com/drive/1_yNCks4BTD5zOnjozppphh5GzMFaMKq_?usp=sharing) | 2.2x faster | 73% less | | **Phi-3 (medium)** | [▶️ Start for free](https://colab.research.google.com/drive/1hhdhBa1j_hsymiW9m-WzxQtgqTH_NHqi?usp=sharing) | 2x faster | 50% less | | **Phi-3 (mini)** | [▶️ Start for free](https://colab.research.google.com/drive/1lN6hPQveB_mHSnTOYifygFcrO8C1bxq4?usp=sharing) | 2x faster | 50% less | | **Gemma (7B)** | [▶️ Start for free](https://colab.research.google.com/drive/10NbwlsRChbma1v55m8LAPYG15uQv6HLo?usp=sharing) | 2.4x faster | 71% less | @@ -38,7 +38,7 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and ## 🦥 Unsloth.ai News - 📣 NEW! [Phi-3 medium](https://colab.research.google.com/drive/1hhdhBa1j_hsymiW9m-WzxQtgqTH_NHqi?usp=sharing) and [Phi-3 mini](https://colab.research.google.com/drive/1lN6hPQveB_mHSnTOYifygFcrO8C1bxq4?usp=sharing) support is here! -- 📣 NEW! [Mistral v3 Base](https://colab.research.google.com/drive/1_yNCks4BTD5zOnjozppphh5GzMFaMKq_?usp=sharing) and [Mistral v3 Instruct](https://colab.research.google.com/drive/15F1xyn8497_dUbxZP4zWmPZ3PJx1Oymv?usp=sharing) support is here! +- 📣 NEW! [Mistral v0.3 Base](https://colab.research.google.com/drive/1_yNCks4BTD5zOnjozppphh5GzMFaMKq_?usp=sharing) and [Mistral v0.3 Instruct](https://colab.research.google.com/drive/15F1xyn8497_dUbxZP4zWmPZ3PJx1Oymv?usp=sharing) support is here! - 📣 NEW! Qwen1.5-7B, Qwen1.5-14B, Qwen1.5-32B, Qwen1.5-72B now work, courtesy of Firefly's PR [#428](https://github.com/unslothai/unsloth/pull/428) - 📣 NEW! [Llama-3 8b](https://colab.research.google.com/drive/135ced7oHytdxu3N2DNe1Z0kqjyYIkDXp?usp=sharing) now works! Llama-3 70b also works (change the model name in the notebook). - 📣 NEW! [ORPO support](https://colab.research.google.com/drive/11t4njE3c4Lxl-07OD8lJSMKkfyJml3Tn?usp=sharing) is here! diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index a2d4d50c24..22fb51144c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -381,7 +381,7 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function): def forward(ctx, forward_function, hidden_states, *args): saved_hidden_states = hidden_states.to("cpu", non_blocking = True) with torch.no_grad(): - (output,) = forward_function(hidden_states, *args) + output = forward_function(hidden_states, *args) ctx.save_for_backward(saved_hidden_states) ctx.forward_function = forward_function ctx.args = args diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 1d6a282a55..9aeb55e4ea 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -657,7 +657,7 @@ def LlamaModel_fast_forward( past_key_values, output_attentions, use_cache, - ) + )[0] elif gradient_checkpointing: def create_custom_forward(module): @@ -1166,7 +1166,7 @@ class FastLlamaModel: except: raise RuntimeError( "Our OSS was designed for people with few GPU resources to level the playing field.\n" - "The OSS Apache 2 license only supports four GPUs - please obtain a commercial license from our website.\n" + "The OSS Apache 2 license only supports one GPU - please obtain a commercial license.\n" "We're a 2 person team, so we still have to fund our development costs - thanks!\n" "If you don't, please consider at least sponsoring us through Ko-fi! Appreciate it!", ) @@ -1194,7 +1194,7 @@ class FastLlamaModel: f"O^O/ \\_/ \\ Batch size per device = {self._train_batch_size:,} | Gradient Accumulation steps = {args.gradient_accumulation_steps}\\n"\\ f"\\ / Total batch size = {total_train_batch_size:,} | Total steps = {max_steps:,}\\n"\\ f' "-____-" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}' - logger.warning_once(debug_info) + logger.warning(debug_info) import gc for _ in range(3): gc.collect() @@ -1209,7 +1209,7 @@ class FastLlamaModel: if n_total_devices > 2: logger.warning_once( "Our OSS was designed for people with few GPU resources to level the playing field.\\n" - "The OSS Apache 2 license only supports four GPUs - please obtain a commercial license from our website.\\n" + "The OSS Apache 2 license only supports one GPU - please obtain a commercial license.\\n" "We're a 2 person team, so we still have to fund our development costs - thanks!\\n" "If you don't, please consider at least sponsoring us through Ko-fi! Appreciate it!", ) @@ -1238,9 +1238,10 @@ class FastLlamaModel: n_total_devices = total_batches // ga // bsz if n_total_devices > 2: logger.warning_once( - "Please consider a commercial license - Unsloth was designed for the GPU Poor.\\n" - "The OSS currently works on 4 GPUs - we're a 2 person team, so please help fund\\n" - "our development costs by supporting us through Ko-fi or buying a license! Thanks!", + "Our OSS was designed for people with few GPU resources to level the playing field.\\n" + "The OSS Apache 2 license only supports one GPU - please obtain a commercial license.\\n" + "We're a 2 person team, so we still have to fund our development costs - thanks!\\n" + "If you don't, please consider at least sponsoring us through Ko-fi! Appreciate it!", ) divisor = n_total_devices / 2 bsz = self._train_batch_size = max(int(bsz / divisor), 1) @@ -1267,7 +1268,7 @@ class FastLlamaModel: if "n_total_devices >" not in inner_training_loop: raise RuntimeError( "Our OSS was designed for people with few GPU resources to level the playing field.\n" - "The OSS Apache 2 license only supports four GPUs - please obtain a commercial license from our website.\n" + "The OSS Apache 2 license only supports one GPU - please obtain a commercial license.\n" "We're a 2 person team, so we still have to fund our development costs - thanks!\n" "If you don't, please consider at least sponsoring us through Ko-fi! Appreciate it!", ) @@ -1703,7 +1704,7 @@ class FastLlamaModel: if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop": raise RuntimeError( "Our OSS was designed for people with few GPU resources to level the playing field.\n" - "The OSS Apache 2 license only supports four GPUs - please obtain a commercial license from our website.\n" + "The OSS Apache 2 license only supports one GPU - please obtain a commercial license.\n" "We're a 2 person team, so we still have to fund our development costs - thanks!\n" "If you don't, please consider at least sponsoring us through Ko-fi! Appreciate it!", ) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 86a0f5d7bf..b2f0e4efdc 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -45,18 +45,18 @@ def _get_model_name(model_name, load_in_4bit = True): elif not load_in_4bit and model_name in INT_TO_FLOAT_MAPPER: new_model_name = INT_TO_FLOAT_MAPPER[model_name] - logger.warning_once( - f"Unsloth: You passed in `{model_name}` which is a 4bit model, yet you set\n"\ - f"`load_in_4bit = False`. We shall load `{new_model_name}` instead." - ) + # logger.warning_once( + # f"Unsloth: You passed in `{model_name}` which is a 4bit model, yet you set\n"\ + # f"`load_in_4bit = False`. We shall load `{new_model_name}` instead." + # ) model_name = new_model_name elif load_in_4bit and SUPPORTS_FOURBIT and model_name in FLOAT_TO_INT_MAPPER: new_model_name = FLOAT_TO_INT_MAPPER[model_name] - logger.warning_once( - f"Unsloth: You passed in `{model_name}` and `load_in_4bit = True`.\n"\ - f"We shall load `{new_model_name}` for 4x faster loading." - ) + # logger.warning_once( + # f"Unsloth: You passed in `{model_name}` and `load_in_4bit = True`.\n"\ + # f"We shall load `{new_model_name}` for 4x faster loading." + # ) model_name = new_model_name pass diff --git a/unsloth/save.py b/unsloth/save.py index 033f6eb1e0..7af6280910 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -94,7 +94,7 @@ def check_if_sentencepiece_model(model, temporary_location = "_unsloth_sentencep temp_tokenizer = model._saved_temp_tokenizer sentencepiece_model = False - file_location = f"{temporary_location}/{temp_tokenizer.name_or_path}" + file_location = os.path.join(temporary_location, temp_tokenizer.name_or_path) if not os.path.exists(file_location): os.makedirs(file_location) pass @@ -1390,6 +1390,22 @@ def unsloth_save_pretrained_gguf( model_type = self.config.model_type is_sentencepiece_model = check_if_sentencepiece_model(self) + + # Check if BOS added already, then warn + print_bos_token_message = False + if (tokenizer("A").input_ids[0] == getattr(tokenizer, "bos_token_id", None)): + chat_template = getattr(tokenizer, "chat_template", None) + if chat_template is not None and \ + (tokenizer.bos_token in chat_template or "{bos_token}" in chat_template.replace(" ", "")): + print_bos_token_message = True + logger.warning( + f"Unsloth: ##### The current model type of {model_type} auto adds a BOS token.\n"\ + "Unsloth: ##### If you're using Ollama or GGUF etc, do not add a BOS in the chat template." + ) + pass + pass + + # Save to GGUF file_location = save_to_gguf(model_type, is_sentencepiece_model, new_save_directory, quantization_method, first_conversion, makefile, ) @@ -1405,6 +1421,13 @@ def unsloth_save_pretrained_gguf( new_save_directory.lstrip('/.') print(f"Saved GGUF to https://huggingface.co/{link}") pass + + if print_bos_token_message: + logger.warning( + f"Unsloth: ##### The current model type of {model_type} auto adds a BOS token.\n"\ + "Unsloth: ##### If you're using Ollama or GGUF etc, do not add a BOS in the chat template." + ) + pass pass @@ -1513,6 +1536,22 @@ def unsloth_push_to_hub_gguf( model_type = self.config.model_type is_sentencepiece_model = check_if_sentencepiece_model(self) + + # Check if BOS added already, then warn + print_bos_token_message = False + if (tokenizer("A").input_ids[0] == getattr(tokenizer, "bos_token_id", None)): + chat_template = getattr(tokenizer, "chat_template", None) + if chat_template is not None and \ + (tokenizer.bos_token in chat_template or "{bos_token}" in chat_template.replace(" ", "")): + print_bos_token_message = True + logger.warning( + f"Unsloth: ##### The current model type of {model_type} auto adds a BOS token.\n"\ + "Unsloth: ##### If you're using Ollama or GGUF etc, do not add a BOS in the chat template." + ) + pass + pass + + # Save to GGUF file_location = save_to_gguf(model_type, is_sentencepiece_model, new_save_directory, quantization_method, first_conversion, makefile, ) @@ -1525,7 +1564,15 @@ def unsloth_push_to_hub_gguf( link = f"{username}/{new_save_directory.lstrip('/.')}" \ if username not in new_save_directory else \ new_save_directory.lstrip('/.') + print(f"Saved GGUF to https://huggingface.co/{link}") + + if print_bos_token_message: + logger.warning( + f"Unsloth: ##### The current model type of {model_type} auto adds a BOS token.\n"\ + "Unsloth: ##### If you're using Ollama or GGUF etc, do not add a BOS in the chat template." + ) + pass pass diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index a03491669f..03f3e341d7 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -20,6 +20,10 @@ import os from transformers.models.llama.modeling_llama import logger from peft import PeftModelForCausalLM import torch +import itertools +import collections +import numpy as np +import gc __all__ = [ "load_correct_tokenizer", @@ -274,12 +278,10 @@ def fix_sentencepiece_gguf(saved_location): user defined tokens. Inspiration from https://github.com/ggerganov/llama.cpp/blob/master/convert-hf-to-gguf.py """ - import numpy as np from copy import deepcopy from transformers.utils import sentencepiece_model_pb2 import json from enum import IntEnum - import os class SentencePieceTokenTypes(IntEnum): NORMAL = 1 @@ -554,44 +556,128 @@ pass @torch.inference_mode -def fix_untrained_tokens(model, eps = 1e-16): +def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16): """ Llama-3 for eg has untrained vectors in the base model. These include <|eot_id|>, <|start_header_id|>, <|end_header_id|> We reset them to the mean of the rest of the tokens """ - embedding_matrix = model.get_input_embeddings ().weight.data - lm_head_matrix = model.get_output_embeddings().weight.data + embedding_matrix = model.get_input_embeddings ().weight + lm_head_matrix = model.get_output_embeddings().weight # Get untrained tokens indicator_untrained = torch.amax(embedding_matrix, axis = 1) <= eps where_untrained = torch.where(indicator_untrained)[0] n_untrained = where_untrained.shape[0] n_trained = embedding_matrix.shape[0] - n_untrained - if n_untrained != 0: - print( - f"Unsloth: Not an error, but your model has {n_untrained} untrained tokens.\n"\ - "We shall set them to the mean of the other trained tokens." + + # Get set and actual tokens + where_untrained = where_untrained.tolist() + if len(where_untrained) == 0: return + + where_untrained_set = frozenset(where_untrained) + actual_bad_tokens = tokenizer.convert_ids_to_tokens(where_untrained) + + # Check if tokenizer and training datasets have bad tokens + if_bad_first = False + if_bad_second = False + # Check tokenizer's chat template for any untrained tokens + chat_template = getattr(tokenizer, "chat_template", None) + if chat_template is not None: + if_bad_first = any(x in chat_template for x in actual_bad_tokens) + pass + + # Check the first 250, last 250 input_ids + size_dataset = len(train_dataset) + size = min(size_dataset, 250) + for j in range(size): + input_ids = train_dataset[j] + if "input_ids" in input_ids: + input_ids = input_ids["input_ids"] + if_bad = any(item in where_untrained_set for item in input_ids) + if if_bad: + if_bad_second = True + break + pass + pass + pass + + # Check last 250 + if not if_bad_second: + left = max(size_dataset-250, 0) + for j in range(left, size_dataset): + input_ids = train_dataset[j] + if "input_ids" in input_ids: + input_ids = input_ids["input_ids"] + if_bad = any(item in where_untrained_set for item in input_ids) + if if_bad: + if_bad_second = True + break + pass + pass + pass + pass + + # Check if bad tokens exists! + if not if_bad_first and not if_bad_second: return + + # Check if lm_head / embed_token are trainable! + bad_not_trainable = False + if not embedding_matrix.requires_grad: bad_not_trainable = True + if not lm_head_matrix .requires_grad: bad_not_trainable = True + + if bad_not_trainable: + raise ValueError( + 'Unsloth: Untrained tokens found, but embed_tokens & lm_head not trainable, causing NaNs. '\ + 'Restart then add `embed_tokens` & `lm_head` to '\ + '`FastLanguageModel.get_peft_model(target_modules = [..., "embed_tokens", "lm_head",])`', ) pass - # First set untrained to all 0s - sometimes it's not! 1e-23 for bfloat16 - embedding_matrix[where_untrained] = 0 - lm_head_matrix [where_untrained] = 0 + # Count all the possible bad tokens + final_counts = np.zeros(len(tokenizer), dtype = np.int64) + def mapping(examples): + input_ids = examples["input_ids"] + counter = np.fromiter(itertools.chain.from_iterable(input_ids), dtype = np.int32) + np.add.at(final_counts, counter, 1) + pass + train_dataset.map(mapping, batched = True, desc = "Counting untrained tokens") - # Find sum - sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0) - sum_lm_head = torch.sum(lm_head_matrix, dtype = torch.float32, axis = 0) + # Get sum of all items + sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0) + sum_lm_head = torch.sum(lm_head_matrix, dtype = torch.float32, axis = 0) + + # Remove bad tokens + sum_embedding -= torch.sum(embedding_matrix[where_untrained], dtype = torch.float32, axis = 0) + sum_lm_head -= torch.sum(lm_head_matrix [where_untrained], dtype = torch.float32, axis = 0) # Find correct average by dividing by sum of trained tokens - mean_embedding = (sum_embedding / n_trained).to(embedding_matrix.dtype) - mean_lm_head = (sum_lm_head / n_trained).to(lm_head_matrix .dtype) + mean_embedding = (sum_embedding / n_trained) + mean_lm_head = (sum_lm_head / n_trained) + + # Scale each to be equal to 1/max_frequency. Also set some to 0 if none seen + scaling = final_counts[where_untrained] / max(final_counts.max(), 1) + scaling = torch.tensor(scaling, device = mean_embedding.device).unsqueeze(1) + mean_embedding = mean_embedding.repeat((n_untrained, 1,)) * scaling + mean_lm_head = mean_lm_head .repeat((n_untrained, 1,)) * scaling + where_null = scaling.ravel() == 0 + mean_embedding[where_null] = 0 + mean_lm_head [where_null] = 0 # Set them to the mean - embedding_matrix[where_untrained] = mean_embedding - lm_head_matrix [where_untrained] = mean_lm_head + logger.warning( + "Unsloth: Setting embed_tokens & lm_head untrained tokens to "\ + "mean(trained) to counteract NaNs during training." + ) + embedding_matrix[where_untrained] = mean_embedding.to(embedding_matrix.dtype) + lm_head_matrix [where_untrained] = mean_lm_head .to(lm_head_matrix .dtype) - return mean_embedding, mean_lm_head + # Clean up + for _ in range(3): + gc.collect() + torch.cuda.empty_cache() + pass + return pass @@ -602,32 +688,32 @@ def mean_of_trained_tokens(model, eps = 1e-16): These include <|eot_id|>, <|start_header_id|>, <|end_header_id|> We reset them to the mean of the rest of the tokens """ - embedding_matrix = model.get_input_embeddings ().weight.data.clone() - lm_head_matrix = model.get_output_embeddings().weight.data.clone() + embedding_matrix = model.get_input_embeddings ().weight.clone() + lm_head_matrix = model.get_output_embeddings().weight.clone() # Get untrained tokens indicator_untrained = torch.amax(embedding_matrix, axis = 1) <= eps where_untrained = torch.where(indicator_untrained)[0] n_untrained = where_untrained.shape[0] n_trained = embedding_matrix.shape[0] - n_untrained - if n_untrained != 0: - print( - f"Unsloth: Not an error, but your model has {n_untrained} untrained tokens.\n"\ - "We shall set them to the mean of the other trained tokens." - ) - pass + # if n_untrained != 0: + # print( + # f"Unsloth: Not an error, but your model has {n_untrained} untrained tokens.\n"\ + # "We shall set them to the mean of the other trained tokens." + # ) + # pass - # First set untrained to all 0s - sometimes it's not! 1e-23 for bfloat16 - embedding_matrix[where_untrained] = 0 - lm_head_matrix [where_untrained] = 0 + # Get sum of all items + sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0) + sum_lm_head = torch.sum(lm_head_matrix, dtype = torch.float32, axis = 0) - # Find sum - sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0) - sum_lm_head = torch.sum(lm_head_matrix, dtype = torch.float32, axis = 0) + # Remove bad tokens + sum_embedding -= torch.sum(embedding_matrix[where_untrained], dtype = torch.float32, axis = 0) + sum_lm_head -= torch.sum(lm_head_matrix [where_untrained], dtype = torch.float32, axis = 0) # Find correct average by dividing by sum of trained tokens - mean_embedding = (sum_embedding / n_trained).to(embedding_matrix.dtype) - mean_lm_head = (sum_lm_head / n_trained).to(lm_head_matrix .dtype) + mean_embedding = (sum_embedding / n_trained) + mean_lm_head = (sum_lm_head / n_trained) return mean_embedding, mean_lm_head pass @@ -676,8 +762,8 @@ def add_new_tokens( # If we use interpolation, we interpolate between the mean embeddings and # the Word2Vec sum of the other vectors - embedding_matrix = model.get_input_embeddings ().weight.data - lm_head_matrix = model.get_output_embeddings().weight.data + embedding_matrix = model.get_input_embeddings ().weight + lm_head_matrix = model.get_output_embeddings().weight if method == "interpolation": print( @@ -718,6 +804,7 @@ pass from inspect import getsource import trl.trainer.sft_trainer from trl.trainer.sft_trainer import * +from transformers.trainer import * def patch_sft_trainer_tokenizer(): """ @@ -749,6 +836,41 @@ def patch_sft_trainer_tokenizer(): exec(f"trl.trainer.sft_trainer.SFTTrainer.{function_name} = {function_name}", globals()) pass + + # Patch train with fix_untrained_tokens + function_name, replacer = "train", "if resume_from_checkpoint is False:" + function = getsource(eval(f"trl.trainer.sft_trainer.SFTTrainer.{function_name}")) + where = function.find("def") + function = function.split("\n") + function = "\n".join(x[where:] for x in function) + + check_text = \ + "\n"\ + "if self._inner_training_loop.__name__ != '_fast_inner_training_loop':\n"\ + " raise RuntimeError(\n"\ + " 'Do not edit specific areas of the Unsloth codebase or you will get CUDA segfaults.'\n"\ + " )\n"\ + "pass\n"\ + "n_devices = torch.cuda.device_count()\n"\ + "more_than = 0\n"\ + "for j in range(n_devices):\n"\ + " vram = torch.cuda.max_memory_reserved(torch.cuda.device(j)) / 1024 / 1024 / 1024\n"\ + " more_than += (vram > 4)\n"\ + "if more_than > 1: raise RuntimeError('Error: More than 1 GPUs have a lot of VRAM usage.')\n"\ + "for _ in range(3):\n"\ + " gc.collect()\n"\ + " torch.cuda.empty_cache()\n"\ + "pass\n"\ + "\n"\ + "fix_untrained_tokens(self.model, self.tokenizer, self.train_dataset, eps = 1e-16)\n\n" + + check_text = check_text.split("\n") + check_text = "\n".join(" "*where + x for x in check_text) + + function = function.replace(replacer, check_text + replacer) + exec(function, globals()) + + exec(f"trl.trainer.sft_trainer.SFTTrainer.{function_name} = {function_name}", globals()) pass patch_sft_trainer_tokenizer()