make mentions of mps in docs. ty good people in issue karpathy#28

Leo-tech-1 · Jan 20, 2023 · 1f77d03 · 1f77d03
1 parent a6bffee
commit 1f77d03
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -125,6 +125,8 @@ $ python train.py config/train_shakespeare_char.py --device=cpu --compile=False
 
 Where we decrease the context length to just 64 characters and only use a batch size of 8.
 
+Finally, on Apple Silicon Macbooks you can use device `--device mps` ("Metal Performance Shaders"), which can significantly accelerate training (2-3X). You will need a specific version of PyTorch. See [Issue 28](https://github.com/karpathy/nanoGPT/issues/28).
+
 ## benchmarking
 
 For model benchmarking `bench.py` might be useful. It's identical to what happens in the meat of the training loop of `train.py`, but omits much of the other complexities.

diff --git a/train.py b/train.py
@@ -67,7 +67,7 @@
 # DDP settings
 backend = 'nccl' # 'nccl', 'gloo', etc.
 # system
-device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
+device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
 dtype = 'bfloat16' # 'float32' or 'bfloat16'
 compile = True # use PyTorch 2.0 to compile the model to be faster
 # -----------------------------------------------------------------------------