Skip to content

Experimental: allow fp16 in mps #961

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Oct 29, 2022
Merged

Experimental: allow fp16 in mps #961

Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
cd221bc
Docs: refer to pre-RC version of PyTorch 1.13.0.
pcuenca Oct 17, 2022
c3960ed
Remove temporary workaround for unavailable op.
pcuenca Oct 17, 2022
8674cf7
Update comment to make it less ambiguous.
pcuenca Oct 17, 2022
762650f
Remove use of contiguous in mps.
pcuenca Oct 17, 2022
f55fcf5
Merge remote-tracking branch 'origin/main' into m1-torch-1.13
pcuenca Oct 17, 2022
de5258f
Merge remote-tracking branch 'origin/main' into m1-torch-1.13
pcuenca Oct 20, 2022
78dbfea
Merge remote-tracking branch 'origin/main' into m1-torch-1.13
pcuenca Oct 22, 2022
ef41ddd
Special case: use einsum for much better performance in mps
pcuenca Oct 24, 2022
8d47740
Update mps docs.
pcuenca Oct 24, 2022
2f8fe8e
MPS: make pipeline work in half precision.
pcuenca Oct 24, 2022
ff9553c
Merge remote-tracking branch 'origin/main' into mps-allow-fp16
pcuenca Oct 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/diffusers/models/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,12 @@ def __init__(self, dim_in: int, dim_out: int):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out * 2)

def gelu(self, gate):
if gate.device.type != "mps":
return F.gelu(gate)
# mps: gelu is not implemented for float16
return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)

def forward(self, hidden_states):
hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
return hidden_states * F.gelu(gate)
return hidden_states * self.gelu(gate)