Skip to content

Latest commit

 

History

History
70 lines (55 loc) · 2.31 KB

README.md

File metadata and controls

70 lines (55 loc) · 2.31 KB
Raw

pytorch-darwin-env

PyTorch development environment with Metal acceleration for Darwin

Description

This project provides a Nix-based development environment for PyTorch with Metal acceleration support on macOS (Darwin) systems with Apple Silicon. It includes a flake configuration for easy setup and a verification script to check PyTorch installation and benchmark Metal acceleration performance.

This repository exists as a workaround for the following issue in the Nixpkgs repository:

Package request: pytorch on darwin with GPU (MPS) support #243868

The current PyTorch package in Nixpkgs for aarch64-darwin does not support GPU acceleration via Metal Performance Shaders (MPS). This project provides a solution by using the PyTorch-provided wheel that includes MPS support, allowing users to leverage GPU acceleration on Apple Silicon Macs.

Features

  • Nix-based development environment
  • PyTorch 2.0.1 with Metal acceleration support
  • Performance benchmark script comparing CPU and MPS (Metal Performance Shaders) execution

Requirements

  • macOS on Apple Silicon (M1/M2) machine
  • Nix package manager with flakes support

Usage

Using the flake directly (without cloning)

You can use this environment directly without cloning the repository:

nix develop "github:david-r-cox/pytorch-darwin-env"

Or to run the verification script directly:

nix run "github:david-r-cox/pytorch-darwin-env#verificationScript"

Local setup

  1. Clone this repository:
git clone https://github.com/david-r-cox/pytorch-darwin-env.git
cd pytorch-darwin-env
  1. Enter the Nix development environment:
nix develop
  1. Run the verification and benchmark script:
verify-pytorch-metal

This script will:

  • Check the PyTorch version
  • Verify MPS (Metal Performance Shaders) availability
  • Run a matrix multiplication benchmark on both CPU and MPS
  • Display the performance comparison between CPU and MPS
$ verify-pytorch-metal # M1 Max
PyTorch version: 2.0.1
MPS available: True
CPU time: 1.3895 seconds
MPS time: 0.0061 seconds
MPS speedup: 228.17x

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.