From a4485302690352b75ca319f5e0972bff3c5cfdb0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 1 Mar 2025 00:13:11 -0800 Subject: [PATCH] Prelim release --- unsloth/__init__.py | 19 ------ unsloth/kernels/layernorm.py | 8 +-- unsloth/kernels/rms_layernorm.py | 7 +- unsloth/kernels/swiglu.py | 2 +- unsloth/kernels/utils.py | 95 ++++++++++++++------------ unsloth/models/_utils.py | 66 +++++++----------- unsloth/models/gemma.py | 10 +-- unsloth/models/gemma2.py | 13 ++-- unsloth/models/granite.py | 13 ++-- unsloth/models/llama.py | 113 ++++++++++--------------------- unsloth/models/mistral.py | 24 ++++--- unsloth/tokenizer_utils.py | 15 ---- 12 files changed, 153 insertions(+), 232 deletions(-) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index d18aaac0a0..e33d16577a 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -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"] = \ diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index a5f7926e2e..11f82b8ff3 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -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), diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 6310f7f392..7487c10eeb 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -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), diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index f81b7aae9b..688e9f9a48 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -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 diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index f052914f98..f743e12f59 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 19b09e803c..5088b79d23 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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()) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index bc29c46abc..873bdcf2eb 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -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 diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index be6b0469d9..316b4e8f0e 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -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 diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index fb7e96d8d2..bfaea9a555 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9515a41cd9..815bad8a3e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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 diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 779ff83496..303c3d9589 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -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, diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 048bee7797..9c5f825a0c 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -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