Skip to content

GitLab

Unverified Commit 1f2ca43a authored by Titus's avatar Titus Committed by GitHub
Browse files

Merge pull request #1222 from EtienneDosSantos/main 

Add erroneously missing optimizers to `str2optimizer32bit`
parents d9b1125c 7a338db2
Show whitespace changes
Inline Side-by-side
Showing with 89 additions and 33 deletions
bitsandbytes/functional.py
View file @ 1f2ca43a
...@@ -27,79 +27,135 @@ name2qmap = {} ...@@ -27,79 +27,135 @@ name2qmap = {}
if lib and lib.compiled_with_cuda: if lib and lib.compiled_with_cuda:
"""C FUNCTIONS FOR OPTIMIZERS""" """C FUNCTIONS FOR OPTIMIZERS"""
str2optimizer32bit = { str2optimizer32bit = {
"adagrad": (
lib.cadagrad32bit_grad_fp32,
lib.cadagrad32bit_grad_fp16,
),
"adam": ( "adam": (
lib.cadam32bit_grad_fp32, lib.cadam32bit_grad_fp32,
lib.cadam32bit_grad_fp16, lib.cadam32bit_grad_fp16,
lib.cadam32bit_grad_bf16, lib.cadam32bit_grad_bf16,
), ),
"momentum": ( "pagedadam": (
lib.cmomentum32bit_grad_32, lib.cpagedadam32bit_grad_fp32,
lib.cmomentum32bit_grad_16, lib.cpagedadam32bit_grad_fp16,
lib.cpagedadam32bit_grad_bf16,
), ),
"rmsprop": ( "adamw": (
lib.crmsprop32bit_grad_32, lib.cadam32bit_grad_fp32,
lib.crmsprop32bit_grad_16, lib.cadam32bit_grad_fp16,
lib.cadam32bit_grad_bf16,
),
"pagedadamw": (
lib.cpagedadam32bit_grad_fp32,
lib.cpagedadam32bit_grad_fp16,
lib.cpagedadam32bit_grad_bf16,
),
"lamb": (
lib.cadam32bit_grad_fp32,
lib.cadam32bit_grad_fp16,
),
"lars": (
lib.clars32bit_grad_fp32,
lib.clars32bit_grad_fp16,
), ),
"lion": ( "lion": (
lib.clion32bit_grad_fp32, lib.clion32bit_grad_fp32,
lib.clion32bit_grad_fp16, lib.clion32bit_grad_fp16,
lib.clion32bit_grad_bf16, lib.clion32bit_grad_bf16,
), ),
"adagrad": ( "momentum": (
lib.cadagrad32bit_grad_32, lib.cmomentum32bit_grad_fp32,
lib.cadagrad32bit_grad_16, lib.cmomentum32bit_grad_fp16,
),
"rmsprop": (
lib.crmsprop32bit_grad_fp32,
lib.crmsprop32bit_grad_fp16,
), ),
} }
str2optimizer8bit = { str2optimizer8bit = {
"adagrad": (
lib.cadagrad8bit_grad_fp32,
lib.cadagrad8bit_grad_fp16,
),
"adam": ( "adam": (
lib.cadam_static_8bit_grad_32, lib.cadam_static_8bit_grad_fp32,
lib.cadam_static_8bit_grad_16, lib.cadam_static_8bit_grad_fp16,
), ),
"momentum": ( "pagedadam": (
lib.cmomentum_static_8bit_grad_32, lib.cpagedadam8bit_grad_fp32,
lib.cmomentum_static_8bit_grad_16, lib.cpagedadam8bit_grad_fp16,
lib.cpagedadam8bit_grad_bf16,
), ),
"rmsprop": ( "adamw": (
lib.crmsprop_static_8bit_grad_32, lib.cadam_static_8bit_grad_fp32,
lib.crmsprop_static_8bit_grad_16, lib.cadam_static_8bit_grad_fp16,
), ),
"lion": ( "pagedadamw": (
lib.clion_static_8bit_grad_32, lib.cpagedadam8bit_grad_fp32,
lib.clion_static_8bit_grad_16, lib.cpagedadam8bit_grad_fp16,
lib.cpagedadam8bit_grad_bf16,
), ),
"lamb": ( "lamb": (
lib.cadam_static_8bit_grad_32, lib.cadam_static_8bit_grad_fp32,
lib.cadam_static_8bit_grad_16, lib.cadam_static_8bit_grad_fp16,
), ),
"lars": ( "lars": (
lib.cmomentum_static_8bit_grad_32, lib.clars8bit_grad_fp32,
lib.cmomentum_static_8bit_grad_16, lib.clars8bit_grad_fp16,
),
"lion": (
lib.clion_static_8bit_grad_fp32,
lib.clion_static_8bit_grad_fp16,
),
"momentum": (
lib.cmomentum_static_8bit_grad_fp32,
lib.cmomentum_static_8bit_grad_fp16,
),
"rmsprop": (
lib.crmsprop_static_8bit_grad_fp32,
lib.crmsprop_static_8bit_grad_fp16,
), ),
} }
str2optimizer8bit_blockwise = { str2optimizer8bit_blockwise = {
"adagrad": (
lib.cadagrad_8bit_blockwise_grad_fp32,
lib.cadagrad_8bit_blockwise_grad_fp16,
),
"adam": ( "adam": (
lib.cadam_8bit_blockwise_grad_fp32, lib.cadam_8bit_blockwise_grad_fp32,
lib.cadam_8bit_blockwise_grad_fp16, lib.cadam_8bit_blockwise_grad_fp16,
lib.cadam_8bit_blockwise_grad_bf16, lib.cadam_8bit_blockwise_grad_bf16,
), ),
"momentum": ( "pagedadam": (
lib.cmomentum_8bit_blockwise_grad_fp32, lib.cpagedadam8bit_blockwise_fp32,
lib.cmomentum_8bit_blockwise_grad_fp16, lib.cpagedadam8bit_blockwise_fp16,
lib.cpagedadam8bit_blockwise_bf16,
), ),
"rmsprop": ( "adamw": (
lib.crmsprop_8bit_blockwise_grad_fp32, lib.cadam_8bit_blockwise_grad_fp32,
lib.crmsprop_8bit_blockwise_grad_fp16, lib.cadam_8bit_blockwise_grad_fp16,
lib.cadam_8bit_blockwise_grad_bf16,
),
"pagedadamw": (
lib.cpagedadam8bit_blockwise_fp32,
lib.cpagedadam8bit_blockwise_fp16,
lib.cpagedadam8bit_blockwise_bf16,
), ),
"lion": ( "lion": (
lib.clion_8bit_blockwise_grad_fp32, lib.clion_8bit_blockwise_grad_fp32,
lib.clion_8bit_blockwise_grad_fp16, lib.clion_8bit_blockwise_grad_fp16,
lib.clion_8bit_blockwise_grad_bf16, lib.clion_8bit_blockwise_grad_bf16,
), ),
"adagrad": ( "momentum": (
lib.cadagrad_8bit_blockwise_grad_fp32, lib.cmomentum_8bit_blockwise_grad_fp32,
lib.cadagrad_8bit_blockwise_grad_fp16, lib.cmomentum_8bit_blockwise_grad_fp16,
),
"rmsprop": (
lib.crmsprop_8bit_blockwise_grad_fp32,
lib.crmsprop_8bit_blockwise_grad_fp16,
), ),
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment