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

Update utils.py

Browse source
This commit is contained in:
Daniel Han 2025-02-05 06:01:38 -08:00
commit dd82d124c1
1 changed files with 8 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

16
unsloth/kernels/utils.py
View file

@ -157,8 +157,8 @@ if HAS_CUDA_STREAM:
global WEIGHT_BUFFER
global ABSMAX_BUFFER
if WEIGHT_BUFFER is None:
WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0")
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0")
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)
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax)
@ -167,11 +167,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")
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")
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False)
pass
# NF4 dequantization of statistics
@ -224,8 +224,8 @@ else:
global WEIGHT_BUFFER
global ABSMAX_BUFFER
if WEIGHT_BUFFER is None:
WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0")
ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = dtype, device = "cuda:0")
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)
@ -234,11 +234,11 @@ else:
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
else:
if out is None:
out = torch.empty(shape, dtype = dtype, device = "cuda:0")
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")
out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False)
pass
# Do dequantization