pytorch-apple-silicon-benchmarks

Benchmarks of PyTorch on Apple Silicon.

This is a work in progress, if there is a dataset or model you would like to add just open an issue or a PR.

Prepare environment

Create conda env with python compiled for osx-arm64 and activate it with:

CONDA_SUBDIR=osx-arm64 conda create -n native python -c conda-forge
conda activate native

and install pytorch nightly build with:

pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

and finally install datasets and transformers with:

pip install transformers datasets

Devices

• M1 Max CPU 32GB: 10 cores, 2 efficient + 8 performance up to ~3GHz; Peak measured power consuption: 30W.
• M1 Max GPU 32GB: 32 cores; Peak measured power consuption: 46W.
• NVIDIA V100 16GB (SXM2): 5,120 CUDA cores + 640 tensor cores; Peak measured power consuption: 310W.
• Tesla T4 (using Google Colab Pro): Runtime settings: GPU & High RAM

Results

BERT Transformers in Sequence Classification.

Run the experiments by yourself with:

python tests/transformers_sequence_classification.py \
    --device <cpu|cuda|mps> \
    --pre_trained_name <bert-base-cased|bert-large-cased> \
    --batch_size <32|64|128> \
    --mode <training|inference> \
    --steps 100 \
    --sequence_length <128|512>

The following tables show the time needed to complete 100 steps without gradient accumulation. - means that the script went out of memory. All experiments have been run with float32.

bert-base-cased

• M1 Max 32-core (32GB) tested with 1.13.0.dev20220601
• M1 Ultra 48-core (64GB) tested with 1.13.0.dev20220604
• M2 Ultra GPU 60-core (64GB) tested with 2.2.0.dev20231113
• M3 Pro GPU 14-core (18GB) tested with 2.2.0.dev20231113

Training:

Inference:

Considerations

• This is the first alpha ever to support the M1 family of processors, so you should expect performance to increase further in the next months since many optimizations will be added to the MPS backed.
• At the moment I experienced a progressive slowdown with MPS such that the first iteration took more than half the time than the last. (seems solved in latest release)
• Before deciding whether the M1 Max could be your best choice, consider that it has no float64 support and neither fp16 tensor cores.
• It seems that there is no real limit to the batch size with the M1 Max because it is able to use the swap also for the 'GPU' memory. However, this really slows down training. Just check how time increases linearly with batch size until the swap is used, after which run time grows exponentially.

FAQ

• If you cannot install tokenizers because rust is missing, do the following:

  brew install rustup
  rustup-init
  source ~/.cargo/env