diff --git a/unsloth/__init__.py b/unsloth/__init__.py index cf4dbe09c0..58c2a19814 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -89,7 +89,7 @@ DEVICE_TYPE : str = get_device_type() # Reduce VRAM usage by reducing fragmentation # And optimize pinning of memory -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="0": os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \ "expandable_segments:True,"\ "roundup_power2_divisions:[32:256,64:128,256:64,>:32]" diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9c2cdb5689..125bc7e610 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1698,18 +1698,18 @@ class FastLlamaModel: @staticmethod def from_pretrained( - model_name = "unsloth/llama-3-8b-bnb-4bit", - max_seq_length = None, - dtype = None, - load_in_4bit = True, - token = None, - device_map = "sequential", - rope_scaling = None, - fix_tokenizer = True, - model_patcher = None, - tokenizer_name = None, - trust_remote_code = False, - revision = None, + model_name = "unsloth/llama-3-8b-bnb-4bit", + max_seq_length = None, + dtype = None, + load_in_4bit = True, + token = None, + device_map = "sequential", + rope_scaling = None, + fix_tokenizer = True, + model_patcher = None, + tokenizer_name = None, + trust_remote_code = False, + revision = None, fast_inference = False, # uses vLLM gpu_memory_utilization = 0.5, @@ -1717,6 +1717,7 @@ class FastLlamaModel: random_state = 3407, max_lora_rank = 16, disable_log_stats = False, + unsloth_vllm_standby = False, num_labels = None, **kwargs, ): @@ -1737,6 +1738,8 @@ class FastLlamaModel: if major_version < 7: print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!") fast_inference = False + if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "0": + raise RuntimeError("Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!") pass if token is None: token = get_token() @@ -1898,6 +1901,7 @@ class FastLlamaModel: max_lora_rank = max_lora_rank, disable_log_stats = disable_log_stats, use_bitsandbytes = load_in_4bit, + unsloth_vllm_standby = unsloth_vllm_standby, ) for allowed_arg in allowed_args: if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs: diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 171e75d197..30dc0d220d 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -171,24 +171,61 @@ RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss) def grpo_trainer__prepare_inputs(function_name, function): if function_name != "_prepare_inputs": return function - if "with torch.inference_mode()" not in function: return function + import re + # Try to find the function signature and insert after it + # This matches the function signature and any decorators/comments, then finds the first non-empty line after the signature + pattern = r"(def _prepare_inputs\s*\([^\)]*\)\s*(->\s*[^:]+)?\s*:\s*\n)" + match = re.search(pattern, function) + if match: + sig_end = match.end(1) + rest = function[sig_end:] + rest = re.sub(r"^[ \t]*self\.llm\.wake_up\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*torch\.cuda\.empty_cache\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest) + insert = ( + " if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" + " self.llm.wake_up()\n" + ) + function = function[:sig_end] + insert + rest + else: + pattern2 = r"(def _prepare_inputs\(.*?\):\n(?:[ ]+#[^\n]*\n)+)" + match2 = re.search(pattern2, function, flags=re.DOTALL) + if match2: + header_and_comments = match2.group(1) + rest = function[len(header_and_comments):] + rest = re.sub(r"^[ \t]*self\.llm\.wake_up\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*torch\.cuda\.empty_cache\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest) + rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest) + insert = ( + " if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" + " self.llm.wake_up()\n" + ) + function = header_and_comments + insert + rest # Add mixed precision training function = function.replace( "with torch.inference_mode():", - "with torch.inference_mode(), "\ "torch.amp.autocast(device_type = 'cuda', "\ "dtype = ((torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16) "\ "if not torch.is_autocast_enabled('cuda') else nullcontext())"\ "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '0' else torch.float16):", ) - - # Disable attaching a float32 conversion hook which upcasts logits to FP32 function = function.replace( "self.accelerator.unwrap_model(self.model)", "self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False)", ) + sleep_and_cache = ( + "if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" + " self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n" + " " + ) + if re.search(r"\n\s*return ", function): + function = re.sub(r"(\n\s*)return ", f"\\1{sleep_and_cache}return ", function, count=1) + else: + function = function.rstrip() + "\n " + sleep_and_cache return function pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs)