GRPO, vLLM, Bug Fixes, Reinforcement Learning (#1620)
* use exact model name
* Update save.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* print
* Update _utils.py
* Update _utils.py
* Update llama.py
* Update _utils.py
* Update vision.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update loader.py
* accurate_accumulation
* Update loader.py
* Update loader.py
* Update _utils.py
* Update loader.py
* Update loader.py
* Update loader.py
* Update loader.py
* Update pyproject.toml
* Update __init__.py
* Update pyproject.toml
* Update __init__.py
* Update __init__.py
* Fix Triton heuristics
https://github.com/triton-lang/triton/issues/5224
* Update __init__.py
* Update __init__.py
* Update __init__.py
* Update __init__.py
* Xformers
* Update loader.py
* Update loader.py
* Rewind
* Update _utils.py
* Update _utils.py
* requires grad
* Update loader.py
* Update _utils.py
* Update loader.py
* changing model to base_model if peft model is already used
* Improve debugging experience (#1512)
* Create CONTRIBUTING.md (#1472)
Creating contributing guidelines
* Update CONTRIBUTING.md
improved sentence
* Improve logging control in `unsloth_compile_transformers` by conditionally redirecting stdout based on UNSLOTH_DISABLE_LOGGER environment variable
---------
Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: Nino Risteski <95188570+NinoRisteski@users.noreply.github.com>
* Update loader.py
* Update llama.py
* Update llama.py
* Revert "Update llama.py"
This reverts commit 67bb995878.
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Auto change is_bfloat16_supported
* Update llama.py
* Force data-type
* Update llama.py
* All attention refactor fix (#1491)
* change initilization of n_heads, n_kv_heads, hidden_size in llama.py
* do the same for cohere, mistral, gemma2, granite
* do the same for flexattention,cohere, mistral, granite
* Update llama.py
* Update llama.py
* Update granite to work with latest post_patch methods (#1502)
* Update granite to work with latest post_patch methods
* Pass position_embeddings for granite even if transformers<4.47
* Update llama.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Minor fixes for granite models (#1503)
* Update granite.py
Grab residual multiplier directly from layer
* Update llama.py
Version should read >= 4.47.1 as that is the version requiring the changes
* Update granite.py
* Update llama.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* support modelscope models and datasets (#1481)
* support modelscope
* change modelscope args
* remove useless import
* remove useless import
* fix
* wip
* fix
* remove useless code
* add readme
* add some comments
* change print to raise error
* update comment
* Update loader.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Merge branch 'main' into nightly
* Phi 4
* Update llama.py
* Torch.Cuda Is Available Condition and Warning (#1545)
* check for torch.cuda and triton if available
on my machine(mac m3) the cuda were not available
* Update pyproject.toml
* Update __init__.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Update mistral.py
* Update mistral.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Update _utils.py
* Fix
* Bug fixes
* Update mapper.py
* Add dropout to granite to match HF's implementation (#1557)
Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com>
* Update llama.py
* Update llama.py
* Bug fixes
* fix: flash_attn_detection_error (#1556)
* fix: flash_attn_detection_error
* Update _utils.py
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Update mapper.py
* Update gemma.py
* Update gemma.py
* Update gemma.py
* Update gemma.py
* dim fix
* Update _utils.py
* Torch 2.6 support
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Update llama.py
* Faster inference?
* Update llama.py
* Update llama.py
* Update utils.py
* Update llama.py
* Update llama.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update utils.py
* Update mapper.py
* Fast Inference via vLLM
* Update llama.py
* Update llama.py
* Update utils.py
* Create rl.py
* PatchRL
* Update rl.py
* Update rl.py
* Update rl.py
* PatchRLStatistics
* Update rl.py
* Update rl.py
* Update rl.py
* Update utils.py
* Update utils.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* RL metrics
* Update rl.py
* RL metrics
* Update __init__.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update chat_templates.py
* Update mapper.py
* Fp8 cache
* Update llama.py
* Update llama.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update rl.py
* Update __init__.py
* Update loader.py
---------
Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com>
Co-authored-by: Itsuro Tajima <tajima@georepublic.de>
Co-authored-by: Muhammad Osama <muhammadosama1994@gmail.com>
Co-authored-by: Edd <68678137+Erland366@users.noreply.github.com>
Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: Nino Risteski <95188570+NinoRisteski@users.noreply.github.com>
Co-authored-by: Kareem <81531392+KareemMusleh@users.noreply.github.com>
Co-authored-by: Datta Nimmaturi <datta.nimmaturi@nutanix.com>
Co-authored-by: Z <coffeevampirebusiness@gmail.com>
Co-authored-by: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com>
Co-authored-by: AminWhat <88392440+aminwhat@users.noreply.github.com>
Co-authored-by: Zhe Zhang <2631992879@qq.com>
This commit is contained in:
Daniel Han
GitHub
parent
8bc38fe60c
commit
4a3b424e8d
13 changed files with 918 additions and 283 deletions
111
pyproject.toml
111
pyproject.toml
|
|
@ -131,6 +131,12 @@ cu124onlytorch240 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu118onlytorch250 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.28.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu121onlytorch250 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.28.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
|
|
@ -147,6 +153,12 @@ cu124onlytorch250 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.28.post2-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu118onlytorch251 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post1-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu121onlytorch251 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
|
|
@ -163,6 +175,28 @@ cu124onlytorch251 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu118onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu124onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
|
||||
]
|
||||
cu126onlytorch260 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'",
|
||||
]
|
||||
cu118 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
|
|
@ -223,21 +257,31 @@ cu121-torch240 = [
|
|||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu121onlytorch240]",
|
||||
]
|
||||
cu121-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu121onlytorch250]",
|
||||
]
|
||||
cu124-torch240 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu124onlytorch240]",
|
||||
]
|
||||
cu118-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu118onlytorch250]",
|
||||
]
|
||||
cu121-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu121onlytorch250]",
|
||||
]
|
||||
cu124-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu124onlytorch250]",
|
||||
]
|
||||
cu118-torch251 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu118onlytorch251]",
|
||||
]
|
||||
cu121-torch251 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
|
|
@ -248,6 +292,21 @@ cu124-torch251 = [
|
|||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu124onlytorch251]",
|
||||
]
|
||||
cu118-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu118onlytorch260]",
|
||||
]
|
||||
cu124-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu124onlytorch260]",
|
||||
]
|
||||
cu126-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu126onlytorch260]",
|
||||
]
|
||||
kaggle = [
|
||||
"unsloth[huggingface]",
|
||||
]
|
||||
|
|
@ -381,24 +440,36 @@ cu121-ampere-torch240 = [
|
|||
"unsloth[cu121onlytorch240]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu121-ampere-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu121onlytorch250]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu124-ampere-torch240 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu124onlytorch240]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu118-ampere-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu118onlytorch250]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu121-ampere-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu121onlytorch250]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu124-ampere-torch250 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu124onlytorch250]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu118-ampere-torch251 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
"unsloth[cu118onlytorch251]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu121-ampere-torch251 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.43.3",
|
||||
|
|
@ -411,6 +482,24 @@ cu124-ampere-torch251 = [
|
|||
"unsloth[cu124onlytorch251]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu118-ampere-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu118onlytorch260]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu124-ampere-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu124onlytorch260]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
cu126-ampere-torch260 = [
|
||||
"unsloth[huggingface]",
|
||||
"bitsandbytes>=0.45.1",
|
||||
"unsloth[cu126onlytorch260]",
|
||||
"unsloth[flashattention]",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
homepage = "http://www.unsloth.ai"
|
||||
|
|
|
|||
|
|
@ -196,7 +196,7 @@ pass
|
|||
# Check for unsloth_zoo
|
||||
try:
|
||||
unsloth_zoo_version = importlib_version("unsloth_zoo")
|
||||
if Version(unsloth_zoo_version) < Version("2025.1.4"):
|
||||
if Version(unsloth_zoo_version) < Version("2025.2.1"):
|
||||
try:
|
||||
os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo")
|
||||
except:
|
||||
|
|
|
|||
|
|
@ -18,14 +18,16 @@ from packaging.version import Version as V
|
|||
v = V(torch.__version__)
|
||||
cuda = str(torch.version.cuda)
|
||||
is_ampere = torch.cuda.get_device_capability()[0] >= 8
|
||||
if cuda != "12.1" and cuda != "11.8" and cuda != "12.4": raise RuntimeError(f"CUDA = {cuda} not supported!")
|
||||
if cuda != "12.1" and cuda != "11.8" and cuda != "12.4" and cuda != "12.6": raise RuntimeError(f"CUDA = {cuda} not supported!")
|
||||
if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
|
||||
elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
|
||||
elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
|
||||
elif v < V('2.3.0'): x = 'cu{}{}-torch220'
|
||||
elif v < V('2.4.0'): x = 'cu{}{}-torch230'
|
||||
elif v < V('2.5.0'): x = 'cu{}{}-torch240'
|
||||
elif v < V('2.6.0'): x = 'cu{}{}-torch250'
|
||||
elif v < V('2.5.1'): x = 'cu{}{}-torch250'
|
||||
elif v <= V('2.5.1'): x = 'cu{}{}-torch251'
|
||||
elif v < V('2.7.0'): x = 'cu{}{}-torch260'
|
||||
else: raise RuntimeError(f"Torch = {v} too new!")
|
||||
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
|
||||
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
|
||||
|
|
@ -759,6 +759,10 @@ DEFAULT_SYSTEM_MESSAGE["llama-3.1"] = "" # Llama3.1 default system message is em
|
|||
|
||||
CHAT_TEMPLATES["llama-31"] = (llama31_template, llama31_template_eos_token, False, llama31_ollama,)
|
||||
DEFAULT_SYSTEM_MESSAGE["llama-31"] = "" # Llama3.1 default system message is empty + the dates
|
||||
|
||||
for version in ("llama-3.2", "llama-3.3", "llama-32", "llama-33"):
|
||||
CHAT_TEMPLATES[version] = CHAT_TEMPLATES["llama-3.1"]
|
||||
DEFAULT_SYSTEM_MESSAGE[version] = ""
|
||||
pass
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@
|
|||
import triton
|
||||
MAX_FUSED_SIZE : int = 65536
|
||||
next_power_of_2 = triton.next_power_of_2
|
||||
import functools
|
||||
|
||||
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
|
||||
import torch
|
||||
|
|
@ -66,6 +67,8 @@ global CUDA_STREAM
|
|||
CUDA_STREAM = None
|
||||
get_ptr = bnb.functional.get_ptr
|
||||
import ctypes
|
||||
ctypes_c_int = ctypes.c_int
|
||||
ctypes_c_int32 = ctypes.c_int32
|
||||
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
|
||||
cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
|
||||
cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
|
||||
|
|
@ -98,25 +101,31 @@ pass
|
|||
|
||||
def get_lora_parameters_bias(proj):
|
||||
# For DPO or disabled adapters
|
||||
base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj)
|
||||
base_layer = getattr(proj, "base_layer", proj) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
|
||||
W = base_layer.weight
|
||||
bias = base_layer.bias
|
||||
|
||||
if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
|
||||
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
|
||||
if getattr(proj, "disable_adapters", True) or proj.merged:
|
||||
return W, QUANT_STATE(W), None, None, None, bias
|
||||
pass
|
||||
|
||||
active_adapter = proj.active_adapters[0] if \
|
||||
hasattr(proj, "active_adapters") else proj.active_adapter
|
||||
getattr(proj, "active_adapters", ) else proj.active_adapter
|
||||
A = proj.lora_A [active_adapter].weight
|
||||
B = proj.lora_B [active_adapter].weight
|
||||
s = proj.scaling[active_adapter]
|
||||
return W, QUANT_STATE(W), A, B, s, bias
|
||||
pass
|
||||
|
||||
global WEIGHT_BUFFER
|
||||
WEIGHT_BUFFER = None
|
||||
global ABSMAX_BUFFER
|
||||
ABSMAX_BUFFER = None
|
||||
|
||||
if HAS_CUDA_STREAM:
|
||||
def fast_dequantize(W, quant_state = None, out = None):
|
||||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
if quant_state is None: return W
|
||||
if type(quant_state) is not list:
|
||||
# New quant_state as a class
|
||||
|
|
@ -139,36 +148,54 @@ if HAS_CUDA_STREAM:
|
|||
global CUDA_STREAM
|
||||
if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0")
|
||||
|
||||
n_elements_absmax = absmax.numel()
|
||||
|
||||
# Create weight matrix
|
||||
if out is None:
|
||||
out = torch.empty(shape, dtype = dtype, device = "cuda:0")
|
||||
if use_global_buffer:
|
||||
|
||||
# Use same buffers for faster inference
|
||||
size = shape[0]*shape[1]
|
||||
global WEIGHT_BUFFER
|
||||
global ABSMAX_BUFFER
|
||||
if WEIGHT_BUFFER is None:
|
||||
WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0", requires_grad = False)
|
||||
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False)
|
||||
|
||||
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
|
||||
if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax)
|
||||
|
||||
out = WEIGHT_BUFFER[:size].view(shape)
|
||||
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
|
||||
else:
|
||||
assert(out.shape == shape)
|
||||
assert(out.dtype == dtype)
|
||||
if out is None:
|
||||
out = torch.empty(shape, dtype = dtype, device = "cuda:0", requires_grad = False)
|
||||
else:
|
||||
assert(out.shape == shape)
|
||||
assert(out.dtype == dtype)
|
||||
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False)
|
||||
pass
|
||||
|
||||
# NF4 dequantization of statistics
|
||||
n_elements_absmax = absmax.numel()
|
||||
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
|
||||
|
||||
# Do dequantization
|
||||
ptr_out_absmax = get_ptr(out_absmax)
|
||||
cdequantize_blockwise_fp32(
|
||||
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax,
|
||||
ctypes.c_int(blocksize2), ctypes.c_int(n_elements_absmax), CUDA_STREAM,
|
||||
ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), CUDA_STREAM,
|
||||
)
|
||||
out_absmax += offset
|
||||
|
||||
# Dequantize W
|
||||
fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \
|
||||
cdequantize_blockwise_bf16_nf4
|
||||
fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out),
|
||||
ctypes.c_int(blocksize), ctypes.c_int(out.numel()), CUDA_STREAM,)
|
||||
ctypes_c_int(blocksize), ctypes_c_int(out.numel()), CUDA_STREAM,)
|
||||
|
||||
# Careful returning transposed data
|
||||
is_transposed = (True if W.shape[0] == 1 else False)
|
||||
return out.t() if is_transposed else out
|
||||
pass
|
||||
else:
|
||||
def fast_dequantize(W, quant_state = None, out = None):
|
||||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
if quant_state is None: return W
|
||||
if type(quant_state) is not list:
|
||||
# New quant_state as a class
|
||||
|
|
@ -189,29 +216,45 @@ else:
|
|||
absmax2, code2, blocksize2, _, _, _, _ = state2
|
||||
pass
|
||||
|
||||
# Create weight matrix
|
||||
if out is None:
|
||||
out = torch.empty(shape, dtype = dtype, device = "cuda:0")
|
||||
else:
|
||||
assert(out.shape == shape)
|
||||
assert(out.dtype == dtype)
|
||||
|
||||
# NF4 dequantization of statistics
|
||||
n_elements_absmax = absmax.numel()
|
||||
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
|
||||
|
||||
# Create weight matrix
|
||||
if use_global_buffer:
|
||||
|
||||
# Use same buffers for faster inference
|
||||
size = shape[0]*shape[1]
|
||||
global WEIGHT_BUFFER
|
||||
global ABSMAX_BUFFER
|
||||
if WEIGHT_BUFFER is None:
|
||||
WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0", requires_grad = False)
|
||||
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = dtype, device = "cuda:0", requires_grad = False)
|
||||
|
||||
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
|
||||
if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax)
|
||||
|
||||
out = WEIGHT_BUFFER[:size].view(shape)
|
||||
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
|
||||
else:
|
||||
if out is None:
|
||||
out = torch.empty(shape, dtype = dtype, device = "cuda:0", requires_grad = False)
|
||||
else:
|
||||
assert(out.shape == shape)
|
||||
assert(out.dtype == dtype)
|
||||
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False)
|
||||
pass
|
||||
|
||||
# Do dequantization
|
||||
ptr_out_absmax = get_ptr(out_absmax)
|
||||
cdequantize_blockwise_fp32(
|
||||
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax,
|
||||
ctypes.c_int(blocksize2), ctypes.c_int(n_elements_absmax),
|
||||
ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax),
|
||||
)
|
||||
out_absmax += offset
|
||||
|
||||
fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \
|
||||
cdequantize_blockwise_bf16_nf4
|
||||
fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out),
|
||||
ctypes.c_int(blocksize), ctypes.c_int(out.numel()),)
|
||||
ctypes_c_int(blocksize), ctypes_c_int(out.numel()),)
|
||||
|
||||
# Careful returning transposed data
|
||||
is_transposed = (True if W.shape[0] == 1 else False)
|
||||
|
|
@ -263,17 +306,17 @@ if HAS_CUDA_STREAM:
|
|||
lda = shape[0]
|
||||
ldc = shape[0]
|
||||
ldb = (hd+1)//2
|
||||
m = ctypes.c_int32(m)
|
||||
n = ctypes.c_int32(n)
|
||||
k = ctypes.c_int32(k)
|
||||
lda = ctypes.c_int32(lda)
|
||||
ldb = ctypes.c_int32(ldb)
|
||||
ldc = ctypes.c_int32(ldc)
|
||||
m = ctypes_c_int32(m)
|
||||
n = ctypes_c_int32(n)
|
||||
k = ctypes_c_int32(k)
|
||||
lda = ctypes_c_int32(lda)
|
||||
ldb = ctypes_c_int32(ldb)
|
||||
ldc = ctypes_c_int32(ldc)
|
||||
|
||||
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
|
||||
cdequantize_blockwise_fp32(
|
||||
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df),
|
||||
ctypes.c_int(blocksize2), ctypes.c_int(df.numel()), CUDA_STREAM,
|
||||
ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), CUDA_STREAM,
|
||||
)
|
||||
df += offset
|
||||
absmax = df
|
||||
|
|
@ -281,7 +324,7 @@ if HAS_CUDA_STREAM:
|
|||
fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \
|
||||
cgemm_4bit_inference_naive_bf16
|
||||
|
||||
blocksize = ctypes.c_int32(blocksize)
|
||||
blocksize = ctypes_c_int32(blocksize)
|
||||
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
|
||||
lda, ldb, ldc, blocksize, CUDA_STREAM,)
|
||||
|
||||
|
|
@ -327,17 +370,17 @@ else:
|
|||
lda = shape[0]
|
||||
ldc = shape[0]
|
||||
ldb = (hd+1)//2
|
||||
m = ctypes.c_int32(m)
|
||||
n = ctypes.c_int32(n)
|
||||
k = ctypes.c_int32(k)
|
||||
lda = ctypes.c_int32(lda)
|
||||
ldb = ctypes.c_int32(ldb)
|
||||
ldc = ctypes.c_int32(ldc)
|
||||
m = ctypes_c_int32(m)
|
||||
n = ctypes_c_int32(n)
|
||||
k = ctypes_c_int32(k)
|
||||
lda = ctypes_c_int32(lda)
|
||||
ldb = ctypes_c_int32(ldb)
|
||||
ldc = ctypes_c_int32(ldc)
|
||||
|
||||
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
|
||||
cdequantize_blockwise_fp32(
|
||||
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df),
|
||||
ctypes.c_int(blocksize2), ctypes.c_int(df.numel()),
|
||||
ctypes_c_int(blocksize2), ctypes_c_int(df.numel()),
|
||||
)
|
||||
df += offset
|
||||
absmax = df
|
||||
|
|
@ -345,7 +388,7 @@ else:
|
|||
fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \
|
||||
cgemm_4bit_inference_naive_bf16
|
||||
|
||||
blocksize = ctypes.c_int32(blocksize)
|
||||
blocksize = ctypes_c_int32(blocksize)
|
||||
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
|
||||
lda, ldb, ldc, blocksize,)
|
||||
|
||||
|
|
@ -354,6 +397,9 @@ else:
|
|||
pass
|
||||
|
||||
|
||||
torch_mm = torch.mm
|
||||
torch_mv = torch.mv
|
||||
torch_matmul = torch.matmul
|
||||
def fast_linear_forward(proj, X, temp_lora = None, out = None):
|
||||
|
||||
W, W_quant, lora_A, lora_B, lora_S, bias = get_lora_parameters_bias(proj)
|
||||
|
|
@ -361,12 +407,12 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
|
|||
if q_len != 1: return matmul_lora(X, W, W_quant, lora_A, lora_B, lora_S)
|
||||
|
||||
if W_quant is None:
|
||||
out = torch.matmul(X, W.t(), out = out)
|
||||
out = torch_matmul(X, W.t(), out = out)
|
||||
elif bsz == 1 and q_len == 1:
|
||||
out = fast_gemv(X, W, W_quant, out = out)
|
||||
else:
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
out = torch.matmul(X, W, out = out)
|
||||
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
|
||||
out = torch_matmul(X, W, out = out)
|
||||
pass
|
||||
|
||||
# Add in LoRA weights
|
||||
|
|
@ -381,11 +427,11 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
|
|||
|
||||
if bsz == 1:
|
||||
out = out.view(out_dim)
|
||||
temp_lora = torch.mv(lora_A._fast_lora, X.ravel(), out = temp_lora)
|
||||
temp_lora = torch_mv(lora_A._fast_lora, X.ravel(), out = temp_lora)
|
||||
out.addmv_(lora_B._fast_lora, temp_lora, alpha = lora_S)
|
||||
else:
|
||||
out = out.view(bsz, out_dim)
|
||||
temp_lora = torch.mm(X.view(bsz, in_dim), lora_A._fast_lora.t(), out = temp_lora)
|
||||
temp_lora = torch_mm(X.view(bsz, in_dim), lora_A._fast_lora.t(), out = temp_lora)
|
||||
out.addmm_(temp_lora, lora_B._fast_lora.t(), alpha = lora_S)
|
||||
pass
|
||||
out = out.view(bsz, 1, out_dim)
|
||||
|
|
@ -399,7 +445,7 @@ pass
|
|||
|
||||
def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
||||
dtype = X.dtype
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
|
||||
|
||||
if X.dim() == 3:
|
||||
batch, seq_len, d = X.shape
|
||||
|
|
@ -409,7 +455,7 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
|||
reshape = False
|
||||
pass
|
||||
|
||||
out = torch.matmul(X, W, out = out)
|
||||
out = torch_matmul(X, W, out = out)
|
||||
if W_quant is not None: del W
|
||||
|
||||
if A is not None:
|
||||
|
|
|
|||
|
|
@ -20,3 +20,4 @@ from .mistral import FastMistralModel
|
|||
from .qwen2 import FastQwen2Model
|
||||
from .dpo import PatchDPOTrainer, PatchKTOTrainer
|
||||
from ._utils import is_bfloat16_supported
|
||||
from .rl import PatchFastRL
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2025.1.8"
|
||||
__version__ = "2025.2.1"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
|
|||
|
|
@ -17,115 +17,8 @@ __all__ = [
|
|||
"PatchKTOTrainer",
|
||||
]
|
||||
|
||||
try:
|
||||
from transformers.utils.notebook import (
|
||||
IntervalStrategy,
|
||||
NotebookTrainingTracker,
|
||||
NotebookProgressCallback,
|
||||
)
|
||||
HAS_NOTEBOOK = True
|
||||
except:
|
||||
HAS_NOTEBOOK = False
|
||||
pass
|
||||
import torch
|
||||
from ._utils import torch_compile_options
|
||||
import inspect
|
||||
import torch.nn as nn
|
||||
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
|
||||
from .rl import PatchFastRL
|
||||
|
||||
def PatchDPOTrainer(): PatchFastRL("DPO")
|
||||
|
||||
DPOTrainer_metrics = [
|
||||
"rewards/chosen",
|
||||
"rewards/rejected",
|
||||
"rewards/accuracies",
|
||||
"rewards/margins",
|
||||
"logps/rejected",
|
||||
"logps/chosen",
|
||||
"logits/rejected",
|
||||
"logits/chosen",
|
||||
]
|
||||
set_DPOTrainer_metrics = frozenset(DPOTrainer_metrics)
|
||||
|
||||
|
||||
def NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs):
|
||||
self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
|
||||
self.training_loss = 0
|
||||
self.last_log = 0
|
||||
column_names = [self.first_column] + ["Training Loss"]
|
||||
if args.eval_strategy != IntervalStrategy.NO:
|
||||
column_names.append("Validation Loss")
|
||||
column_names += [x.replace("/", " / ") for x in DPOTrainer_metrics]
|
||||
self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
|
||||
pass
|
||||
|
||||
|
||||
def NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs):
|
||||
# Only for when there is no evaluation
|
||||
if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
|
||||
values = {"Training Loss": logs["loss"]}
|
||||
for metric in DPOTrainer_metrics:
|
||||
values[metric.replace("/", " / ")] = logs[metric]
|
||||
pass
|
||||
# First column is necessarily Step since we're not in epoch eval strategy
|
||||
values["Step"] = state.global_step
|
||||
self.training_tracker.write_line(values)
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
def NotebookTrainingTracker_write_line(self, values):
|
||||
"""
|
||||
Write the values in the inner table.
|
||||
|
||||
Args:
|
||||
values (`Dict[str, float]`): The values to display.
|
||||
"""
|
||||
if self.inner_table is None:
|
||||
self.inner_table = [list(values.keys()), list(values.values())]
|
||||
else:
|
||||
columns = self.inner_table[0]
|
||||
new_values = {}
|
||||
for key, value in values.items():
|
||||
lowered = key.lower()
|
||||
if lowered in set_DPOTrainer_metrics:
|
||||
new_values[lowered.replace("/", " / ")] = value
|
||||
else:
|
||||
new_values[key] = value
|
||||
pass
|
||||
values = new_values
|
||||
|
||||
self.inner_table[0] = columns
|
||||
if len(self.inner_table) > 1:
|
||||
last_values = self.inner_table[-1]
|
||||
first_column = self.inner_table[0][0]
|
||||
if last_values[0] != values[first_column]:
|
||||
# write new line
|
||||
self.inner_table.append([values[c] if c in values else "No Log" for c in columns])
|
||||
else:
|
||||
# update last line
|
||||
new_values = values
|
||||
for c in columns:
|
||||
if c not in new_values.keys():
|
||||
new_values[c] = last_values[columns.index(c)]
|
||||
self.inner_table[-1] = [new_values[c] for c in columns]
|
||||
else:
|
||||
# Edit for evaluation purposes
|
||||
self.inner_table.append([values[c] if c in values else 0 for c in columns])
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
def PatchDPOTrainer():
|
||||
if HAS_NOTEBOOK:
|
||||
from transformers.trainer import is_in_notebook
|
||||
if is_in_notebook():
|
||||
# Patch DPO notebook printing
|
||||
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
|
||||
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
||||
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
|
||||
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
PatchKTOTrainer = PatchDPOTrainer
|
||||
def PatchKTOTrainer(): PatchFastRL("KTO")
|
||||
|
|
|
|||
|
|
@ -210,7 +210,15 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
|
|||
config = None, # [TODO] Hack to pass in config - need to remove later
|
||||
):
|
||||
super().__init__()
|
||||
if config is not None: return # [TODO] Hack to pass in config - need to remove later
|
||||
if config is not None:
|
||||
# [TODO] Hack to pass in config - need to remove later
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
dim = getattr(config, "head_dim", None)
|
||||
if dim is None: dim = int((config.hidden_size // config.num_attention_heads))
|
||||
device = "cuda"
|
||||
max_position_embeddings = config.max_position_embeddings
|
||||
pass
|
||||
self.dim = dim
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.base = base
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from ._utils import *
|
|||
from ._utils import __version__
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from transformers import __version__ as transformers_version
|
||||
from unsloth_zoo.utils import Version
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
transformers_version = Version(transformers_version)
|
||||
# Transformers moved rotary embeddings out of all attention layers
|
||||
IS_ATTENTION_REFACTOR = transformers_version > Version("4.47.1")
|
||||
|
|
@ -70,7 +70,8 @@ except:
|
|||
from huggingface_hub.utils._token import get_token
|
||||
pass
|
||||
from triton import __version__ as triton_version
|
||||
BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if xformers is not None else None
|
||||
HAS_XFORMERS = xformers is not None
|
||||
BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None
|
||||
|
||||
|
||||
def original_apply_qkv(self, X):
|
||||
|
|
@ -89,6 +90,8 @@ pass
|
|||
from math import sqrt as math_sqrt
|
||||
KV_CACHE_INCREMENT = 256 # KV Cache update size
|
||||
torch_nn_functional_softmax = torch.nn.functional.softmax
|
||||
# SDPA has GQA internally
|
||||
SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__
|
||||
|
||||
# Fix new HF's inference code
|
||||
def _fast_prepare_inputs_for_generation(self, input_ids, **kwargs,):
|
||||
|
|
@ -243,7 +246,7 @@ def LlamaAttention_fast_forward_inference(
|
|||
|
||||
# Grouped query attention
|
||||
_, _, cached_len, _ = Knn.shape
|
||||
if n_groups != 1:
|
||||
if bsz == 1 or not SDPA_HAS_GQA and n_groups != 1:
|
||||
Knn = Knn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
|
||||
Vnn = Vnn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
|
||||
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
|
||||
|
|
@ -262,7 +265,10 @@ def LlamaAttention_fast_forward_inference(
|
|||
A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
|
||||
A = torch_matmul(A, Vnn, out = Qn)
|
||||
else:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False)
|
||||
if SDPA_HAS_GQA:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True)
|
||||
else:
|
||||
A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False)
|
||||
pass
|
||||
A = A.transpose(1, 2)
|
||||
A = A.reshape(bsz, 1, attention_size)
|
||||
|
|
@ -272,15 +278,15 @@ pass
|
|||
|
||||
|
||||
torch_nn_functional_silu = torch.nn.functional.silu
|
||||
def fast_swiglu_inference(self, X):
|
||||
def fast_swiglu_inference(self, X, temp_gate = None, temp_up = None):
|
||||
# gate = self.gate_proj(X)
|
||||
# up = self.up_proj(X)
|
||||
bsz, _, hd = X.shape
|
||||
# mlp_size = self.config.intermediate_size
|
||||
# temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
|
||||
gate = fast_linear_forward(self.gate_proj, X)#, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X)#, out = temp[1])
|
||||
gate = fast_linear_forward(self.gate_proj, X, out = temp_gate)
|
||||
up = fast_linear_forward(self. up_proj, X, out = temp_up)
|
||||
gate = torch_nn_functional_silu(gate, inplace = True)
|
||||
gate *= up
|
||||
|
||||
|
|
@ -289,14 +295,23 @@ def fast_swiglu_inference(self, X):
|
|||
return down
|
||||
pass
|
||||
|
||||
|
||||
def fast_rms_layernorm_inference(self, X):
|
||||
torch_square = torch.square
|
||||
torch_mean = torch.mean
|
||||
def fast_rms_layernorm_inference(self, X, XX = None, XX2 = None, variance = None):
|
||||
old_dtype = X.dtype
|
||||
XX = X.to(torch.float32)
|
||||
variance = XX.square().mean(-1, keepdim = True)
|
||||
if XX is None:
|
||||
XX = X.to(torch.float32)
|
||||
variance = XX.square().mean(-1, keepdim = True)
|
||||
else:
|
||||
XX.copy_(X)
|
||||
torch_mean(torch_square(XX, out = XX2), -1, keepdim = True, out = variance)
|
||||
pass
|
||||
variance += self.variance_epsilon
|
||||
XX *= variance.rsqrt_()
|
||||
X = XX.to(old_dtype) # Must preserve due to residual
|
||||
|
||||
if XX is None: X = XX.to(old_dtype)
|
||||
else: X.copy_(XX)
|
||||
|
||||
X *= self.weight
|
||||
return X
|
||||
pass
|
||||
|
|
@ -403,7 +418,7 @@ def LlamaAttention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if (not HAS_FLASH_ATTENTION and attention_mask is None):
|
||||
if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None):
|
||||
# Xformers memory efficient attention
|
||||
# Also has Flash Attention v2 dispatching
|
||||
Q = Q.transpose(1, 2)
|
||||
|
|
@ -902,15 +917,29 @@ def LlamaModel_fast_forward_inference(
|
|||
attention_mask = None,
|
||||
):
|
||||
input_ids = input_ids[:,:self.max_seq_length]
|
||||
hidden_states = self.model.embed_tokens(input_ids)
|
||||
hidden_states = hidden_states.to(self.config.torch_dtype)
|
||||
bsz, q_len, hd = hidden_states.shape
|
||||
bsz, q_len = input_ids.shape
|
||||
hd = self.config.hidden_size
|
||||
mlp_size = self.config.intermediate_size
|
||||
|
||||
X = self.model.embed_tokens(input_ids)
|
||||
X = X.to(self.config.torch_dtype)
|
||||
bsz, q_len, hd = X.shape
|
||||
assert(q_len == 1)
|
||||
|
||||
# Get saved buffers to reduce memory movement
|
||||
residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
_XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
|
||||
XX, XX2 = _XX[0], _XX[1]
|
||||
variance = torch.empty((bsz, q_len, 1), dtype = torch.float32, device = "cuda:0")
|
||||
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
|
||||
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
|
||||
|
||||
seq_len = past_key_values[0][0].shape[-2]
|
||||
if bsz != 1:
|
||||
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask,
|
||||
(bsz, q_len),
|
||||
hidden_states,
|
||||
X,
|
||||
seq_len,
|
||||
sliding_window = getattr(self.config, "sliding_window", None),
|
||||
)
|
||||
|
|
@ -919,30 +948,54 @@ def LlamaModel_fast_forward_inference(
|
|||
pass
|
||||
|
||||
next_decoder_cache = []
|
||||
|
||||
for idx, decoder_layer in enumerate(self.model.layers):
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(decoder_layer.input_layernorm, hidden_states)
|
||||
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
residual.copy_(X) # residual = X
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.input_layernorm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X, present_key_value = LlamaAttention_fast_forward_inference(
|
||||
decoder_layer.self_attn,
|
||||
hidden_states = hidden_states,
|
||||
hidden_states = X,
|
||||
past_key_value = past_key_values[idx],
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"),
|
||||
)
|
||||
hidden_states += residual
|
||||
X += residual
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm_inference(decoder_layer.post_attention_layernorm, hidden_states)
|
||||
hidden_states = fast_swiglu_inference(decoder_layer.mlp, hidden_states)
|
||||
hidden_states += residual
|
||||
residual.copy_(X) # residual = X
|
||||
X = fast_rms_layernorm_inference(
|
||||
decoder_layer.post_attention_layernorm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
X = fast_swiglu_inference(
|
||||
decoder_layer.mlp,
|
||||
X,
|
||||
temp_gate = temp_gate,
|
||||
temp_up = temp_up,
|
||||
)
|
||||
X += residual
|
||||
|
||||
next_decoder_cache.append(present_key_value)
|
||||
pass
|
||||
hidden_states = fast_rms_layernorm_inference(self.model.norm, hidden_states)
|
||||
X = fast_rms_layernorm_inference(
|
||||
self.model.norm,
|
||||
X,
|
||||
XX = XX,
|
||||
XX2 = XX2,
|
||||
variance = variance,
|
||||
)
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state = hidden_states,
|
||||
last_hidden_state = X,
|
||||
past_key_values = next_decoder_cache,
|
||||
hidden_states = [],
|
||||
attentions = [],
|
||||
|
|
@ -977,7 +1030,7 @@ def CausalLM_fast_forward(fast_forward_inference):
|
|||
attention_mask = attention_mask,
|
||||
)
|
||||
else:
|
||||
causal_mask = xformers.attn_bias.LowerTriangularMask()
|
||||
causal_mask = xformers.attn_bias.LowerTriangularMask() if HAS_XFORMERS else None
|
||||
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
|
|
@ -1159,7 +1212,8 @@ class LlamaRotaryEmbedding(torch.nn.Module):
|
|||
# [TODO] Hack to pass in config - need to remove later
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
dim = int((config.hidden_size // config.num_attention_heads))
|
||||
dim = getattr(config, "head_dim", None)
|
||||
if dim is None: dim = int((config.hidden_size // config.num_attention_heads))
|
||||
device = "cuda"
|
||||
max_position_embeddings = config.max_position_embeddings
|
||||
pass
|
||||
|
|
@ -1580,9 +1634,18 @@ class FastLlamaModel:
|
|||
model_patcher = None,
|
||||
tokenizer_name = None,
|
||||
trust_remote_code = False,
|
||||
|
||||
fast_inference = False, # uses vLLM
|
||||
gpu_memory_utilization = 0.5,
|
||||
float8_kv_cache = False,
|
||||
random_state = 3407,
|
||||
max_lora_rank = 16,
|
||||
disable_log_stats = False,
|
||||
**kwargs,
|
||||
):
|
||||
if trust_remote_code:
|
||||
if fast_inference:
|
||||
raise NotImplementedError("Unsloth: Fast inference does not support `trust_remote_code` yet.")
|
||||
print(
|
||||
"Unsloth: WARNING `trust_remote_code` is True.\n"\
|
||||
"Are you certain you want to do remote code execution?"
|
||||
|
|
@ -1596,9 +1659,9 @@ class FastLlamaModel:
|
|||
|
||||
statistics = \
|
||||
f"==((====))== Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers: {transformers_version}.\n"\
|
||||
f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\
|
||||
f"O^O/ \_/ \\ Torch: {torch.__version__}. CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {torch.version.cuda}. Triton: {triton_version}\n"\
|
||||
f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\
|
||||
f" {chr(92)}{chr(92)} /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\
|
||||
f"O^O/ {chr(92)}_/ {chr(92)} Torch: {torch.__version__}. CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {torch.version.cuda}. Triton: {triton_version}\n"\
|
||||
f"{chr(92)} / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\
|
||||
f' "-____-" Free Apache license: http://github.com/unslothai/unsloth'
|
||||
print(statistics)
|
||||
|
||||
|
|
@ -1626,7 +1689,11 @@ class FastLlamaModel:
|
|||
assert(dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32)
|
||||
|
||||
# RoPE Scaling
|
||||
model_config = AutoConfig.from_pretrained(model_name, token = token)
|
||||
model_config = AutoConfig.from_pretrained(
|
||||
model_name,
|
||||
token = token,
|
||||
attn_implementation = "sdpa",
|
||||
)
|
||||
model_max_seq_length = model_config.max_position_embeddings
|
||||
|
||||
# Check if RoPE Scaling is even allowed
|
||||
|
|
@ -1647,6 +1714,9 @@ class FastLlamaModel:
|
|||
|
||||
rope_scaling = max_seq_length / model_max_seq_length
|
||||
|
||||
if fast_inference:
|
||||
raise NotImplementedError("Unsloth: Fast inference does not yet work with RoPE Scaling.")
|
||||
|
||||
logger.warning_once(
|
||||
f"Unsloth: {model_name} can only handle sequence lengths of at most "\
|
||||
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "\
|
||||
|
|
@ -1688,17 +1758,54 @@ class FastLlamaModel:
|
|||
# Cannot be None, since HF now checks for the config
|
||||
if load_in_4bit: kwargs["quantization_config"] = bnb_config
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
device_map = device_map,
|
||||
torch_dtype = dtype,
|
||||
# quantization_config = bnb_config,
|
||||
token = token,
|
||||
max_position_embeddings = max_position_embeddings,
|
||||
trust_remote_code = trust_remote_code,
|
||||
attn_implementation = "eager",
|
||||
**kwargs,
|
||||
)
|
||||
if not fast_inference:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
device_map = device_map,
|
||||
torch_dtype = dtype,
|
||||
# quantization_config = bnb_config,
|
||||
token = token,
|
||||
max_position_embeddings = max_position_embeddings,
|
||||
trust_remote_code = trust_remote_code,
|
||||
attn_implementation = "eager",
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
from unsloth_zoo.vllm_utils import (
|
||||
load_vllm,
|
||||
get_vllm_state_dict,
|
||||
convert_vllm_to_huggingface,
|
||||
generate_batches,
|
||||
)
|
||||
allowed_args = inspect.getfullargspec(load_vllm).args
|
||||
load_vllm_kwargs = dict(
|
||||
model_name = model_name,
|
||||
config = model_config,
|
||||
gpu_memory_utilization = gpu_memory_utilization,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
float8_kv_cache = float8_kv_cache,
|
||||
enable_lora = True,
|
||||
max_lora_rank = max_lora_rank,
|
||||
disable_log_stats = disable_log_stats,
|
||||
)
|
||||
for allowed_arg in allowed_args:
|
||||
if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs:
|
||||
load_vllm_kwargs[allowed_arg] = kwargs[allowed_arg]
|
||||
pass
|
||||
|
||||
# Load vLLM first
|
||||
llm = load_vllm(**load_vllm_kwargs)
|
||||
|
||||
# Convert to HF format
|
||||
_, quant_state_dict = get_vllm_state_dict(llm, config = model_config)
|
||||
model = convert_vllm_to_huggingface(quant_state_dict, model_config, dtype)
|
||||
model.vllm_engine = llm
|
||||
model.fast_generate = model.vllm_engine.generate
|
||||
|
||||
from functools import partial
|
||||
model.fast_generate_batches = partial(generate_batches, model.vllm_engine)
|
||||
pass
|
||||
# Return old flag
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
|
||||
# We currently only support NVIDIA GPUs - AMD / Intel is a work in progress!
|
||||
|
|
@ -2194,6 +2301,20 @@ class FastLlamaModel:
|
|||
modules_to_save = list(set(modules_to_save))
|
||||
pass
|
||||
|
||||
vllm_engine = None
|
||||
if hasattr(model, "vllm_engine"):
|
||||
# Fast inference!
|
||||
vllm_engine = model.vllm_engine
|
||||
vllm_fast_generate = model.fast_generate
|
||||
vllm_fast_generate_batches = model.fast_generate_batches
|
||||
|
||||
if modules_to_save is not None:
|
||||
raise NotImplementedError("Unsloth: Currently fast inference does not work with training embeddings or lm_head.")
|
||||
|
||||
if bias != "none":
|
||||
raise NotImplementedError("Unsloth: Currently fast inference does not work with using biases for LoRA.")
|
||||
pass
|
||||
|
||||
# Get LoRA
|
||||
arguments = dict(
|
||||
r = r,
|
||||
|
|
@ -2300,6 +2421,19 @@ class FastLlamaModel:
|
|||
torch.cuda.empty_cache()
|
||||
pass
|
||||
|
||||
# Patch for fast inference
|
||||
if vllm_engine is not None:
|
||||
model.vllm_engine = vllm_engine
|
||||
model.fast_generate = vllm_fast_generate
|
||||
model.fast_generate_batches = vllm_fast_generate_batches
|
||||
|
||||
# Also saving and loading LoRA
|
||||
from functools import partial
|
||||
from unsloth_zoo.vllm_utils import save_lora, load_lora
|
||||
model.save_lora = partial(save_lora, model)
|
||||
model.load_lora = partial(load_lora, model)
|
||||
pass
|
||||
|
||||
return model
|
||||
pass
|
||||
|
||||
|
|
@ -2509,18 +2643,24 @@ class FastLlamaModel:
|
|||
# return
|
||||
# pass
|
||||
|
||||
internal_model = model
|
||||
internal_model.gradient_checkpointing = False
|
||||
internal_model.training = False
|
||||
|
||||
while hasattr(internal_model, "model"):
|
||||
internal_model = internal_model.model
|
||||
internal_model.gradient_checkpointing = False
|
||||
internal_model.training = False
|
||||
pass
|
||||
if hasattr(internal_model, "training"):
|
||||
internal_model.training = False
|
||||
m = model
|
||||
while hasattr(m, "model"):
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = False
|
||||
if hasattr(m, "training"):
|
||||
m.training = False
|
||||
# Pad tokenizer to the left
|
||||
if hasattr(m, "_saved_temp_tokenizer"):
|
||||
m._saved_temp_tokenizer.padding_side = "left"
|
||||
m = m.model
|
||||
pass
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = False
|
||||
if hasattr(m, "training"):
|
||||
m.training = False
|
||||
# Pad tokenizer to the left
|
||||
if hasattr(m, "_saved_temp_tokenizer"):
|
||||
m._saved_temp_tokenizer.padding_side = "left"
|
||||
|
||||
# Also check if lm_head / embeddings are trained
|
||||
internal_model = model
|
||||
|
|
@ -2529,30 +2669,13 @@ class FastLlamaModel:
|
|||
pass
|
||||
lm_head = internal_model.lm_head.weight
|
||||
device_type = lm_head.device.type
|
||||
dtype = model.config.torch_dtype
|
||||
|
||||
if type(dtype) is str:
|
||||
if dtype == "float16": dtype = torch.float16
|
||||
elif dtype == "bfloat16": dtype = torch.bfloat16
|
||||
pass
|
||||
dtype = _get_dtype(model.config.torch_dtype)
|
||||
|
||||
# Wrap model.generate
|
||||
if model.generate.__name__ != "_fast_generate":
|
||||
model._unwrapped_old_generate = model.generate
|
||||
model.generate = _wrap_fast_inference(model.generate, device_type, dtype, model)
|
||||
pass
|
||||
|
||||
# Patch tokenizer to pad to the left
|
||||
internal_model = model
|
||||
while hasattr(internal_model, "model"):
|
||||
if hasattr(internal_model, "_saved_temp_tokenizer"):
|
||||
internal_model._saved_temp_tokenizer.padding_side = "left"
|
||||
pass
|
||||
internal_model = internal_model.model
|
||||
pass
|
||||
if hasattr(internal_model, "_saved_temp_tokenizer"):
|
||||
internal_model._saved_temp_tokenizer.padding_side = "left"
|
||||
pass
|
||||
|
||||
# Also disable training for embeddings for NEFTune
|
||||
if hasattr(model, "get_input_embeddings"):
|
||||
|
|
@ -2570,9 +2693,6 @@ class FastLlamaModel:
|
|||
|
||||
@staticmethod
|
||||
def for_training(model, use_gradient_checkpointing = True):
|
||||
internal_model = model
|
||||
internal_model.gradient_checkpointing = use_gradient_checkpointing
|
||||
internal_model.training = True
|
||||
|
||||
# Delete all fast inference loras
|
||||
for param in model.parameters():
|
||||
|
|
@ -2580,14 +2700,24 @@ class FastLlamaModel:
|
|||
del param._fast_lora
|
||||
pass
|
||||
|
||||
while hasattr(internal_model, "model"):
|
||||
internal_model = internal_model.model
|
||||
internal_model.gradient_checkpointing = use_gradient_checkpointing
|
||||
internal_model.training = True
|
||||
pass
|
||||
if hasattr(internal_model, "training"):
|
||||
internal_model.training = True
|
||||
m = model
|
||||
while hasattr(m, "model"):
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = use_gradient_checkpointing
|
||||
if hasattr(m, "training"):
|
||||
m.training = True
|
||||
# Pad tokenizer to the right
|
||||
if hasattr(m, "_saved_temp_tokenizer"):
|
||||
m._saved_temp_tokenizer.padding_side = "right"
|
||||
m = m.model
|
||||
pass
|
||||
if hasattr(m, "gradient_checkpointing"):
|
||||
m.gradient_checkpointing = use_gradient_checkpointing
|
||||
if hasattr(m, "training"):
|
||||
m.training = True
|
||||
# Pad tokenizer to the right
|
||||
if hasattr(m, "_saved_temp_tokenizer"):
|
||||
m._saved_temp_tokenizer.padding_side = "right"
|
||||
|
||||
# Also revert model.generate
|
||||
if hasattr(model, "_unwrapped_old_generate"):
|
||||
|
|
@ -2595,18 +2725,6 @@ class FastLlamaModel:
|
|||
del model._unwrapped_old_generate
|
||||
pass
|
||||
|
||||
# Patch tokenizer to pad to the right
|
||||
internal_model = model
|
||||
while hasattr(internal_model, "model"):
|
||||
if hasattr(internal_model, "_saved_temp_tokenizer"):
|
||||
internal_model._saved_temp_tokenizer.padding_side = "right"
|
||||
pass
|
||||
internal_model = internal_model.model
|
||||
pass
|
||||
if hasattr(internal_model, "_saved_temp_tokenizer"):
|
||||
internal_model._saved_temp_tokenizer.padding_side = "right"
|
||||
pass
|
||||
|
||||
# Also re-enable training for embeddings for NEFTune
|
||||
if hasattr(model, "get_input_embeddings"):
|
||||
embeddings = model.get_input_embeddings()
|
||||
|
|
|
|||
|
|
@ -30,11 +30,11 @@ except:
|
|||
from huggingface_hub.utils._token import get_token
|
||||
pass
|
||||
from huggingface_hub import HfFileSystem
|
||||
import importlib.util
|
||||
|
||||
# [TODO] Move USE_MODELSCOPE to utils
|
||||
USE_MODELSCOPE = os.environ.get("UNSLOTH_USE_MODELSCOPE", "0") == "1"
|
||||
if USE_MODELSCOPE:
|
||||
import importlib
|
||||
if importlib.util.find_spec("modelscope") is None:
|
||||
raise ImportError(f'You are using the modelscope hub, please install modelscope by `pip install modelscope -U`')
|
||||
pass
|
||||
|
|
@ -73,9 +73,25 @@ class FastLanguageModel(FastLlamaModel):
|
|||
resize_model_vocab = None,
|
||||
revision = None,
|
||||
use_exact_model_name = False,
|
||||
|
||||
fast_inference = False, # uses vLLM
|
||||
gpu_memory_utilization = 0.5,
|
||||
float8_kv_cache = False,
|
||||
random_state = 3407,
|
||||
max_lora_rank = 64,
|
||||
disable_log_stats = True,
|
||||
*args, **kwargs,
|
||||
):
|
||||
if token is None: token = get_token()
|
||||
|
||||
if fast_inference:
|
||||
if importlib.util.find_spec("vllm") is None:
|
||||
raise ImportError(
|
||||
"Unsloth: Please install vLLM before enabling `fast_inference`!\n"\
|
||||
"You can do this in a terminal via `pip install vllm`"
|
||||
)
|
||||
pass
|
||||
pass
|
||||
|
||||
old_model_name = model_name
|
||||
if not use_exact_model_name:
|
||||
|
|
@ -255,6 +271,24 @@ class FastLanguageModel(FastLlamaModel):
|
|||
tokenizer_name = None
|
||||
pass
|
||||
|
||||
if fast_inference:
|
||||
from unsloth_zoo.vllm_utils import (
|
||||
patch_vllm,
|
||||
vllm_dynamic_quant_supported,
|
||||
)
|
||||
patch_vllm()
|
||||
if model_name.endswith("unsloth-bnb-4bit"):
|
||||
if not vllm_dynamic_quant_supported(model_name, model_config):
|
||||
# Instead use -bnb-4bit variant
|
||||
print(
|
||||
f"Unsloth: Switching from Unsloth dynamic quant to normal quant since\n"\
|
||||
f"we do not yet support fast inference for {model_name}"
|
||||
)
|
||||
model_name = model_name[:-len("unsloth-bnb-4bit")] + "bnb-4bit"
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
model, tokenizer = dispatch_model.from_pretrained(
|
||||
model_name = model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
|
|
@ -268,6 +302,13 @@ class FastLanguageModel(FastLlamaModel):
|
|||
tokenizer_name = tokenizer_name,
|
||||
trust_remote_code = trust_remote_code,
|
||||
revision = revision if not is_peft else None,
|
||||
|
||||
fast_inference = fast_inference,
|
||||
gpu_memory_utilization = gpu_memory_utilization,
|
||||
float8_kv_cache = float8_kv_cache,
|
||||
random_state = random_state,
|
||||
max_lora_rank = max_lora_rank,
|
||||
disable_log_stats = disable_log_stats,
|
||||
*args, **kwargs,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -304,25 +304,30 @@ __INT_TO_FLOAT_MAPPER = \
|
|||
"unsloth/Mistral-Small-Instruct-2409",
|
||||
"mistralai/Mistral-Small-Instruct-2409",
|
||||
),
|
||||
"unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-0.5B-Instruct-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-0.5B-Instruct",
|
||||
"Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"unsloth/Qwen2.5-0.5B-Instruct-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-1.5B-Instruct-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-1.5B-Instruct",
|
||||
"Qwen/Qwen2.5-1.5B-Instruct",
|
||||
"unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-3B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-3B-Instruct-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-3B-Instruct",
|
||||
"Qwen/Qwen2.5-3B-Instruct",
|
||||
"unsloth/Qwen2.5-3B-Instruct-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-7B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-7B-Instruct-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-7B-Instruct",
|
||||
"Qwen/Qwen2.5-7B-Instruct",
|
||||
"unsloth/Qwen2.5-7B-Instruct-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-14B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-14B-Instruct-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-14B-Instruct",
|
||||
"Qwen/Qwen2.5-14B-Instruct",
|
||||
"unsloth/Qwen2.5-14B-Instruct-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-32B-Instruct-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-32B-Instruct",
|
||||
|
|
@ -332,25 +337,30 @@ __INT_TO_FLOAT_MAPPER = \
|
|||
"unsloth/Qwen2.5-72B-Instruct",
|
||||
"Qwen/Qwen2.5-72B-Instruct",
|
||||
),
|
||||
"unsloth/Qwen2.5-0.5B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-0.5B-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-0.5B",
|
||||
"Qwen/Qwen2.5-0.5B",
|
||||
"unsloth/Qwen2.5-0.5B-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-1.5B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-1.5B-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-1.5B",
|
||||
"Qwen/Qwen2.5-1.5B",
|
||||
"unsloth/Qwen2.5-1.5B-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-3B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-3B-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-3B",
|
||||
"Qwen/Qwen2.5-3B",
|
||||
"unsloth/Qwen2.5-3B-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-7B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-7B-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-7B",
|
||||
"Qwen/Qwen2.5-7B",
|
||||
"unsloth/Qwen2.5-7B-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-14B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-14B-unsloth-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-14B",
|
||||
"Qwen/Qwen2.5-14B",
|
||||
"unsloth/Qwen2.5-14B-bnb-4bit",
|
||||
),
|
||||
"unsloth/Qwen2.5-32B-bnb-4bit" : (
|
||||
"unsloth/Qwen2.5-32B",
|
||||
|
|
@ -555,12 +565,12 @@ __INT_TO_FLOAT_MAPPER = \
|
|||
"deepseek-ai/DeepSeek-R1-Distill-Llama-70B",
|
||||
),
|
||||
"unsloth/Mistral-Small-24B-Base-2501-unsloth-bnb-4bit" : (
|
||||
"unsloth/Mistral-Small-24B-Base",
|
||||
"unsloth/Mistral-Small-24B-Base-2501",
|
||||
"mistralai/Mistral-Small-24B-Base-2501",
|
||||
"unsloth/Mistral-Small-24B-Base-2501-bnb-4bit",
|
||||
),
|
||||
"unsloth/Mistral-Small-24B-Instruct-2501-unsloth-bnb-4bit" : (
|
||||
"unsloth/Mistral-Small-24B-Instruct",
|
||||
"unsloth/Mistral-Small-24B-Instruct-2501",
|
||||
"mistralai/Mistral-Small-24B-Instruct-2501",
|
||||
"unsloth/Mistral-Small-24B-Instruct-2501-bnb-4bit",
|
||||
),
|
||||
|
|
|
|||
423
unsloth/models/rl.py
Normal file
423
unsloth/models/rl.py
Normal file
|
|
@ -0,0 +1,423 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__all__ = [
|
||||
"PatchFastRL",
|
||||
]
|
||||
|
||||
METRICS_MOVE_TO_END = [
|
||||
"nll",
|
||||
"aux",
|
||||
"beta",
|
||||
"alpha",
|
||||
]
|
||||
import torch
|
||||
try:
|
||||
from transformers.utils.notebook import (
|
||||
IntervalStrategy,
|
||||
NotebookTrainingTracker,
|
||||
NotebookProgressCallback,
|
||||
)
|
||||
HAS_NOTEBOOK = True
|
||||
except:
|
||||
HAS_NOTEBOOK = False
|
||||
pass
|
||||
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import functools
|
||||
from unsloth_zoo.compiler import create_new_function
|
||||
|
||||
|
||||
def PatchRL(FastLanguageModel):
|
||||
|
||||
from trl.models.utils import unwrap_model_for_generation
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def unsloth_unwrap_model_for_generation(model, accelerator):
|
||||
with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
|
||||
# Put the model in inference mode.
|
||||
FastLanguageModel.for_inference(unwrapped_model)
|
||||
|
||||
# We must use .clone for Unsloth since we force inference_mode
|
||||
# Rather we should have used no_grad
|
||||
original_generate = unwrapped_model.generate
|
||||
def generate_with_clone(*args, **kwargs):
|
||||
out = original_generate(*args, **kwargs)
|
||||
if isinstance(out, torch.Tensor):
|
||||
return out.clone()
|
||||
return out
|
||||
pass
|
||||
unwrapped_model.generate = generate_with_clone
|
||||
|
||||
try:
|
||||
yield unwrapped_model
|
||||
finally:
|
||||
# Restore generate and return
|
||||
unwrapped_model.generate = original_generate
|
||||
FastLanguageModel.for_training(model)
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
import trl.trainer
|
||||
trainers = dir(trl.trainer)
|
||||
trainers = [x for x in trainers if x.endswith("_trainer")]
|
||||
unwrap = "unwrap_model_for_generation"
|
||||
for trainer in trainers:
|
||||
if hasattr(eval(f"trl.trainer.{trainer}"), unwrap):
|
||||
exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}")
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
def NotebookProgressCallback_on_train_begin(Trainer_metrics):
|
||||
def _NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs):
|
||||
self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
|
||||
self.training_loss = 0
|
||||
self.last_log = 0
|
||||
column_names = [self.first_column] + ["Training Loss"]
|
||||
if args.eval_strategy != IntervalStrategy.NO:
|
||||
column_names.append("Validation Loss")
|
||||
column_names += [x.replace("/", " / ") for x in Trainer_metrics]
|
||||
self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
|
||||
pass
|
||||
return _NotebookProgressCallback_on_train_begin
|
||||
pass
|
||||
|
||||
|
||||
def NotebookProgressCallback_on_log(Trainer_metrics):
|
||||
def _NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs):
|
||||
# Only for when there is no evaluation
|
||||
if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
|
||||
values = {"Training Loss": logs["loss"]}
|
||||
for metric in Trainer_metrics:
|
||||
# Sometimes metric is not inside logs
|
||||
try: values[metric.replace("/", " / ")] = logs[metric]
|
||||
except: pass
|
||||
pass
|
||||
# First column is necessarily Step since we're not in epoch eval strategy
|
||||
values["Step"] = state.global_step
|
||||
self.training_tracker.write_line(values)
|
||||
pass
|
||||
pass
|
||||
return _NotebookProgressCallback_on_log
|
||||
pass
|
||||
|
||||
|
||||
def NotebookTrainingTracker_write_line(Trainer_metrics):
|
||||
set_Trainer_metrics = set(Trainer_metrics)
|
||||
def _NotebookTrainingTracker_write_line(self, values):
|
||||
"""
|
||||
Write the values in the inner table.
|
||||
|
||||
Args:
|
||||
values (`Dict[str, float]`): The values to display.
|
||||
"""
|
||||
if self.inner_table is None:
|
||||
self.inner_table = [list(values.keys()), list(values.values())]
|
||||
else:
|
||||
columns = self.inner_table[0]
|
||||
new_values = {}
|
||||
for key, value in values.items():
|
||||
lowered = key.lower()
|
||||
if lowered in set_Trainer_metrics:
|
||||
new_values[lowered.replace("/", " / ")] = value
|
||||
else:
|
||||
new_values[key] = value
|
||||
pass
|
||||
values = new_values
|
||||
|
||||
self.inner_table[0] = columns
|
||||
if len(self.inner_table) > 1:
|
||||
last_values = self.inner_table[-1]
|
||||
first_column = self.inner_table[0][0]
|
||||
if last_values[0] != values[first_column]:
|
||||
# write new line
|
||||
self.inner_table.append([values[c] if c in values else "No Log" for c in columns])
|
||||
else:
|
||||
# update last line
|
||||
new_values = values
|
||||
for c in columns:
|
||||
if c not in new_values.keys():
|
||||
new_values[c] = last_values[columns.index(c)]
|
||||
self.inner_table[-1] = [new_values[c] for c in columns]
|
||||
else:
|
||||
# Edit for evaluation purposes
|
||||
self.inner_table.append([values[c] if c in values else 0 for c in columns])
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
return _NotebookTrainingTracker_write_line
|
||||
pass
|
||||
|
||||
|
||||
def _PatchRLStatistics(metrics, algorithm):
|
||||
if HAS_NOTEBOOK:
|
||||
if len(metrics) == 0:
|
||||
raise RuntimeError(f"Unsloth: RL statistics for {algorithm} failed with no metrics seen?")
|
||||
from transformers.trainer import is_in_notebook
|
||||
if is_in_notebook():
|
||||
# Patch DPO notebook printing
|
||||
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line(metrics)
|
||||
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
||||
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin(metrics)
|
||||
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log(metrics)
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_trl_metrics():
|
||||
# Gets metrics so we can output them in notebooks
|
||||
|
||||
import trl.trainer
|
||||
trainers = dir(trl.trainer)
|
||||
trainers = [x for x in trainers if x.endswith("_trainer")]
|
||||
filepath = inspect.getfile(trl.trainer)
|
||||
filepath = os.path.split(filepath)[0]
|
||||
|
||||
all_metrics = dict()
|
||||
for trainer in trainers:
|
||||
filename = os.path.join(filepath, f"{trainer}.py")
|
||||
if not os.path.exists(filename): continue
|
||||
with open(filename, "r") as file: file = file.read()
|
||||
|
||||
# Get metrics['kl'] or stats['kl']
|
||||
metrics = re.findall(r"metrics\[[\"\']([^\"\']{1,})[\"\']\]", file)
|
||||
stats = re.findall(r"stats\[[\"\']([^\"\']{1,})[\"\']\]", file)
|
||||
metrics = metrics + stats
|
||||
|
||||
# Get optional f-strings
|
||||
metrics_f = re.findall(r"metrics\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file)
|
||||
stats_f = re.findall(r"stats\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file)
|
||||
metrics_f = metrics_f + stats_f
|
||||
# Filter out prefixes if seen
|
||||
# metrics[f"{prefix}rewards/chosen"]
|
||||
left_prefix = 'prefix = "eval_" if train_eval == "eval" else ""' in file
|
||||
if left_prefix: metrics += metrics_f
|
||||
|
||||
# Move all eval_ things to the end and reward to the front
|
||||
beginning = []
|
||||
middle = []
|
||||
end = []
|
||||
for x in metrics:
|
||||
lowered = x.lower()
|
||||
if "reward" in lowered:
|
||||
beginning.append(x)
|
||||
elif x.lower().startswith("eval"):
|
||||
end.append(x)
|
||||
else:
|
||||
# Check if we want to move to the end
|
||||
moved = False
|
||||
for move_end in METRICS_MOVE_TO_END:
|
||||
if move_end in lowered:
|
||||
end.append(x)
|
||||
moved = True
|
||||
break
|
||||
if not moved:
|
||||
middle.append(x)
|
||||
pass
|
||||
pass
|
||||
metrics = beginning + middle + end
|
||||
|
||||
all_metrics[trainer[:trainer.find("_")].upper()] = metrics
|
||||
pass
|
||||
return all_metrics
|
||||
pass
|
||||
|
||||
|
||||
def PatchRLStatistics(algorithm = "GRPO"):
|
||||
# Get notebook statistics columns to show up
|
||||
algorithm = algorithm.upper()
|
||||
all_metrics = get_trl_metrics()
|
||||
if algorithm not in all_metrics:
|
||||
print(
|
||||
f"Unsloth for {algorithm.upper()} is not yet implemented! Just ignore this function.\n"\
|
||||
f"We support: `{list(all_metrics.keys())}`"
|
||||
)
|
||||
pass
|
||||
_PatchRLStatistics(all_metrics[algorithm], algorithm)
|
||||
pass
|
||||
|
||||
|
||||
def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
||||
# Patch for vLLM and Unsloth PEFT
|
||||
import trl
|
||||
import trl.trainer
|
||||
|
||||
trainer = eval(f"trl.trainer.{trainer_file}")
|
||||
name = [x for x in dir(trainer) if x.endswith("Trainer") and x != "Trainer" and trainer_file.split("_")[0] in x.lower()]
|
||||
assert(len(name) == 1)
|
||||
RLTrainer_name = name[0]
|
||||
RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}")
|
||||
|
||||
try:
|
||||
__init__ = inspect.getsource(RLTrainer.__init__)
|
||||
except:
|
||||
# Already patched most likely!
|
||||
return
|
||||
old__init__ = __init__
|
||||
all_imports = dir(trainer)
|
||||
assert("Union" in all_imports)
|
||||
imports = [x for x in all_imports if not x.startswith("_")]
|
||||
imports += ["Trainer"]
|
||||
|
||||
spaces = __init__.find("def")
|
||||
__init__ = __init__.split("\n")
|
||||
__init__ = "\n".join(x[spaces:] for x in __init__)
|
||||
|
||||
# Replace vLLM sections since we already have it done!
|
||||
vllm_part = re.findall(
|
||||
r"(\n[\s]{4}"\
|
||||
r"if (self|args)\.use_vllm\:.+?"\
|
||||
r"\n[\s]{4,}"\
|
||||
"else:\n)",
|
||||
__init__,
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
if (len(vllm_part) != 1): return
|
||||
|
||||
vllm_part, args = vllm_part[0][0], vllm_part[0][1]
|
||||
# Strip all comments
|
||||
new_vllm_part = re.sub(r"\#[^\n]{1,}\n", "", vllm_part)
|
||||
|
||||
# Get SamplingParams
|
||||
sampling_params = re.findall(
|
||||
r"\n[\s]{4,}(self\.[^\s]{1,}[\s]{0,}\=[\s]{0,}"\
|
||||
r"SamplingParams\(.+?\))",
|
||||
new_vllm_part,
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
if len(sampling_params) != 1: return
|
||||
|
||||
sampling_params = sampling_params[0]
|
||||
# Replace with our vLLM engine
|
||||
sampling_params = \
|
||||
" "*8 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \
|
||||
sampling_params # Add spaces
|
||||
new_vllm_part = f"\n if {args}.use_vllm:\n{sampling_params}\n else:\n"
|
||||
__init__ = __init__.replace(vllm_part, new_vllm_part)
|
||||
|
||||
# Remove peft_config
|
||||
__init__ = __init__.replace("elif peft_config is None:", "elif False:")
|
||||
__init__ = __init__.replace("elif peft_config is not None:", "elif False:")
|
||||
__init__ = __init__.replace("if peft_config is None:", "if False:")
|
||||
__init__ = __init__.replace("if peft_config is not None:", "if False:")
|
||||
__init__ = __init__.replace("get_peft_model(model, peft_config)", "model")
|
||||
|
||||
# Add spaces back into __init__
|
||||
__init__ = __init__.split("\n")
|
||||
__init__ = "\n".join(' '*spaces + x for x in __init__)
|
||||
|
||||
# Search for vLLM calling in all child functions
|
||||
functions = dir(RLTrainer)
|
||||
RLTrainer_source = inspect.getsource(RLTrainer)
|
||||
functions = [x for x in functions if f"def {x}" in RLTrainer_source]
|
||||
|
||||
changed = {"__init__" : (old__init__, __init__,)}
|
||||
for function in functions:
|
||||
if not hasattr(RLTrainer, function): continue
|
||||
fx = getattr(RLTrainer, function)
|
||||
try:
|
||||
source = inspect.getsource(fx)
|
||||
except:
|
||||
continue
|
||||
original_source = source
|
||||
|
||||
# llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
|
||||
source = re.sub(
|
||||
r"(\n[\s]{4,}).+?model_executor\.driver_worker.+?\n",
|
||||
r"\n\1pass\n",
|
||||
source,
|
||||
)
|
||||
|
||||
# llm_model.load_weights(model.state_dict().items())
|
||||
source = re.sub(
|
||||
r"(\n[\s]{4,}).+?load_weights\(.+?\n",
|
||||
r"\n\1pass\n",
|
||||
source,
|
||||
)
|
||||
|
||||
# .state_dict()
|
||||
source = re.sub(
|
||||
r"\.state_dict\(\)",
|
||||
r"",
|
||||
source,
|
||||
)
|
||||
|
||||
# Replace self.llm.generate and self.llm.chat
|
||||
lora_name = trainer_file + "_lora_model"
|
||||
source = re.sub(
|
||||
r"(self\.llm\.(?:generate|chat)\([^\)]{1,})\)",
|
||||
r"\1, lora_request = model.load_lora('" + lora_name + r"', load_tensors = True))",
|
||||
source
|
||||
)
|
||||
|
||||
# Skip if no changes done
|
||||
if source == original_source: continue
|
||||
|
||||
# Find all imports
|
||||
imports += [x for x in all_imports if not x.startswith("_") and x in source]
|
||||
|
||||
changed[function] = (original_source, source,)
|
||||
pass
|
||||
|
||||
# Import all functions
|
||||
imports = list(set(imports))
|
||||
|
||||
# Patch all functions
|
||||
for function in changed:
|
||||
old, new = changed[function]
|
||||
RLTrainer_source = RLTrainer_source.replace(old, new)
|
||||
pass
|
||||
RLTrainer_source = RLTrainer_source.replace(
|
||||
f"class {RLTrainer_name}", f"class Unsloth{RLTrainer_name}", 1
|
||||
)
|
||||
|
||||
# Create new class in compiled cache and import it
|
||||
module = create_new_function(
|
||||
RLTrainer_name,
|
||||
RLTrainer_source,
|
||||
f"trl.trainer.{trainer_file}",
|
||||
imports,
|
||||
)
|
||||
|
||||
# Patch over modules
|
||||
exec(f"trl.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals())
|
||||
exec(f"trl.trainer.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals())
|
||||
exec(f"trl.trainer.{trainer_file}.{RLTrainer_name} = module.Unsloth{RLTrainer_name}", locals(), globals())
|
||||
return module
|
||||
pass
|
||||
|
||||
|
||||
def patch_trl_rl_trainers():
|
||||
# Patch all TRL modules if they have vLLM or PEFT
|
||||
import trl.trainer
|
||||
all_trainers = dir(trl.trainer)
|
||||
all_trainers = [x for x in all_trainers if x.islower() and x.endswith("_trainer")]
|
||||
for trainer in all_trainers:
|
||||
_patch_trl_rl_trainers(trainer)
|
||||
return
|
||||
pass
|
||||
|
||||
|
||||
def PatchFastRL(algorithm = "GRPO", FastLanguageModel = None):
|
||||
if FastLanguageModel is not None: PatchRL(FastLanguageModel)
|
||||
patch_trl_rl_trainers()
|
||||
PatchRLStatistics(algorithm)
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue