Training on M1 "MPS"

[okpatil4u]

Most of the people do not have access to 8XA100 40GB systems. But a single M1 Max laptop with 64 GB memory could host the training. How difficult is it to port this code to "MPS" ?

[okpatil4u]

I take it back. Seems like these are 8 x 40 GB systems.

There is a good paper on cramming [Cramming: Training a Language Model on a Single GPU in One Day]
https://arxiv.org/abs/2212.14034

I thought some work on these lines was done here as well.

[karpathy]

Actually I think this issue is great to keep open, in case anyone investigates nanoGPT in mps context. I haven't tried yet.

[okpatil4u]

What is the actual memory requirement ? Will Mac Studio with 128 GB RAM be sufficient for training ?

[jwkirchenbauer]

Refining the above comment slightly, do you currently have any (rough is fine) estimates on the relative sizes of the memory footprint for the just the model parameters, params plus the forward activations as a fn of bsz, versus the backward graph as a fn of bsz, on the 8xA100 40gb configuration? Where does it peak across the server during training?

That might start to inform some people on how to go about laying this out on the resources they have.

[NightMachinery]

Also relevant for inference.

[personsg]

I haven't had a chance to do any benchmarking yet but training starts just fine on M1 Ultra with --device=mps.

[acheong08]

I use Google Colab with the smaller model
https://github.com/acheong08/nanoGPT/tree/Colab - https://colab.research.google.com/github/acheong08/nanoGPT/blob/Colab/nanoGPT.ipynb

[itakafu]

I tried out "i only have a MacBook" from README but with --device="mps" and it seems to run faster. With CPU, one iteration is roughly about 100ms whereas with mps is about ~40ms. My machine is a base line Mac Studio.

[okpatil4u]

That's for training a very small transformer. My machine is 64 gb RAM, M1 Max. For bert-medium like architecture, this is how it goes.

Overriding: dataset = shakespeare
Overriding: n_layer = 8
Overriding: n_head = 512
Overriding: n_embd = 512
Overriding: device = mps
Overriding: compile = False
Overriding: eval_iters = 1
Overriding: block_size = 64
Overriding: batch_size = 128

Initializing a new model from scratch
number of parameters: 50.98M
step 0: train loss 10.9816, val loss 10.9783
iter 0: loss 10.9711, time 4613.50ms
iter 1: loss 10.9673, time 5791.48ms
iter 2: loss 10.9647, time 7842.40ms
iter 3: loss 10.9646, time 10196.35ms
iter 4: loss 10.9604, time 11602.34ms
iter 5: loss 10.9495, time 9393.25ms
iter 6: loss 10.9615, time 10373.34ms

[karpathy]

@itakafu thank you for reporting, i'll add mentions of mps to the readme&code.

[SiyuanHuang95]

test on MacBook Air M2, without charger:

with mps: roughly 150~200ms for one iteration
without mps: roughly 450 ~ 500ms for one iteration

just for one reference

@itakafu thank you for reporting, i'll add mentions of mps to the readme&code.

[tomeck]

Confirmed works great w/device='mps'. But make sure to install this version of pytorch:

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

I'm getting <40ms

Thank you SO MUCH for this