diff options
| author | Tim Dettmers <dettmers@cs.washington.edu> | 2021-10-20 19:15:47 -0700 |
|---|---|---|
| committer | Tim Dettmers <dettmers@cs.washington.edu> | 2021-10-20 19:15:47 -0700 |
| commit | a6eae2e7f2bf03f268fcb6b055201ff6827684c4 (patch) | |
| tree | d2f72792251c9feaef1cf9dcddc3c79e6312a93a /bitsandbytes/functional.py | |
| parent | bb34fd50a1fec74e62beb6e23d51f0142c7d0ab6 (diff) | |
Added skip_zeros; tests are passing.
Diffstat (limited to 'bitsandbytes/functional.py')
| -rw-r--r-- | bitsandbytes/functional.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index 48ab40c..9fe1345 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -486,13 +486,13 @@ def optimizer_update_8bit_blockwise(optimizer_name: str, g: Tensor, p: Tensor, s str2optimizer8bit_blockwise[optimizer_name][0](get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), - get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), + get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel())) elif g.dtype == torch.float16 and state1.dtype == torch.uint8: str2optimizer8bit_blockwise[optimizer_name][1](get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), - get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), + get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel())) else: raise ValueError(f'Gradient+optimizer bit data type combination not supported: grad {g.dtype}, optimizer {state1.dtype}') |