diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index b677f864a4..044629eaed 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1143,57 +1143,73 @@ def patch_sft_trainer_tokenizer(): 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) + for path_to_trainer in \ + ("sft_trainer.SFTTrainer", "dpo_trainer.DPOTrainer",): - check_text = \ - "\n"\ - "if self._inner_training_loop.__name__ != '_fast_inner_training_loop':\n"\ - " raise RuntimeError(\n"\ - " 'Please do not edit specific areas of the Unsloth codebase or you will get CUDA segfaults.'\n"\ - " )\n"\ - "pass\n"\ - "import subprocess, re, gc, numpy as np\n"\ - "a = np.array([0,])\n"\ - "try:\n"\ - " a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True)\n"\ - " a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a)\n"\ - " a = np.array([int(x.decode('utf-8'))/1024 for x in a])\n"\ - "except:\n"\ - " if not torch.cuda.is_available():\n"\ - " raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')\n"\ - "if ((a - PRE_CHECK) >= 1).sum() > 1:\n"\ - " raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')\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" + function_name, replacer = "train", "if resume_from_checkpoint is False:" + function = getsource(eval(f"trl.trainer.{path_to_trainer}.{function_name}")) + where = function.find("def") + function = function.split("\n") + function = "\n".join(x[where:] for x in function) - # Add NEFTune since it doesn't seem to work?? We need to manually inject it - check_text += \ - "\n"\ - "if hasattr(self, 'neftune_hook_handle'):\n"\ - " self.neftune_hook_handle.remove()\n"\ - " if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle\n"\ - "\n"\ - "if getattr(self, 'neftune_noise_alpha', None) is not None:\n"\ - " self.model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha\n"\ - " self.neftune_hook_handle = self.model.get_input_embeddings().register_forward_hook(neftune_post_forward_hook)\n"\ - "pass\n"\ - "\n" + check_text = \ + "\n"\ + "if self._inner_training_loop.__name__ != '_fast_inner_training_loop':\n"\ + " raise RuntimeError(\n"\ + " 'Please do not edit specific areas of the Unsloth codebase or you will get CUDA segfaults.'\n"\ + " )\n"\ + "pass\n"\ + "import subprocess, re, gc, numpy as np\n"\ + "a = np.array([0,])\n"\ + "try:\n"\ + " a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True)\n"\ + " a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a)\n"\ + " a = np.array([int(x.decode('utf-8'))/1024 for x in a])\n"\ + "except:\n"\ + " if not torch.cuda.is_available():\n"\ + " raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')\n"\ + "if ((a - PRE_CHECK) >= 1).sum() > 1:\n"\ + " raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')\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) + # Add NEFTune since it doesn't seem to work?? We need to manually inject it + check_text += \ + "\n"\ + "if hasattr(self, 'neftune_hook_handle'):\n"\ + " self.neftune_hook_handle.remove()\n"\ + " if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle\n"\ + "\n"\ + "if getattr(self, 'neftune_noise_alpha', None) is not None:\n"\ + " self.model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha\n"\ + " self.neftune_hook_handle = self.model.get_input_embeddings().register_forward_hook(neftune_post_forward_hook)\n"\ + "pass\n"\ + "\n" - function = function.replace(replacer, check_text + replacer) - exec(function, globals()) + # Also DPO weirdly tokenizes non numeric columns? Delete them! + check_text += \ + "\n"\ + "column_names = set(self.train_dataset.column_names)\n"\ + "check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n"\ + " 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n"\ + " 'prompt_input_ids', 'prompt_attention_mask']\n"\ + "if all(x in column_names for x in check):\n"\ + " self.train_dataset = self.train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n"\ + "del check, column_names\n"\ + "\n" - exec(f"trl.trainer.sft_trainer.SFTTrainer.{function_name} = {function_name}", globals()) + 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.{path_to_trainer}.{function_name} = {function_name}", globals()) + pass pass patch_sft_trainer_tokenizer()