MPS Support on Apple Silicon devices #547
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…ilicon This commit adds full support for training AI models on Apple Silicon Macs using MPS, including fixes for multiprocessing tensor sharing issues and UI compatibility. Major changes: - Created toolkit/device_utils.py with comprehensive MPS device management - Fixed PyTorch multiprocessing issues on MPS by forcing num_workers=0 - Added Apple Silicon GPU detection in UI API - Fixed UI client-side errors with type checking and error handling - Added Apple Silicon-optimized training configs for Flux and Z-Image Key fixes: - Resolved "_share_filename_: only available on CPU" error - Fixed .toFixed() errors on string values from MPS API - Added defensive JSON parsing with fallbacks - Updated all training processes to use device-specific dataloader settings
- Update existing job configuration from adamw8bit to adamw for MPS compatibility - Add device-aware optimizer selection in UI (MPS shows compatible optimizers only) - Update default optimizer to use adamw on Mac, adamw8bit on CUDA - Add backend validation to automatically convert 8-bit optimizers for MPS - Fix multiple UI components with .toFixed() errors on string values - Add comprehensive MPS device utilities and detection
Author
|
@jaretburkett you are the boss here, up to you if you want to integrate MPS support or not. For sure from PyTorch 2.9 things has improved a lot. |
Contributor
|
I looked over the changes, I don't see anything that looks like it would cause an issue. I'll run some tests on it to make sure. @ivanfioravanti, how much ram on mac was needed to train z-image-turbo? I have a 24GB macbook I can test on. |
Author
|
Let me try today. Keep you posted. |
Author
|
~30GB I can grant you access to an M3 Ultra 512GB, but next week. |
|
Excellent. What settings are you using? I had to disable transformations to get going using PyTorch 2.7.0. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
MPS Support on Apple Silicon devices. This should solve #127
Many changes had to be applied to make this work properly.