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

Merge branch 'main' into nightly

Browse source
This commit is contained in:
Daniel Han 2025-07-10 01:50:14 -07:00
commit 13a32054b7
11 changed files with 165 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -2,6 +2,7 @@
__pycache__/
*.py[cod]
*.class
unsloth_compiled_cache/
# C extensions
*.so

25
unsloth/__init__.py
View file

@ -50,20 +50,6 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
# "pinned_use_cuda_host_register:True,"\
# "pinned_num_register_threads:8"
# Hugging Face Hub faster downloads
if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
pass
# Disable XET Cache for now
os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
os.environ["HF_XET_CHUNK_CACHE_SIZE_BYTES"] = "0"
os.environ["HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY"] = "0"
os.environ["HF_XET_NUM_CONCURRENT_RANGE_GETS"] = "64"
# More verbose HF Hub info
if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
os.environ["HF_HUB_VERBOSITY"] = "info"
# Log Unsloth is being used
os.environ["UNSLOTH_IS_PRESENT"] = "1"
@ -225,12 +211,11 @@ elif DEVICE_TYPE == "xpu":
# Check for unsloth_zoo
try:
unsloth_zoo_version = importlib_version("unsloth_zoo")
if Version(unsloth_zoo_version) < Version("2025.4.1"):
pass
# print(
# "Unsloth: Updating Unsloth-Zoo utilies to the latest version.\n"\
# "To disable this, set `os.environ['UNSLOTH_DISABLE_AUTO_UPDATES'] = '1'`"
# )
if Version(unsloth_zoo_version) < Version("2025.7.1"):
print(
"Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n"\
"Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`"
)
# if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0":
# try:
# os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo")

2
unsloth/kernels/utils.py
View file

@ -78,7 +78,7 @@ HAS_CUDA_STREAM = False
# INTEL GPU specific logic
if DEVICE_TYPE == "xpu":
# TODO: Changed here after adding XPU BNB support
HAS_XPU_STREAM = False
HAS_XPU_STREAM = True
def get_ptr(x: Optional[torch.Tensor]):
raise RuntimeError("XPU BNB support is not implemented yet. This function should not be called.")
else:

6
unsloth/models/__init__.py
View file

@ -19,7 +19,11 @@ from .qwen2 import FastQwen2Model
from .qwen3 import FastQwen3Model
from .qwen3_moe import FastQwen3MoeModel
from .granite import FastGraniteModel
from .falcon_h1 import FastFalconH1Model
try:
from .falcon_h1 import FastFalconH1Model
except:
# transformers_version < 4.53.0 does not have falcon_h1 so silenty skip it for now
pass
from .dpo import PatchDPOTrainer, PatchKTOTrainer
from ._utils import is_bfloat16_supported, is_vLLM_available, __version__
from .rl import PatchFastRL, vLLMSamplingParams

10
unsloth/models/_utils.py
View file

@ -142,6 +142,12 @@ warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "
import logging
logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL+1)
def get_device_num():
if DEVICE_TYPE == "xpu":
return torch.xpu.device_count()
else:
return torch.cuda.device_count()
# Ignore logging messages
class HideLoggingMessage(logging.Filter):
__slots__ = "text",
@ -740,7 +746,7 @@ def get_statistics():
pass
pass
try:
devices = torch.cuda.device_count()
devices = get_device_num()
_get_statistics(f"{devices if devices <= 8 else 9}")
except:
pass
@ -767,7 +773,7 @@ BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace(
)
exec(BitsAndBytesConfig__init__, globals())
if torch.cuda.device_count() == 1:
if get_device_num() == 1:
from accelerate.utils.dataclasses import DistributedType
def _prepare_backend(self, *args, **kwargs): return None, DistributedType.NO
import accelerate.state

2
unsloth/models/cohere.py
View file

@ -128,7 +128,7 @@ def CohereAttention_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)

2
unsloth/models/granite.py
View file

@ -115,7 +115,7 @@ def GraniteAttention_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
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)

73
unsloth/models/llama.py
View file

@ -85,6 +85,11 @@ from triton import __version__ as triton_version
HAS_XFORMERS = xformers is not None
BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None
def clean_gpu_cache():
if DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()
else:
torch.cuda.empty_cache()
def original_apply_qkv(self, X):
Q = self.q_proj(X)
@ -319,6 +324,13 @@ def LlamaAttention_fast_forward_inference(
# Knn, Vnn = Knn, Vnn
# pass
# when qlen==vlen and attn_mask is None, we should use causal attention
Q_len = Qn.shape[-2]
K_len = Knn.shape[-2]
if attention_mask is None and Q_len == K_len:
is_causal = True
else:
is_causal = False
# Attention
if bsz == 1:
Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963
@ -519,11 +531,18 @@ def LlamaAttention_fast_forward(
V = V.transpose(1, 2)
A = flash_attn_func(Q, K, V, causal = True)
else:
# when qlen==vlen and attn_mask is None, we should use causal attention
Q_len = Q.shape[-2]
K_len = K.shape[-2]
if attention_mask is None and Q_len == K_len:
is_causal = True
else:
is_causal = False
# Grouped query attention
if SDPA_HAS_GQA:
# Needs (batch_size, n_heads, seq_len, head_dim)
# is_casual and attention_mask must not be both set!
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False, enable_gqa = n_groups != 1)
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = n_groups != 1)
# Go back to (batch_size, seq_len, n_heads, head_dim)
A = A.transpose(1, 2)#.contiguous()
else:
@ -538,7 +557,7 @@ def LlamaAttention_fast_forward(
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
# Needs (batch_size, n_heads, seq_len, head_dim)
# is_casual and attention_mask must not be both set!
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal)
# Go back to (batch_size, seq_len, n_heads, head_dim)
A = A.transpose(1, 2).contiguous()
pass
@ -1275,9 +1294,8 @@ def PeftModel_fast_forward(
logits_to_keep = 0,
**kwargs,
):
is_classification = "Classification" in str(type( self.base_model.model))
is_classification = "Classification" in str(type(self.base_model.model))
if is_classification:
#causal_mask = causal_mask,
return self.base_model(
input_ids = input_ids,
attention_mask = attention_mask,
@ -1287,7 +1305,7 @@ def PeftModel_fast_forward(
output_hidden_states = output_hidden_states,
return_dict = return_dict,
**kwargs,
)
)
else:
return self.base_model(
input_ids = input_ids,
@ -1752,10 +1770,11 @@ class FastLlamaModel:
if not is_vLLM_available():
print("Unsloth: vLLM is not installed! Will use Unsloth inference!")
fast_inference = False
major_version, minor_version = torch.cuda.get_device_capability()
if major_version < 7:
print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!")
fast_inference = False
if DEVICE_TYPE == "cuda":
major_version, minor_version = torch.cuda.get_device_capability()
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
@ -1779,8 +1798,8 @@ class FastLlamaModel:
num_gpus = torch.xpu.device_count()
gpu_stats_snippet = f"Intel Toolkit: {gpu_version}."
# TODO: After adding vLLM support for XPU, changed this
vllm_version = ""
try: vllm_version = f" vLLM: {importlib_version('vllm')}."
except: vllm_version = ""
else:
raise ValueError(f"Unsloth: Unsupported device type: {DEVICE_TYPE}")
@ -2020,7 +2039,10 @@ class FastLlamaModel:
import gc
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()"""
if DEVICE_TYPE == "xpu":
torch.xpu.empty_cache()
else:
torch.cuda.empty_cache()"""
debug_info = debug_info.split('\n')
debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]])
@ -2508,7 +2530,7 @@ class FastLlamaModel:
# Remove old items to save VRAM
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()
clean_gpu_cache()
pass
if train_lm_head:
@ -2519,7 +2541,7 @@ class FastLlamaModel:
# Remove old items to save VRAM
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()
clean_gpu_cache()
pass
pass
@ -2580,7 +2602,7 @@ class FastLlamaModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()
clean_gpu_cache()
pass
# Patch for fast inference
@ -2695,10 +2717,21 @@ class FastLlamaModel:
if lora_dropout == 0 and bias == "none":
for idx, layer in enumerate(model.model.model.layers):
# Determine MLP module name (falcon_h1 has feed_forward, llama style has mlp)
if hasattr(layer, "mlp"):
mlp_module_name = "mlp"
elif hasattr(layer, "feed_forward"):
mlp_module_name = "feed_forward"
else:
logger.warning_once(f"Unsloth: No MLP module found in layer {idx} so skipping peft mlp patching")
continue
mlp_module = getattr(layer, mlp_module_name)
# MLP patching
gate_proj = layer.mlp.gate_proj
up_proj = layer.mlp. up_proj
down_proj = layer.mlp.down_proj
gate_proj = mlp_module.gate_proj
up_proj = mlp_module. up_proj
down_proj = mlp_module.down_proj
if hasattr(gate_proj, "lora_A") and \
hasattr( up_proj, "lora_A") and \
@ -2711,7 +2744,7 @@ class FastLlamaModel:
(len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0):
# https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
layer.mlp.forward = types.MethodType(_apply_lora_mlp, layer.mlp)
mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module)
n_mlp += 1
else:
logger.warning_once(
@ -2796,7 +2829,7 @@ class FastLlamaModel:
# Clear deleted GPU items
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()
clean_gpu_cache()
pass
# Patch for fast inference

40
unsloth/models/mistral.py
View file

@ -99,7 +99,7 @@ def MistralAttention_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
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
@ -191,15 +191,35 @@ def MistralForCausalLM_fast_forward(
if causal_mask is None and past_key_values is None:
bsz, q_len = input_ids.shape
sliding_window = getattr(self.config, "sliding_window", None)
if sliding_window is None or sliding_window == "null" or sliding_window <= 0:
causal_mask = xformers.attn_bias.LowerTriangularMask()
elif q_len <= sliding_window:
causal_mask = xformers.attn_bias.LowerTriangularMask()
else:
causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\
.from_seqlens([q_len]*bsz)\
.make_local_attention(window_size = sliding_window)
pass
if HAS_XFORMERS and attention_mask is None:
if sliding_window is None or sliding_window == "null" or sliding_window <= 0:
causal_mask = xformers.attn_bias.LowerTriangularMask()
elif q_len <= sliding_window:
causal_mask = xformers.attn_bias.LowerTriangularMask()
else:
causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\
.from_seqlens([q_len]*bsz)\
.make_local_attention(window_size = sliding_window)
elif not HAS_XFORMERS and attention_mask is None:
if sliding_window is None or sliding_window == "null" or sliding_window <= 0 or q_len <= sliding_window:
# Fully causal mask
mask = torch.full((q_len, q_len), -torch.inf, device=input_ids.device)
mask = torch.triu(mask, diagonal=1)
attention_mask = mask.expand(bsz, 1, q_len, q_len)
else:
# Sliding window attention
q_indices = torch.arange(q_len, device=input_ids.device).view(-1, 1)
k_indices = torch.arange(q_len, device=input_ids.device).view(1, -1)
causal_bool_mask = k_indices <= q_indices
window_bool_mask = (q_indices - k_indices) < sliding_window
mask = torch.where(causal_bool_mask & window_bool_mask, 0.0, -torch.inf)
attention_mask = mask[None, None, :, :].expand(bsz, 1, q_len, q_len)
attention_mask = attention_mask.to(dtype=_get_dtype(self.config.torch_dtype))
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (

2
unsloth/models/qwen3.py
View file

@ -126,7 +126,7 @@ def Qwen3Attention_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
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)

62
unsloth/models/rl_replacements.py
View file

@ -26,6 +26,8 @@ import torch
import inspect
from collections import defaultdict
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
from unsloth import DEVICE_TYPE
RL_EXTRA_ARGS = defaultdict(list)
RL_FUNCTIONS = defaultdict(list)
RL_PRE_ITEMS = defaultdict(list)
@ -258,7 +260,7 @@ def grpo_trainer__get_per_token_logps(function_name, function):
if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):
with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype):
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
logits = model(
input_ids = input_ids,
@ -291,6 +293,58 @@ def grpo_trainer__get_per_token_logps(function_name, function):
pass
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps)
def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
if function_name != "_get_per_token_logps_and_entropies": return function
# Just copy over from _get_per_token_logps replacement function above. For now this returns None anyway
def _get_per_token_logps_and_entropies(self, model, input_ids, attention_mask, logits_to_keep, batch_size = None, compute_entropy = False):
if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
return {"logps": None, "entropies": None} # Unsloth efficient GRPO
# Otherwise, calculate normally:
if not hasattr(self, '_autocast_dtype'):
self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16
if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': self._autocast_dtype = torch.float16
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
logits = model(
input_ids = input_ids,
attention_mask = attention_mask,
logits_to_keep = logits_to_keep + 1,
).logits
entropies = None
if compute_entropy:
from trl.trainer.utils import entropy_from_logits
entropies = entropy_from_logits(logits)
# logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
return {"logps": logits, "entropies": entropies}
# input_ids = input_ids[:, -logits_to_keep:]
# For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
# See https://github.com/huggingface/trl/issues/2770
# logits = logits[:, -logits_to_keep:]
# return logits
# See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details
# logits = logits / self.temperature
# logps = selective_log_softmax(logits, input_ids)
# row_indices, col_indices = torch.where(logps < -20)
# # Method 1: Check if tensors have elements
# if len(row_indices) > 0 and len(col_indices) > 0:
# breakpoint() # Breakpoint triggered here
# print("Found high values!")
# return logps # compute logprobs for the input tokens
pass
pass
function = inspect.getsource(_get_per_token_logps_and_entropies)
return function
pass
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies)
grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"]
grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"]
UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"]
@ -319,14 +373,16 @@ def grpo_trainer_compute_loss(function_name, function):
_input_ids = input_ids
_logits_to_keep = logits_to_keep
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
get_logps_func = lambda model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False: self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep, batch_size) if hasattr(self, "_get_per_token_logps") else self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)['logps']
per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep)
# Compute the KL divergence between the model and the reference model
# _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves.
# https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328
if self.beta != 0.0:
with torch.inference_mode(), model.disable_adapter():
ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
ref_per_token_logps = per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep)
else:
ref_per_token_logps = None
# per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1