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-02 16:29:26 -08:00
commit 91c1964ba4
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/utils.py
View file

@ -154,12 +154,13 @@ if HAS_CUDA_STREAM:
size = shape[0]*shape[1]
global WEIGHT_BUFFER
global ABSMAX_BUFFER
print(size, shape)
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")
if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel(): WEIGHT_BUFFER.resize_(n_elements_absmax)
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]
@ -219,7 +220,6 @@ else:
# Use same buffers for faster inference
size = shape[0]*shape[1]
print(shape, size)
global WEIGHT_BUFFER
global ABSMAX_BUFFER
if WEIGHT_BUFFER is None: