PyTorch Training on iOS Devices

Question

Created Jan ’23

Replies 1

Boosts 0

Views 2.2k

Participants 2

Is there any way to train PyTorch models directly on-device on an iPhone via the GPU? PyTorch Mobile docs seems to be completely focused on inference only, as do the iOS app examples (https://github.com/pytorch/ios-demo-app). I did find this article about using MPS backend on Macs (https://developer.apple.com/metal/pytorch/), but not sure if this is at all viable for iOS devices. There's also this prototype article about using iOS GPU for PyTorch mobile (https://pytorch.org/tutorials/prototype/ios_gpu_workflow.html), but it too seems to be focused on inference only.

We are attempting to train a large language model on the iPhone 14 and we would like to a) discard intermediate activations and recompute them, and b) manage memory directly to write some intermediate activations to the filesystem and later read them back. We suspect that converting a PyTorch model to CoreML format and using CoreML for training would prevent us from making these low-level modifications, but PyTorch might have the APIs necessary for this. If there's any examples/pointers that anyone can link to that would be great.

Boost

Answer 1

aerdm OP

Feb ’23

Yes, it is possible to train PyTorch models directly on an iPhone via the GPU. PyTorch provides support for Metal, which is Apple's low-level, low-overhead graphics API. You can use the PyTorch Metal API to write custom metal shader functions that can be executed on the GPU, allowing you to perform computationally intensive tasks, such as training a deep learning model.

You can check Metal here: https://developer.apple.com/metal/

However, training a deep learning model on an iPhone can be challenging due to the limited computational resources available on the device. Additionally, it may also be difficult to manage the data transfer to and from the device. It's generally recommended to perform model training on a more powerful machine and then deploy the trained model to the iPhone for inference.

1