# Explicitly disable bf16 for some layers

**URL:** <https://discuss.huggingface.co/t/explicitly-disable-bf16-for-some-layers/159578>\
**Category:** 🤗Transformers\
**Created:** [June 16, 2025, 7:36pm UTC](https://discuss.huggingface.co/t/explicitly-disable-bf16-for-some-layers/159578 "2025-06-16T19:36:21Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![BoltzmachineQ](https://avatars.discourse-cdn.com/v4/letter/b/e47c2d/32.png) [@BoltzmachineQ](https://discuss.huggingface.co/u/BoltzmachineQ)\
**Post date:** [June 16, 2025, 7:36pm UTC](https://discuss.huggingface.co/t/explicitly-disable-bf16-for-some-layers/159578/1 "2025-06-16T19:36:21Z")

</div>

I am using huggingface’s Trainer with `--bfl16` flag enabled and deepspeed enabled. However, I want to force float32 for a specific layer. How to do it?

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [June 17, 2025, 3:18am UTC](https://discuss.huggingface.co/t/explicitly-disable-bf16-for-some-layers/159578/2 "2025-06-17T03:18:59Z")

</div>

Hmm… Mixed precision training?

> **[Performance and Scalability: How To Fit a Bigger Model and Train It Faster](https://huggingface.co/docs/transformers/v4.18.0/en/performance)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> <https://github.com/huggingface/transformers/issues/28921>
>
> Normally when saving a model using the Trainer class, the \`dtype\` of the saved m…odel is (and should be) the same as the original model. This is also true when using mixed precision, and when using DeepSpeed. However, when using mixed precision \*\*together\*\* with DeepSpeed, the output is in float16 no matter the model input \`dtype\`.
> 
> The Trainer class has custom handling for DeepSpeed, depending on the ZeRO stage:
> https://github.com/huggingface/transformers/blob/5f9685576149fb45a61d0dcec9a260930df0a49a/src/transformers/trainer.py#L2914-L2928
> 
> as well does accelerate:
> https://github.com/huggingface/accelerate/blob/06b138d84537ffb2d1d404f2f198a0446e8d7ec3/src/accelerate/accelerator.py#L3042-L3056
> 
> For ZeRO stage \<=2 DeepSpeed holds the model weights in the \`state\_dict\`. Using mixed precision training, these are always in float16. Using full precision training, they are the same dtype as the original model.
> 
> For ZeRO stage 3 the \`state\_dict\` contains just placeholders since the model weights are partitioned. By setting \`stage3\_gather\_16bit\_weights\_on\_model\_save=true\`, DeepSpeed consolidates the weights. When training using mixed precision, float16 is always produced. When training in full precision, despite the name, it follows the dtype of the original model. If \`stage3\_gather\_16bit\_weights\_on\_model\_save=false\`, Trainer saves a full checkpoint instead, and the DeepSpeed \`zero\_to\_fp32.py\` script can be used to recover weights in float32.
> 
> 
> Currently, the only way to save a model that is trained using the Trainer class that applies mixed precision along with DeepSpeed ZeRO stage \<=2 in float32, is to manually save a checkpoint and then use some weight recovery method afterwards. Is this due to a limitation of the DeepSpeed API, or could this be handled in the Trainer class (or preferably in Accelerate)? At least, maybe a flag could be available to either save the float16 weights or a checkpoint at the end of training (kind of how stage 3 with \`stage3\_gather\_16bit\_weights\_on\_model\_save=true\` is handled)?
> 
> 
> \### Who can help
> @pacman100, @muellerzr

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [June 17, 2025, 4:47am UTC](https://discuss.huggingface.co/t/explicitly-disable-bf16-for-some-layers/159578/3 "2025-06-17T04:47:33Z")

</div>

Or similar to this issue…?

> [@Why do you need to re-upcast the norm layers of HF falcon to fb32?](https://discuss.huggingface.co/t/why-do-you-need-to-re-upcast-the-norm-layers-of-hf-falcon-to-fb32/46139):
>
> I saw these lines: bnb\_4bit\_compute\_dtype=torch.float16, ... optim = "paged\_adamw\_32bit" ... for name, module in trainer.model.named\_modules(): if "norm" in name: module = module.to(torch.float32) in the falcon tutorial. These are confusing me. Based on the original QLoRA tutorial, they use 4 bit model + during training they use 16 brain float (not normal float nor float 32). See equation 5: YBF16 = XBF16doubleDequant(c\_FP32\_1, c\_k-bit\_2, WNF4) + XBF16L\_BF16\_1 L\_BF16\_2 doubleDequa…
