8BitAdamW and bitsandbytes.functional.create_dynamic_map

[ArrowM]

Hello! 8BitAdamW has stopped working for Dreambooth since 0.36.0 (loss NaNs very quickly). I took a stab at reviewing the differences between the releases. The if-block around functional.py:218 in 2f2063b may have been accidentally indented. tests.test_functional.test_few_bit_quant doesn't cover the inside of the if block, so maybe this was missed (or it was intentional and it'd be great to have confirmation). I'm hoping one of the maintainers can take a quick look. Thanks!

[feffy380]

I wonder if this is why I get loss=nan with the rocm port (which does include that commit). Have you tried un-indenting the block in question? It's in the python code so it shouldn't require recompilation.

Edit: nope, at least for the rocm port it's not the cause as I'm still getting nan loss after a certain point. It definitely seems to be tied to bitsandbytes though, as switching to the lion scheduler resolved the issue

[ArrowM]

Oh wow, you're right - I can test it pretty easily. It seemed to fix it! I'll make a PR.

[rationalism]

This should be fixed by 0.41.1

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.