NameError: name 'str2optimizer8bit_blockwise' is not defined

[rexroth0619]

Describe the bug

Trying to run migrate colab scripts to runpod, followed along then ran into this error (at Run Training section)

Reproduction

Running colab ipynb scripts on Jupyter Lab. First it says num_processes undefined, I specified num_processes=1 in Run Training section, then it returns this error

Logs

--------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_10686/1781923440.py in <module>
      1 #@title Run training
      2 import accelerate
----> 3 accelerate.notebook_launcher(training_function, num_processes=1, args=(text_encoder, vae, unet))
      4 with torch.no_grad():
      5     torch.cuda.empty_cache()

/opt/conda/lib/python3.7/site-packages/accelerate/launchers.py in notebook_launcher(function, args, num_processes, use_fp16, mixed_precision, use_port)
    132                 print("Launching training on CPU.")
    133             with patch_environment(use_mps_device=use_mps_device):
--> 134                 function(*args)
    135 
    136 

/tmp/ipykernel_10686/1900933013.py in training_function(text_encoder, vae, unet)
    134                 if accelerator.sync_gradients:
    135                     accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
--> 136                 optimizer.step()
    137                 optimizer.zero_grad()
    138 

/opt/conda/lib/python3.7/site-packages/accelerate/optimizer.py in step(self, closure)
    138                 self._is_overflow = scale_after < scale_before
    139             else:
--> 140                 self.optimizer.step(closure)
    141 
    142     def _switch_parameters(self, parameters_map):

/opt/conda/lib/python3.7/site-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)
    107                 profile_name = "Optimizer.step#{}.step".format(obj.__class__.__name__)
    108                 with torch.autograd.profiler.record_function(profile_name):
--> 109                     return func(*args, **kwargs)
    110             return wrapper
    111 

/opt/conda/lib/python3.7/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.clone():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

/opt/conda/lib/python3.7/site-packages/bitsandbytes/optim/optimizer.py in step(self, closure)
    263                     self.init_state(group, p, gindex, pindex)
    264 
--> 265                 self.update_step(group, p, gindex, pindex)
    266 
    267         return loss

/opt/conda/lib/python3.7/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.clone():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

/opt/conda/lib/python3.7/site-packages/bitsandbytes/optim/optimizer.py in update_step(self, group, p, gindex, pindex)
    521                 config["weight_decay"],
    522                 gnorm_scale=gnorm_scale,
--> 523                 skip_zeros=config["skip_zeros"],
    524             )
    525 

/opt/conda/lib/python3.7/site-packages/bitsandbytes/functional.py in optimizer_update_8bit_blockwise(optimizer_name, g, p, state1, state2, beta1, beta2, eps, step, lr, qmap1, qmap2, absmax1, absmax2, weight_decay, gnorm_scale, skip_zeros)
    856 
    857     if g.dtype == torch.float32 and state1.dtype == torch.uint8:
--> 858         str2optimizer8bit_blockwise[optimizer_name][0](
    859             get_ptr(p),
    860             get_ptr(g),

NameError: name 'str2optimizer8bit_blockwise' is not defined

System Info

ubuntu 20.04 on runpod docker

[patrickvonplaten]

Hey @rexroth0619,

Could you please copy-paste a reproducible code snippet and also the full environment info?
I assume you're trying to train stable diffusion dreambooth using the bitsandbytes optimizer?

[patrickvonplaten]

cc @patil-suraj here

[patil-suraj]

Thanks for the issue , could you please post your command, or the trainig argument you used, so I could reproduce this ?

[rexroth0619]

Hi team, thanks for coming in.

I ended up figuring it out myself. Since I'm running a docker container on runpod, bitsandbytes weren't able to find CUDA under the normal installed directory /usr, since it's installed in the base conda environment (under /opt/conda/lib). By directing it to that folder the code was able to run.

[patil-suraj]

Glad you were able to fix this!

[pcuenca]

Awesome @rexroth0619! Can you please close the issue if this is fully resolved? :)

[sandys]

Hi team, thanks for coming in.

I ended up figuring it out myself. Since I'm running a docker container on runpod, bitsandbytes weren't able to find CUDA under the normal installed directory /usr, since it's installed in the base conda environment (under /opt/conda/lib). By directing it to that folder the code was able to run.

@rexroth0619 could you explain how you did this ? what were the lines of code. im hitting the exact same issue.

[sandys]

fixed this. in case anyone else has an issue - they can refer to my notebook
https://github.com/sandys/SimpleDiffuserDreambooth/blob/main/DreamBooth_Stable_Diffusion.ipynb

works on vast.ai and runpod

[adammenges]

fixed this. in case anyone else has an issue - they can refer to my notebook https://github.com/sandys/SimpleDiffuserDreambooth/blob/main/DreamBooth_Stable_Diffusion.ipynb

works on vast.ai and runpod

thank you @sandys !