invoke-ai
[bug]: I don't thin bfloat16 config option is being used.
Is there an existing issue for this problem?
- [X] I have searched the existing issues
Operating system
macOS
GPU vendor
Apple Silicon (MPS)
GPU model
M3 base 10 GPU revision
GPU VRAM
24
Version number
v3.7.0
Browser
Safari 17.2.1
Python dependencies
{ "accelerate": "0.27.2", "compel": "2.0.2", "cuda": null, "diffusers": "0.26.3", "numpy": "1.26.4", "opencv": "4.9.0.80", "onnx": "1.15.0", "pillow": "10.2.0", "python": "3.10.13", "torch": "2.3.0.dev20240221", "torchvision": "0.18.0.dev20240221", "transformers": "4.37.2", "xformers": null }
What happened
Recent PyTorch Nightlies have added some bfloat16 for MPS, and testing Diffusers with them showed there enough support for Stable Diffusion to run and give a small but statically significant decrease in seconds per iteration.
I thats to see if I could use the nightlies with InvokeAI now the basicSR dependancy has been removed and that testing worked fine. So I set the config to bfloat16.
Everything worked but I saw no change to render times.
Digging though the code and I spotted a few bits of code that look like they force the use of float16., including one that I think prevents the use of bfloat16 on all formats.
I made the following code changes
invokeai/backend/util/devices.py
made choose_precision allow use of bfloat16 for MPS
def choose_precision(device: torch.device) -> str:
"""Returns an appropriate precision for the given torch device"""
if device.type == "cuda":
device_name = torch.cuda.get_device_name(device)
if not ("GeForce GTX 1660" in device_name or "GeForce GTX 1650" in device_name):
if config.precision == "bfloat16":
return "bfloat16"
else:
return "float16"
elif device.type == "mps":
if config.precision == "bfloat16":
return "bfloat16"
else:
return "float16"
return "float32"
This broke Compel so had to do this quick fix, this needs a better fix as I think the issue is purely a MPS torch issue.
invokeai/app/invocations/compel.py
comment out both occurrences of dtype_for_device_getter, line 125 and 248
dtype_for_device_getter=torch_dtype,
And finally the one I think prevents bfloat16 usage on all platforms
invokeai/app/services/model_manager/model_manager_default.py line 65
dtype = torch.float32 if precision == "float32" elif precision == "bfloat16" else torch.float16
replaced with
if precision == "float32":
dtype = torch.float32
elif precision == "bfloat16":
dtype = torch.bfloat16
else:
dtype = torch.float16
and that seems to do the trick and I see a small speed up I'd expect.
I suspect there are other places, where similar changes need to be made as the same
invokeai/app/invocations/latent.py
loads the vae as either float16 or float32, not bfloat16. Not sure if that is actually necessary.
Model installation
invokeai/app/services/model_install/model_install_default.py
_guess_variant() looks like it will try and get the 32bit version of models for bfloat16 instead of the fp16 variant by default (andI've noticed one or two bf16 variants recently)
Model loading
invokeai/backend/model_management/models/base.py I think line 285 means for bfloat16 it will load the fp32 model in preference to the fp16 variant if both are installed.
What you expected to happen
Modifying the config file to use bfloat16 should be enough to use bfloat16 precision.
How to reproduce the problem
Changer the invoke.yaml file to use bfloat16, you may need some debug code to conform its actually still using float16, and to run a Diffusers script to test exactly how bfloat16 should behave with your hardware.
Additional context
No response
Discord username
No response
Adding another bfloat16 support issue
[2024-03-09 12:06:28,029]::[InvokeAI]::ERROR --> Traceback (most recent call last):
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/services/invocation_processor/invocation_processor_default.py", line 134, in __process
outputs = invocation.invoke_internal(
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/invocations/baseinvocation.py", line 669, in invoke_internal
output = self.invoke(context)
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/invocations/latent.py", line 773, in invoke
) = pipeline.latents_from_embeddings(
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/backend/stable_diffusion/diffusers_pipeline.py", line 381, in latents_from_embeddings
latents, attention_map_saver = self.generate_latents_from_embeddings(
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/backend/stable_diffusion/diffusers_pipeline.py", line 441, in generate_latents_from_embeddings
callback(
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/invocations/latent.py", line 704, in step_callback
self.dispatch_progress(context, source_node_id, state, self.unet.unet.base_model)
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/invocations/latent.py", line 318, in dispatch_progress
stable_diffusion_step_callback(
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/util/step_callback.py", line 85, in stable_diffusion_step_callback
image = sample_to_lowres_estimated_image(sample, sdxl_latent_rgb_factors, sdxl_smooth_matrix)
File "/Volumes/SSD2TB/AI/InvokeAI/lib/python3.10/site-packages/invokeai/app/util/step_callback.py", line 13, in sample_to_lowres_estimated_image
latent_image = samples[0].permute(1, 2, 0) @ latent_rgb_factors
RuntimeError: Undefined type BFloat16
[2024-03-09 12:06:28,029]::[InvokeAI]::ERROR --> Error while invoking:
Undefined type BFloat16
EDIT: raised a PyTorch issue https://github.com/pytorch/pytorch/issues/121583
EDIT2: Now fixed in PyTorch nightlies, I think, the current nightly totally breaks InvokeAI by breaking Accelerate. The test code that uses the same size tensors works.
This seems to be much improved in version 4, I only had to change invokeai/backend/util/devices.py to get bfloat16 working with MPS which suggests it should just work with CUDA.