Skip to content
This repository was archived by the owner on Sep 10, 2025. It is now read-only.

update torchao pin: optimized mps lowbit shaders #1428

Merged
merged 2 commits into from
Dec 18, 2024
Merged

update torchao pin: optimized mps lowbit shaders #1428

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5ffa1b6
update torchao pin: optimized shaders
manuelcandales Dec 18, 2024
3c0e898
Merge branch 'main' into torchao-mps-opt
manuelcandales Dec 18, 2024
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
2 changes: 1 addition & 1 deletion install/.pins/torchao-pin.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
7d7c14e898eca3fe66138d2a9445755a9270b800
2e032c6b0de960dee554dcb08126ace718b14c6d
2 changes: 2 additions & 0 deletions torchchat/utils/quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -932,13 +932,15 @@ def quantized_model(self) -> nn.Module:
libs = glob.glob(f"{torchao_build_path}/cmake-out/lib/libtorchao_ops_aten.*")
libs = list(filter(lambda l: (l.endswith("so") or l.endswith("dylib")), libs))
torch.ops.load_library(libs[0])
print("Loaded torchao cpu ops.")
except Exception as e:
print("Unabled to load torchao cpu ops library. Slow fallback kernels will be used.")

try:
libname = "libtorchao_ops_mps_aten.dylib"
libpath = f"{torchao_build_path}/cmake-out/lib/{libname}"
torch.ops.load_library(libpath)
print("Loaded torchao mps ops.")
except Exception as e:
print("Unabled to load torchao mps ops library.")

Expand Down
Loading