ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enable vLLM to share memory space (#2712)

Browse source 
* vLLM sleep once generation is done

* Make enable_sleep_model configurable

* Make default to false

Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com>

* Force standby under environment variable

---------

Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com>
This commit is contained in:
Datta Nimmaturi 2025-06-19 16:34:14 +05:30 • committed by GitHub
commit b87ff3f528
3 changed files with 58 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/__init__.py
View file

@ -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]"

28
unsloth/models/llama.py
View file

@ -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:

45
unsloth/models/rl_replacements.py
View file

@ -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)