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

torch_cuda_device

Browse source
This commit is contained in:
Daniel Han 2025-03-02 23:31:16 -08:00
commit 7e096ff9bf
7 changed files with 140 additions and 115 deletions
Download patch file Download diff file Expand all files Collapse all files

91
unsloth/kernels/cross_entropy_loss.py
View file

@ -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

18
unsloth/kernels/geglu.py
View file

@ -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

48
unsloth/kernels/layernorm.py
View file

@ -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

47
unsloth/kernels/rms_layernorm.py
View file

@ -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

42
unsloth/kernels/rope_embedding.py
View file

@ -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

8
unsloth/kernels/swiglu.py
View file

@ -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

1
unsloth/kernels/utils.py
View file

@ -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