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

Prelim release

Browse source
This commit is contained in:
Daniel Han 2025-03-01 00:13:11 -08:00
commit a448530269
12 changed files with 152 additions and 231 deletions
Download patch file Download diff file Expand all files Collapse all files

19
unsloth/__init__.py
View file

@ -46,25 +46,6 @@ pass
# Fixes https://github.com/unslothai/unsloth/issues/1266
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
if "CUDA_VISIBLE_DEVICES" in os.environ:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
devices = os.environ["CUDA_VISIBLE_DEVICES"]
# Check if there are multiple cuda devices set in env
if not devices.isdigit():
first_id = devices.split(",")[0]
warnings.warn(
f"Unsloth: 'CUDA_VISIBLE_DEVICES' is currently {devices} \n"\
"Unsloth currently does not support multi GPU setups - but we are working on it!\n"\
"Multiple CUDA devices detected but we require a single device.\n"\
f"We will override CUDA_VISIBLE_DEVICES to first device: {first_id}."
)
os.environ["CUDA_VISIBLE_DEVICES"] = str(first_id)
else:
# warnings.warn("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.")
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
pass
# Reduce VRAM usage by reducing fragmentation
# And optimize pinning of memory
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \

8
unsloth/kernels/layernorm.py
View file

@ -105,10 +105,10 @@ class Fast_Layernorm(torch.autograd.Function):
X = X.view(-1, dim)
n_rows, n_cols = X.shape
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
mu = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
device = X.device
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = device)
r = torch.empty(n_rows, dtype = torch.float32, device = device)
mu = torch.empty(n_rows, dtype = torch.float32, device = device)
layernorm_forward[(n_rows,)](
Y, Y.stride(0),

7
unsloth/kernels/rms_layernorm.py
View file

@ -148,9 +148,10 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
BLOCK_SIZE : int
num_warps : int
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
device = X.device
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = device)
r = torch.empty(n_rows, dtype = torch.float32, device = device)
fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward
fx[(n_rows,)](
@ -180,7 +181,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
n_cols : int
n_rows, n_cols = dY.shape
# dW = X
dX = torch.empty_like(dY, device = "cuda:0") if ctx.GEMMA else dY
dX = torch.empty_like(dY) if ctx.GEMMA else dY
_rms_layernorm_backward[(n_rows,)](
dY, dY.stride(0),

2
unsloth/kernels/swiglu.py
View file

@ -41,7 +41,7 @@ pass
def swiglu_fg_kernel(e, g):
batch, seq_len, hd = e.shape
n_elements = e.numel()
h = torch.empty((batch, seq_len, hd), dtype = e.dtype, device = "cuda:0")
h = torch.empty((batch, seq_len, hd), dtype = e.dtype, device = e.device)
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
_fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE = 1024,)
return h

95
unsloth/kernels/utils.py
View file

@ -61,12 +61,29 @@ pass
import bitsandbytes as bnb
import ctypes
# https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
global CUDA_STREAM
CUDA_STREAM = None
get_ptr = bnb.functional.get_ptr
import ctypes
# Get array of CUDA streams and other buffers
global CUDA_STREAMS
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
_CUDA_STREAMS = {
(index := torch.cuda.device(i).idx) : ctypes.c_void_p(torch._C._cuda_getCurrentRawStream(index))
for i in range(torch.cuda.device_count())
}
CUDA_STREAMS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
WEIGHT_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
ABSMAX_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
for k, v in _CUDA_STREAMS.items(): CUDA_STREAMS[k] = v
CUDA_STREAMS = tuple(CUDA_STREAMS)
del _CUDA_STREAMS
# Bitsandbytes operations
ctypes_c_int = ctypes.c_int
ctypes_c_int32 = ctypes.c_int32
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
@ -118,11 +135,6 @@ def get_lora_parameters_bias(proj):
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:
@torch.inference_mode
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
@ -145,8 +157,10 @@ if HAS_CUDA_STREAM:
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAM
if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0")
global CUDA_STREAMS
device = W.device
device_index = device.index
CUDA_STREAM = CUDA_STREAMS[device_index]
n_elements_absmax = absmax.numel()
@ -155,11 +169,13 @@ if HAS_CUDA_STREAM:
# Use same buffers for faster inference
size = shape[0]*shape[1]
global WEIGHT_BUFFER
global ABSMAX_BUFFER
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index]
ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index]
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)
WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = device, requires_grad = False)
ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, 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)
@ -168,11 +184,11 @@ if HAS_CUDA_STREAM:
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
else:
if out is None:
out = torch.empty(shape, dtype = dtype, device = "cuda:0", requires_grad = False)
out = torch.empty(shape, dtype = dtype, device = device, 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)
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False)
pass
# NF4 dequantization of statistics
@ -217,31 +233,15 @@ else:
pass
n_elements_absmax = absmax.numel()
device = W.device
# 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]
if out is None:
out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False)
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
assert(out.shape == shape)
assert(out.dtype == dtype)
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False)
# Do dequantization
ptr_out_absmax = get_ptr(out_absmax)
@ -288,14 +288,16 @@ if HAS_CUDA_STREAM:
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAM
if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0")
global CUDA_STREAMS
device = W.device
device_index = device.index
CUDA_STREAM = CUDA_STREAMS[device_index]
# assert(dtype == X.dtype)
bout = shape[0]
if out is None:
out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda:0")
out = torch.empty((1, 1, bout,), dtype = dtype, device = device)
# else:
# assert(out.shape == (1, 1, bout,))
# pass
@ -313,7 +315,7 @@ if HAS_CUDA_STREAM:
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
df = torch.empty(absmax.shape, dtype = torch.float32, device = device)
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,
@ -357,9 +359,10 @@ else:
pass
# assert(dtype == X.dtype)
bout = shape[0]
device = W.device
if out is None:
out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda:0")
out = torch.empty((1, 1, bout,), dtype = dtype, device = device)
# else:
# assert(out.shape == (1, 1, bout,))
# pass
@ -377,7 +380,7 @@ else:
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0")
df = torch.empty(absmax.shape, dtype = torch.float32, device = device)
cdequantize_blockwise_fp32(
get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df),
ctypes_c_int(blocksize2), ctypes_c_int(df.numel()),
@ -400,6 +403,7 @@ pass
torch_mm = torch.mm
torch_mv = torch.mv
torch_matmul = torch.matmul
torch_addmm = torch.addmm
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)
@ -461,7 +465,8 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
if A is not None:
# LoRA is enabled
A, B = A.t(), B.t()
out += (X @ A.to(dtype)) @ (s * B.to(dtype))
out = torch_addmm(X @ A.to(dtype), B.to(dtype), alpha = s, beta = 1.0, out = out)
# out += (X @ A.to(dtype)) @ (s * B.to(dtype))
pass
return out.view(batch, seq_len, -1) if reshape else out

66
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2025.2.15"
__version__ = "2025.3.1"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -37,7 +37,6 @@ __all__ = [
"torch_compile_options",
"patch_linear_scaling",
"patch_llama_rope_scaling",
"check_nvidia",
"create_boolean_mask",
"torch_amp_custom_fwd",
"torch_amp_custom_bwd",
@ -703,9 +702,7 @@ pass
# =============================================
# Fixes Bitsandbytes to remove missing warnings
from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
from inspect import getsource
from accelerate.utils.dataclasses import DistributedType
BitsAndBytesConfig__init__ = getsource(BitsAndBytesConfig.__init__)
BitsAndBytesConfig__init__ = inspect.getsource(BitsAndBytesConfig.__init__)
BitsAndBytesConfig__init__ = re.sub(
r"if[\s]{1,}kwargs\:[\s]{1,}.+?\n",
"",
@ -719,28 +716,30 @@ BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace(
"__init__",
"_BitsAndBytesConfig__init__",
)
def _prepare_backend(
self, cpu = False, sagemaker_dp = False, backend: str = None,
) -> tuple[str, DistributedType]:
return None, DistributedType.NO
pass
import accelerate.state
accelerate.state.PartialState._prepare_backend = _prepare_backend
import accelerate.accelerator
prepare = inspect.getsource(accelerate.accelerator.Accelerator.prepare)
prepare = prepare.split("\n")
spaces = prepare[0].find("def")
prepare = "\n".join(x[spaces:] for x in prepare)
x = "for obj in args:"
s = " "*spaces
prepare = prepare.replace(x, f'self.state.distributed_type = DistributedType.NO\n{s}{x}', 1)
exec(prepare, globals())
accelerate.accelerator.Accelerator.prepare = prepare
exec(BitsAndBytesConfig__init__, globals())
if torch.cuda.device_count() == 1:
from accelerate.utils.dataclasses import DistributedType
def _prepare_backend(
self, cpu = False, sagemaker_dp = False, backend: str = None,
) -> tuple[str, DistributedType]:
return None, DistributedType.NO
pass
import accelerate.state
accelerate.state.PartialState._prepare_backend = _prepare_backend
import accelerate.accelerator
prepare = inspect.getsource(accelerate.accelerator.Accelerator.prepare)
prepare = prepare.split("\n")
spaces = prepare[0].find("def")
prepare = "\n".join(x[spaces:] for x in prepare)
x = "for obj in args:"
s = " "*spaces
prepare = prepare.replace(x, f'self.state.distributed_type = DistributedType.NO\n{s}{x}', 1)
exec(prepare, globals())
accelerate.accelerator.Accelerator.prepare = prepare
pass
import transformers.utils.quantization_config
transformers.utils.quantization_config.BitsAndBytesConfig.__init__ = _BitsAndBytesConfig__init__
# =============================================
@ -963,21 +962,6 @@ def patch_llama_rope_scaling(
pass
def check_nvidia():
# Unsloth doesn't work yet on AMD devices - we're working on it!
output = np.array([0,])
try:
output = subprocess.check_output("nvidia-smi --query-gpu=memory.used --format=csv", shell = True)
output = re.findall(rb'([\d]{1,})[\s]{1,}M', output)
output = np.array([int(x.decode('utf-8'))/1024 for x in output])
except:
if not torch.cuda.is_available():
raise RuntimeError("Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!")
return output
pass
PRE_CHECK = check_nvidia()
def create_boolean_mask(n = 4096, sliding_window = 2048):
# Creates a boolean mask for attention
mask = torch.ones(n, n, dtype = torch.bool)
@ -1122,8 +1106,6 @@ def patch_gradient_accumulation_fix(Trainer):
items_in_trainer = dir(transformers.trainer)
good_items = []
for item in items_in_trainer:
# TODO: Support Deepspeed
if item.startswith(("deepspeed", "xm", "met", "smp")): continue
if item in function: good_items.append(item)
pass
exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals())

10
unsloth/models/gemma.py
View file

@ -245,8 +245,8 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
emb = torch.cat((radians_new, radians_new), dim = -1)
# We must do RoPE in float32!
cos = emb.cos().to(device = "cuda:0", non_blocking = True)#, dtype = dtype)
sin = emb.sin().to(device = "cuda:0", non_blocking = True)#, dtype = dtype)
cos = emb.cos().to(device = "cuda", non_blocking = True)#, dtype = dtype)
sin = emb.sin().to(device = "cuda", non_blocking = True)#, dtype = dtype)
self.register_buffer("cos_cached", cos, persistent = False)
self.register_buffer("sin_cached", sin, persistent = False)
pass
@ -270,7 +270,7 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
if seq_len <= self.current_rope_size: return
# Iteratively grow by increments of 8192
self.current_rope_size = math.ceil(seq_len / 8192) * 8192
self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype)
self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype)
pass
pass
@ -304,8 +304,8 @@ class GemmaFixedLinearScalingRotaryEmbedding(GemmaFixedRotaryEmbedding):
emb = torch.cat((radians_new, radians_new), dim = -1)
# We must do RoPE in float32!
cos = emb.cos().to(device = "cuda:0", non_blocking = True)#, dtype = dtype)
sin = emb.sin().to(device = "cuda:0", non_blocking = True)#, dtype = dtype)
cos = emb.cos().to(device = "cuda", non_blocking = True)#, dtype = dtype)
sin = emb.sin().to(device = "cuda", non_blocking = True)#, dtype = dtype)
self.register_buffer("cos_cached", cos, persistent = False)
self.register_buffer("sin_cached", sin, persistent = False)
pass

13
unsloth/models/gemma2.py
View file

@ -265,21 +265,22 @@ def Gemma2Attention_fast_forward_inference(
attention_size = n_heads*head_dim
seq_len = K1.shape[-2]
kv_seq_len = seq_len + 1
device = hidden_states.device
# Prefill phase
# if not hasattr(self, "paged_attention"):
if do_prefill:
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0")
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device)
self.paged_attention_K = self.paged_attention[:,0]
self.paged_attention_V = self.paged_attention[:,1]
self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)
self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0")
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0")
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device)
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device)
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
# Only for Gemma2
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0")
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0")
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device)
# See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e
# Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below

13
unsloth/models/granite.py
View file

@ -274,21 +274,22 @@ def GraniteAttention_fast_forward_inference(
attention_size = n_heads*head_dim
seq_len = K1.shape[-2]
kv_seq_len = seq_len + 1
device = hidden_states.device
# Prefill phase
# if not hasattr(self, "paged_attention"):
if do_prefill:
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0")
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device)
self.paged_attention_K = self.paged_attention[:,0]
self.paged_attention_V = self.paged_attention[:,1]
self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)
self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0")
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0")
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device)
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
# Only for Gemma2
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0")
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0")
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device)
self.half_head_dim = head_dim // 2

113
unsloth/models/llama.py
View file

@ -167,24 +167,25 @@ def LlamaAttention_fast_forward_inference(
# Prefill phase
# if not hasattr(self, "paged_attention"):
device = hidden_states.device
if do_prefill:
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0")
self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device)
self.paged_attention_K = self.paged_attention[:,0]
self.paged_attention_V = self.paged_attention[:,1]
self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)
self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0")
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0")
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device)
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device)
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
# Mistral Nemo 12b has weird dimensions
if attention_size != hidden_size:
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0")
self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
else:
self.temp_O = self.temp_QA[1][:,:,:hidden_size]
pass
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0")
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device)
self.scalar = 1.0 / math_sqrt(self.head_dim)
self.half_head_dim = head_dim // 2
elif kv_seq_len >= self.paged_attention.shape[0]:
@ -813,13 +814,13 @@ def LlamaModel_fast_forward(
is_causal = True,
sliding_window = self.config.sliding_window,
)\
.to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\
.to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda",)\
.squeeze(0).squeeze(0)
self.GA_mask = AttentionMaskConverter(
is_causal = True,
)\
.to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\
.to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda",)\
.squeeze(0).squeeze(0)
pass
pass
@ -1075,10 +1076,16 @@ def CausalLM_fast_forward(fast_forward_inference):
bsz, q_len, hd = hidden_states.shape
lm_head = self.lm_head.weight
lm_head_device = lm_head.device
logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
logit_scaling = getattr(self.config, "logit_scale", 0)
dtype = lm_head.dtype
num_logits_to_keep = max(num_logits_to_keep, logits_to_keep)
# Move items to same device as lm_head
hidden_states = hidden_states.to(lm_head_device)
if labels is not None: labels = labels.to(lm_head_device)
# Output last hidden states without logits if asked
if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1":
@ -1148,11 +1155,14 @@ def CausalLM_fast_forward(fast_forward_inference):
if labels is not None:
shift_logits = logits
if not hasattr(self, "extra_ignored_labels"):
# Fixes https://github.com/unslothai/unsloth/issues/10
self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
pass
shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
# if not hasattr(self, "extra_ignored_labels"):
# # Fixes https://github.com/unslothai/unsloth/issues/10
# self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
# pass
shift_labels = torch.empty_like(labels)
shift_labels[..., :-1] = labels[..., 1:]
shift_labels[..., -1] = -100
# shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
loss = fast_cross_entropy_loss(
logits = shift_logits,
labels = shift_labels,
@ -1297,7 +1307,7 @@ class LlamaRotaryEmbedding(torch.nn.Module):
if seq_len <= self.current_rope_size: return
# Iteratively grow by increments of 8192
self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192
self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype)
self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype)
pass
pass
@ -1423,7 +1433,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module):
if seq_len <= self.current_rope_size: return
# Iteratively grow by increments of 8192
self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192
self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype)
self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype)
pass
pass
@ -1538,7 +1548,7 @@ class LongRopeRotaryEmbedding(torch.nn.Module):
if seq_len <= self.current_rope_size: return
# Iteratively grow by increments of 8192
self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192
self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype)
self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype)
pass
pass
@ -1771,8 +1781,6 @@ class FastLlamaModel:
# Add to kwargs
kwargs["rope_scaling"] = rope_scaling
pass
# We currently only support NVIDIA GPUs - AMD / Intel is a work in progress!
pre_check = check_nvidia()
bnb_config = None
if load_in_4bit:
@ -1840,8 +1848,6 @@ class FastLlamaModel:
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!
post_check = check_nvidia()
# Counteract saved tokenizers
tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
@ -1882,8 +1888,6 @@ class FastLlamaModel:
items_in_trainer = dir(transformers.trainer)
good_items = []
for item in items_in_trainer:
# TODO: Support Deepspeed
if item.startswith(("deepspeed", "xm", "met", "smp")): continue
if item in inner_training_loop: good_items.append(item)
pass
exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals())
@ -1903,17 +1907,7 @@ class FastLlamaModel:
f"{chr(92)} / Total batch size = {total_train_batch_size:,} | Total steps = {max_steps:,}\\n"\\
f' "-____-" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}'
logger.warning(debug_info)
import subprocess, re, gc, numpy as np
a = np.array([0,])
try:
a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True)
a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a)
a = np.array([int(x.decode('utf-8'))/1024 for x in a])
except:
if not torch.cuda.is_available():
raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')
if ((a - PRE_CHECK) >= 1).sum() > 1:
raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')
import subprocess, re, gc
for _ in range(3):
gc.collect()
torch.cuda.empty_cache()"""
@ -1925,7 +1919,7 @@ class FastLlamaModel:
debug_info = """n_total_devices = total_train_batch_size // \\
args.gradient_accumulation_steps // self._train_batch_size
if n_total_devices > 1:
logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!')
logger.warning_once('Unsloth is running with multi GPUs - the effective batch size is multiplied by ' + str(n_total_devices))
debug_info ="""
debug_info = debug_info.split('\n')
debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]])
@ -1937,31 +1931,6 @@ class FastLlamaModel:
"train_dataloader = tpu_spmd_dataloader(train_dataloader)",
"raise RuntimeError('Unsloth: TPUs are not yet supported!')"
)
inner_training_loop = inner_training_loop.replace(
"self.accelerator.free_memory()",
"self.accelerator.free_memory()\n" + \
front_spaces + "if self.is_deepspeed_enabled:"\
"raise RuntimeError('Unsloth: Deepspeed is not yet supported!')\n", 1,
)
check_batches = """train_dataloader = self.get_train_dataloader()
ga = args.gradient_accumulation_steps
bsz = self._train_batch_size
total_batches = bsz * ga * args.world_size
n_total_devices = total_batches // ga // bsz
if n_total_devices > 1:
logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!')
divisor = n_total_devices / 1
bsz = self._train_batch_size = max(int(bsz / divisor), 1)
if total_batches // ga // bsz > 1:
divisor = n_total_devices / 1
ga = args.gradient_accumulation_steps = max(int(ga / divisor), 1)"""
check_batches = check_batches.split('\n')
check_batches = "\n".join([check_batches[0]] + [front_spaces + x[8:] for x in check_batches[1:]])
inner_training_loop = inner_training_loop.replace(
"train_dataloader = self.get_train_dataloader()",
check_batches, 1,
)
inner_training_loop = inner_training_loop.replace(
"_inner_training_loop",
"_fast_inner_training_loop", 1,
@ -1973,13 +1942,6 @@ class FastLlamaModel:
"is_torch_tpu_available()",
"False",
)
if "n_total_devices >" not in inner_training_loop:
raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')
pass
inner_training_loop = inner_training_loop.replace(
"is_sagemaker_mp_enabled()",
"False",
)
exec(inner_training_loop, globals())
Trainer._inner_training_loop = _fast_inner_training_loop
@ -2136,7 +2098,7 @@ class FastLlamaModel:
pass
model.get_input_embeddings().modules_to_save.default\
.to(device = "cuda:0", dtype = new_dtype, non_blocking = True)
.to(device = "cuda", dtype = new_dtype, non_blocking = True)
model.get_input_embeddings().modules_to_save.default.requires_grad_(True)
# [TODO] Move old embed_tokens to CPU - should be disk!
@ -2156,7 +2118,7 @@ class FastLlamaModel:
pass
model.get_output_embeddings().modules_to_save.default\
.to(device = "cuda:0", dtype = new_dtype, non_blocking = True)
.to(device = "cuda", dtype = new_dtype, non_blocking = True)
model.get_output_embeddings().modules_to_save.default.requires_grad_(True)
# [TODO] Move old lm_head to CPU - should be disk!
@ -2413,7 +2375,7 @@ class FastLlamaModel:
pass
model.get_input_embeddings().modules_to_save.default\
.to(device = "cuda:0", dtype = new_dtype, non_blocking = True)
.to(device = "cuda", dtype = new_dtype, non_blocking = True)
model.get_input_embeddings().modules_to_save.default.requires_grad_(True)
pass
@ -2429,7 +2391,7 @@ class FastLlamaModel:
pass
model.get_output_embeddings().modules_to_save.default\
.to(device = "cuda:0", dtype = new_dtype, non_blocking = True)
.to(device = "cuda", dtype = new_dtype, non_blocking = True)
model.get_output_embeddings().modules_to_save.default.requires_grad_(True)
pass
@ -2515,12 +2477,7 @@ class FastLlamaModel:
from transformers.trainer import Trainer
if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop":
raise RuntimeError(
'Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so '\
'enabling it will require much more work, so we have to prioritize. Please understand!\n'\
'We do have a separate beta version, which you can contact us about!\n'\
'Thank you for your understanding and we appreciate it immensely!'
)
raise RuntimeError("Unsloth: Unsuccessfully patched Trainer! Please file a bug report!")
pass
# Fix loftq issues
@ -2636,8 +2593,8 @@ class FastLlamaModel:
# Patch cross entropy loss labels
# Fixes https://github.com/unslothai/unsloth/issues/10
max_seq_length = model.max_seq_length
extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda:0")
model.model.extra_ignored_labels = extra_ignored_labels
# extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda:0")
# model.model.extra_ignored_labels = extra_ignored_labels
internal_model = model
while hasattr(internal_model, "model"):
internal_model.max_seq_length = max_seq_length

24
unsloth/models/mistral.py
View file

@ -235,6 +235,14 @@ def MistralForCausalLM_fast_forward(
hidden_states = outputs[0]
bsz, q_len, hd = hidden_states.shape
lm_head = self.lm_head.weight
lm_head_device = lm_head.device
# Move items to same device as lm_head
hidden_states = hidden_states.to(lm_head_device)
if labels is not None: labels = labels.to(lm_head_device)
# If we are in GRPO mode, return raw hidden states
if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1":
num_logits_to_keep = max(num_logits_to_keep, logits_to_keep)
@ -249,8 +257,6 @@ def MistralForCausalLM_fast_forward(
)
pass
bsz, q_len, hd = hidden_states.shape
lm_head = self.lm_head.weight
if bsz == 1 and q_len == 1:
logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
logits = logits.unsqueeze(0).unsqueeze(0)
@ -292,12 +298,14 @@ def MistralForCausalLM_fast_forward(
loss = None
if labels is not None:
shift_logits = logits
if not hasattr(self, "extra_ignored_labels"):
# Fixes https://github.com/unslothai/unsloth/issues/10
self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
pass
shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
# if not hasattr(self, "extra_ignored_labels"):
# # Fixes https://github.com/unslothai/unsloth/issues/10
# self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
# pass
# shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
shift_labels = torch.empty_like(labels)
shift_labels[..., :-1] = labels[..., 1:]
shift_labels[..., -1] = -100
loss = fast_cross_entropy_loss(
logits = shift_logits,
labels = shift_labels,

15
unsloth/tokenizer_utils.py
View file

@ -857,21 +857,6 @@ def check_tokenizer(
pass
def check_nvidia():
# Unsloth doesn't work yet on AMD devices - we're working on it!
output = np.array([0,])
try:
output = subprocess.check_output("nvidia-smi --query-gpu=memory.used --format=csv", shell = True)
output = re.findall(rb'([\d]{1,})[\s]{1,}M', output)
output = np.array([int(x.decode('utf-8'))/1024 for x in output])
except:
if not torch.cuda.is_available():
raise RuntimeError("Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!")
return output
pass
PRE_CHECK = check_nvidia()
import inspect
from inspect import getsource
import trl.trainer.sft_trainer