Why does adamw_bnb_8bit skip updating embedding parameters?
i’m using whisper for testing, but this probably applies to other models since in whisper’s source code there is no specific instruction to elicit this behavior.when using adamw_bnb_8bit as the optimizer for seq2seq tasks, i noticed it automatically turns off (at least according to the output below which pops up right at the beginning of the training) embedding parameters.is this a quirk/property of bitsandbytes package? i haven’t read the dettmers et al.'s work, so not sure if this is expected behavior.@sanchit-gandhiany thoughts?(for whisper-small)skipped Embedding(1500, 768): 1.0986328125M params
skipped Embedding(51865, 768, padding_idx=50257): 39.085693359375M params
skipped WhisperPositionalEmbedding(448, 768): 39.413818359375M params
skipped: 39.413818359375M params
2023-02-10T14:59:44Z
Hey@ozanciga! Very interesting observation! My understanding is that embedding layers can become particularly unstable when downcast to lower precisions (fp16/fp8). It might be that there are exception clauses for embedding modules in the bitsandbytes package which prevents them from being downcast.
2023-02-10T15:00:14Z
Indeed! It looks like we need a special embedding layer for 8bit embeddings:GitHub - TimDettmers/bitsandbytes: 8-bit CUDA functions for PyTorchSee step 3: Replace embedding layer if necessary:torch.nn.Embedding(..) -> bnb.nn.Embedding(..)
2023-02-10T19:22:21Z
thank you@sanchit-gandhi, that seems about right. maybe a question for the hugging face devs, should this be handled on the backend automatically?also, do you have any opinion on if freezing embeddings have any significant impact on the outcome? specifically asking for whisper but also interested in general since it’s very desirable to use 8bit in most cases.
2023-02-12T03:36:31Z
actually i went through the source and turns out “skipped” means the optimizer is using 32bits for those parameters. i monkey patched the code to incorporatebnb.nn.stableembeddings, but i doubt it’s worth it for most cases.github.comhuggingface/transformers/blob/67d074874d285e616393c65a0e670088e1b6b74a/src/transformers/trainer.py#L1084else:self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)if optimizer_cls.__name__ == "Adam8bit":import bitsandbytesmanager = bitsandbytes.optim.GlobalOptimManager.get_instance()skipped = 0for module in opt_model.modules():if isinstance(module, nn.Embedding):skipped += sum(dict((p.data_ptr(), p.numel()) for p in module.parameters()).values())print(f"skipped {module}: {skipped/2**20}M params")manager.register_module_override(module, "weight", {"optim_bits": 32})logger.debug(f"bitsandbytes: will optimize {module} in fp32")print(f"skipped: {skipped/2**20}M params")if is_sagemaker_mp_enabled():self.optimizer = smp.DistributedOptimizer(self.optimizer)return self.optimizer
2023-02-17T14:05:55Z
Hey@ozanciga! Sorry for the late reply here. Indeed, it could be worth handling it automatically by the trainer. Feel free to open an issue on thetransformersrepoif you want to discuss how this might look! Or directly a PR if you’ve got an idea on how to fix it already. Think this would make for a nice PR to the repo! Happy to help you any questions around the issue/PR!
2023-02-17T14:08:32Z
ozanciga:do you have any opinion on if freezing embeddings have any significant impact on the outcome?My intuition would be that freezing the embeddings won’t have a significant impact on the outcome. I believe in the Dalle-Mini project freezing the pre-trained embeddings actually gave superior performance vs non-frozen. Probably best to experiment here with say 1 epoch of data and see how eval WER performance compares with frozen / non-frozen embeddings
2023-04-29T09:18:28Z
Now I ques it is automatically handled if trainer is correctly initialized in the main branch?