From 7e096ff9bff7c6b8feb777ecdfaf40c304e56b2d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 2 Mar 2025 23:31:16 -0800 Subject: [PATCH] torch_cuda_device --- unsloth/kernels/cross_entropy_loss.py | 91 +++++++++++++++------------ unsloth/kernels/geglu.py | 18 ++++-- unsloth/kernels/layernorm.py | 48 +++++++------- unsloth/kernels/rms_layernorm.py | 47 +++++++------- unsloth/kernels/rope_embedding.py | 42 +++++++------ unsloth/kernels/swiglu.py | 8 ++- unsloth/kernels/utils.py | 1 + 7 files changed, 140 insertions(+), 115 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 1c9998e1c9..006dfff631 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -15,7 +15,13 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh, triton_cast +from .utils import ( + calculate_settings, + MAX_FUSED_SIZE, + triton_tanh, + triton_cast, + torch_cuda_device, +) from transformers.models.llama.modeling_llama import logger from packaging.version import Version @@ -295,37 +301,39 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): BLOCK_SIZE, num_warps = calculate_settings(vocab_size) logsumexp = torch.empty(n_rows, dtype = torch.float32, device = device) - _cross_entropy_forward[(n_rows,)]( - logits, logits.stride(0), - losses, - logsumexp, - labels, - VOCAB_SIZE = vocab_size, - BLOCK_SIZE = BLOCK_SIZE, - DO_SOFTCAPPING = DO_SOFTCAPPING, - SOFTCAP = logit_softcapping, - DO_LOGIT_SCALING = DO_LOGIT_SCALING, - LOGIT_SCALE = logit_scaling, - num_warps = num_warps, - ) + with torch_cuda_device(device): + _cross_entropy_forward[(n_rows,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE = vocab_size, + BLOCK_SIZE = BLOCK_SIZE, + DO_SOFTCAPPING = DO_SOFTCAPPING, + SOFTCAP = logit_softcapping, + DO_LOGIT_SCALING = DO_LOGIT_SCALING, + LOGIT_SCALE = logit_scaling, + num_warps = num_warps, + ) else: # For large vocabs > 65336 like Gemma 256K logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = device) - _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( - logits, logits.stride(0), - losses, - logsumexp, - labels, - VOCAB_SIZE = vocab_size, - N_CHUNKS = n_chunks, - BLOCK_SIZE = MAX_FUSED_SIZE, - DO_SOFTCAPPING = DO_SOFTCAPPING, - SOFTCAP = logit_softcapping, - DO_LOGIT_SCALING = DO_LOGIT_SCALING, - LOGIT_SCALE = logit_scaling, - num_warps = 32, - ) + with torch_cuda_device(device): + _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE = vocab_size, + N_CHUNKS = n_chunks, + BLOCK_SIZE = MAX_FUSED_SIZE, + DO_SOFTCAPPING = DO_SOFTCAPPING, + SOFTCAP = logit_softcapping, + DO_LOGIT_SCALING = DO_LOGIT_SCALING, + LOGIT_SCALE = logit_scaling, + num_warps = 32, + ) # logsumexp(chunked_logsumexp) - x # Do the -x separately logsumexp = torch.logsumexp(logsumexp, dim = 1) # Row sum @@ -355,19 +363,20 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): div, mod = divmod(vocab_size, BLOCK_SIZE) n_blocks : int = div + (mod != 0) - _cross_entropy_backward[(n_rows, n_blocks,)]( - logits, logits.stride(0), - dlosses, dlosses.stride(0), - logsumexp, - labels, - VOCAB_SIZE = vocab_size, - BLOCK_SIZE = BLOCK_SIZE, - DO_SOFTCAPPING = ctx.DO_SOFTCAPPING, - SOFTCAP = ctx.logit_softcapping, - DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING, - LOGIT_SCALE = ctx.logit_scaling, - num_warps = 8, - ) + with torch_cuda_device(dlosses.device): + _cross_entropy_backward[(n_rows, n_blocks,)]( + logits, logits.stride(0), + dlosses, dlosses.stride(0), + logsumexp, + labels, + VOCAB_SIZE = vocab_size, + BLOCK_SIZE = BLOCK_SIZE, + DO_SOFTCAPPING = ctx.DO_SOFTCAPPING, + SOFTCAP = ctx.logit_softcapping, + DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING, + LOGIT_SCALE = ctx.logit_scaling, + num_warps = 8, + ) return logits, None, None, None, pass pass diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index 9fedae769e..d5a69aa67f 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -15,7 +15,11 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings, triton_tanh +from .utils import ( + calculate_settings, + triton_tanh, + torch_cuda_device, +) @triton.jit @@ -43,7 +47,8 @@ def geglu_exact_forward_kernel(gate, up): n_elements = gate.numel() out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda:0") grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - _exact_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) + with torch_cuda_device(gate.device): + _exact_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) return out pass @@ -99,7 +104,8 @@ def geglu_exact_backward_kernel(DW, e, g): batch_seq_len, hd = e.shape n_elements = e.numel() grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - _exact_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) + with torch_cuda_device(e.device): + _exact_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) return DW, e, g pass @@ -135,7 +141,8 @@ def geglu_approx_forward_kernel(gate, up): n_elements = gate.numel() out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda:0") grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - _approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) + with torch_cuda_device(gate.device): + _approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) return out pass @@ -198,6 +205,7 @@ def geglu_approx_backward_kernel(DW, e, g): batch_seq_len, hd = e.shape n_elements = e.numel() grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - _approx_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) + with torch_cuda_device(e.device): + _approx_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) return DW, e, g pass diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index ffcc5cc13c..26a77f03a0 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -16,7 +16,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings +from .utils import calculate_settings, torch_cuda_device from unsloth_zoo.patching_utils import ( patch_layernorm, ) @@ -111,17 +111,18 @@ class Fast_Layernorm(torch.autograd.Function): 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), - X, X.stride(0), - W, - b, - r, - mu, - n_cols, eps, - BLOCK_SIZE = BLOCK_SIZE, - num_warps = num_warps, - ) + with torch_cuda_device(device): + layernorm_forward[(n_rows,)]( + Y, Y.stride(0), + X, X.stride(0), + W, + b, + r, + mu, + n_cols, eps, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, + ) ctx.eps = eps ctx.BLOCK_SIZE = BLOCK_SIZE ctx.num_warps = num_warps @@ -137,17 +138,18 @@ class Fast_Layernorm(torch.autograd.Function): X, W, b, r, mu = ctx.saved_tensors n_rows, n_cols = dY.shape - layernorm_backward[(n_rows,)]( - dY, dY.stride(0), - X, X .stride(0), - W, - b, - r, - mu, - n_cols, ctx.eps, - BLOCK_SIZE = ctx.BLOCK_SIZE, - num_warps = ctx.num_warps, - ) + with torch_cuda_device(dY.device): + layernorm_backward[(n_rows,)]( + dY, dY.stride(0), + X, X .stride(0), + W, + b, + r, + mu, + n_cols, ctx.eps, + BLOCK_SIZE = ctx.BLOCK_SIZE, + num_warps = ctx.num_warps, + ) dX = dY.view(*shape) return dX, None, None, None, None pass diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 7487c10eeb..1cde6388ea 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -15,8 +15,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings - +from .utils import calculate_settings, torch_cuda_device @triton.jit def _rms_layernorm_forward( @@ -154,15 +153,16 @@ class Fast_RMS_Layernorm(torch.autograd.Function): r = torch.empty(n_rows, dtype = torch.float32, device = device) fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward - fx[(n_rows,)]( - Y, Y.stride(0), - X, X.stride(0), - W, W.stride(0), - r, r.stride(0), - n_cols, eps, - BLOCK_SIZE = BLOCK_SIZE, - num_warps = num_warps, - ) + with torch_cuda_device(device): + fx[(n_rows,)]( + Y, Y.stride(0), + X, X.stride(0), + W, W.stride(0), + r, r.stride(0), + n_cols, eps, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, + ) ctx.eps = eps ctx.BLOCK_SIZE = BLOCK_SIZE ctx.num_warps = num_warps @@ -183,18 +183,19 @@ class Fast_RMS_Layernorm(torch.autograd.Function): # dW = X dX = torch.empty_like(dY) if ctx.GEMMA else dY - _rms_layernorm_backward[(n_rows,)]( - dY, dY.stride(0), - dX, dX.stride(0), - X, X .stride(0), - W, W .stride(0), - r, r .stride(0), - # dW, dW.stride(0), - n_cols, ctx.eps, - GEMMA = ctx.GEMMA, - BLOCK_SIZE = ctx.BLOCK_SIZE, - num_warps = ctx.num_warps, - ) + with torch_cuda_device(dY.device): + _rms_layernorm_backward[(n_rows,)]( + dY, dY.stride(0), + dX, dX.stride(0), + X, X .stride(0), + W, W .stride(0), + r, r .stride(0), + # dW, dW.stride(0), + n_cols, ctx.eps, + GEMMA = ctx.GEMMA, + BLOCK_SIZE = ctx.BLOCK_SIZE, + num_warps = ctx.num_warps, + ) dX = dX.view(*shape) return dX, None, None, None pass diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index 88b9ccadb4..a14a485352 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -15,7 +15,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings +from .utils import calculate_settings, torch_cuda_device ROPE_GROUP_SIZE : int = 4 def _rope_embedding( @@ -100,16 +100,17 @@ class Fast_RoPE_Embedding(torch.autograd.Function): div, mod = divmod(n_heads, ROPE_GROUP_SIZE) n_groups : int = div + (mod != 0) - _rope_embedding[(n_rows, n_groups, )]( - Q, Q.stride(0), - cos, cos.stride(0), - sin, sin.stride(0), - seq_len, - head_dim, n_heads, - BACKWARD_PASS = False, - BLOCK_SIZE = BLOCK_SIZE, - num_warps = num_warps, - ) + with torch_cuda_device(Q.device): + _rope_embedding[(n_rows, n_groups, )]( + Q, Q.stride(0), + cos, cos.stride(0), + sin, sin.stride(0), + seq_len, + head_dim, n_heads, + BACKWARD_PASS = False, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, + ) ctx.BLOCK_SIZE = BLOCK_SIZE ctx.num_warps = num_warps ctx.n_groups = n_groups @@ -134,15 +135,16 @@ class Fast_RoPE_Embedding(torch.autograd.Function): cos = ctx.cos sin = ctx.sin - _rope_embedding[(n_rows, ctx.n_groups, )]( - dY, dY .stride(0), - cos, cos.stride(0), - sin, sin.stride(0), - seq_len, head_dim, n_heads, - BACKWARD_PASS = True, - BLOCK_SIZE = ctx.BLOCK_SIZE, - num_warps = ctx.num_warps, - ) + with torch_cuda_device(dY.device): + _rope_embedding[(n_rows, ctx.n_groups, )]( + dY, dY .stride(0), + cos, cos.stride(0), + sin, sin.stride(0), + seq_len, head_dim, n_heads, + BACKWARD_PASS = True, + BLOCK_SIZE = ctx.BLOCK_SIZE, + num_warps = ctx.num_warps, + ) dY = dY.view(batch, seq_len, n_heads, head_dim) return dY, None, None, pass diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index 688e9f9a48..12f1f5e063 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -15,7 +15,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings +from .utils import calculate_settings, torch_cuda_device @triton.jit @@ -43,7 +43,8 @@ def swiglu_fg_kernel(e, g): n_elements = e.numel() 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,) + with torch_cuda_device(e.device): + _fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE = 1024,) return h pass @@ -94,6 +95,7 @@ def swiglu_DWf_DW_dfg_kernel(DW, e, g): batch_seq_len, hd = e.shape n_elements = e.numel() grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - _DWf_DW_dfg_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) + with torch_cuda_device(e.device): + _DWf_DW_dfg_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) return DW, e, g pass diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 985adaaa44..4439a47f23 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -27,6 +27,7 @@ else: torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda") torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda") pass +torch_cuda_device = torch.cuda.device # tl.math.tanh now is libdevice.tanh