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.
