Two PyTorch ops that stopped a CUDA text model on an Apple GPU

A developer ported the Next-Concept Modeling evaluation scripts to Apple Silicon by replacing three CUDA-only assumptions with MPS-aware device selection, moving float64 penalty math to the CPU, and routing non-divisible area interpolation to the CPU when Metal rejects the shape. The model needed no rewrite: released configs already disable FlexAttention in favor of scaled dot-product attention, which the developer kept after citing M1 Max measurements showing SDPA roughly 19x faster than FlexAttention at 8,192-token causal attention. Weights stayed in float32 to keep sampling close to the CUDA run.

The Next-Concept Modeling https://github.com/HazyResearch/ncm eval scripts would not start on this Mac. PyTorch 2.14.1 reported the Apple GPU as available, and evaluate.py still exited with no GPU visible; pass --gpus . The model itself did not need a rewrite. The released configs already turn FlexAttention off and use scaled dot-product attention, linear layers, embeddings, and RoPE. Three call sites assumed CUDA. Two of those assumptions are real Metal limits, and the third is just device plumbing. evaluate.py asked torch.cuda.device count and then set CUDA VISIBLE DEVICES for each shard. On a Mac that count is zero, so the process stopped before it loaded a checkpoint. A short helper keeps the NVIDIA path and adds the Apple GPU when CUDA is absent: php def accelerator - torch.device: if torch.cuda.is available : return torch.device "cuda" if torch.backends.mps.is available : return torch.device "mps" return torch.device "cpu" An Apple Silicon Mac has one GPU, so the launcher starts one process and does not set CUDA VISIBLE DEVICES . Checkpoint loads and .cuda calls in the shard and the scorer use that same device. GPT-2, which scores generative perplexity after generation finishes, moves there too. XSum and LM1B apply a frequency penalty and an adjacency penalty before the argmax. That loop did logits.double on the model device. On MPS that raises: TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead. The loop is a Python argmax over a short sequence, so the numerics can stay in float64 on the CPU. The transformer forward stays on the GPU. php def as float64 t: torch.Tensor - torch.Tensor: if t.device.type == "mps": t = t.detach .to "cpu" return t.double The unigram tables used by that penalty are loaded on the CPU as well. On CUDA they stay on the GPU, in float64, which is what the original script did. Generation then died inside the residual vector quantizer: RuntimeError: Adaptive pool MPS: input sizes must be divisible by output sizes. Non-divisible input sizes are not implemented on MPS device yet. F.interpolate ..., mode="area" on a 1D tensor is adaptive average pooling. The Sudoku tokenizer's levels end at 89 and step through 50, 67, and the other lengths in between. 89 is not a multiple of 50, so Metal refuses the op. PyTorch's own error says to run that case on the CPU. The tensor is a handful of codebook vectors, so the copy is cheap and the attention stays on the GPU: python def interpolate x, size, mode : out len = size -1 if x.device.type == "mps" and mode == "area" and x.shape -1 % out len = 0: return F.interpolate x.cpu , size=size, mode=mode .to device=x.device, dtype=x.dtype return F.interpolate x, size=size, mode=mode This is the same ATen op the CUDA path uses, run on the CPU only when Metal cannot take the shape. PyTorch issue 96056 https://github.com/pytorch/pytorch/issues/96056 is the tracking bug the error points at. FlexAttention has had Metal kernels since PyTorch 2.13. Kiarina's M1 Max measurements https://dev.to/kiarina/testing-pytorch-213-mps-flexattention-on-m1-max-up-to-783x-faster-for-sparse-attention-1p7d are the reason I did not turn it on. A 32,768-token window of 256 was 7.83× faster than SDPA. Causal attention at 8,192 tokens went the other way: SDPA was about 19× faster. NCM's levels are a few hundred tokens, the released configs set use flex attention: false , and the MPS FlexAttention path Kiarina tested is forward-only. SDPA is the kernel that matches this model. I also left the weights in float32. The penalty is sensitive to ties, and I wanted the first samples close to the CUDA run before trying fp16. On this machine torch.mps.recommended max memory returned 17.76 GiB. Two easy Sudoku boards --n 2 scored 2/2 solved, 100% cell accuracy, and 100% valid grids. One XSum document produced a real summary "The number of beds made available for homeless people living in accommodation in Wales has risen in last year..." and the merge step scored it with GPT-2 on the same GPU: ROUGE-1 27.91, ROUGE-2 4.88, ROUGE-L 23.26, generative perplexity 28.34, 26 words. That is one document. It is evidence the scorer ran, not a comparison to the paper. Those two CPU copies are small. A large offload is a different story. Nefayran measured a 4 GiB MPS tensor moved with .to "cpu" https://dev.to/nefayran/pytorch-on-a-mac-why-memory-looks-freed-after-tocpu-1apl : the Activity Monitor footprint dropped to about 0.1 GiB until the CPU actually read the copy, and then it was 4.2 GiB again. On Apple Silicon the CPU and GPU share one pool, so parking weights on the CPU does not hand memory back to the GPU. The same series shows why a big eval can get slow without ever raising out-of-memory. On a 48 GB Mac the default MPS high-watermark sits above physical RAM https://dev.to/nefayran/mps-backend-out-of-memory-on-a-mac-what-the-numbers-mean-and-how-to-set-the-watermark-43dp , so the job swaps first. Setting the cap takes both ratios, before the first MPS tensor: PYTORCH MPS HIGH WATERMARK RATIO=0.8 PYTORCH MPS LOW WATERMARK RATIO=0.8 python evaluate.py ... I would watch swap on LM1B-large. That generator is 36 layers at width 1280, and the scorer is GPT-2 large in the next process. If swap climbs, drop batch size in the eval config. Generation and scoring do not overlap, which already helps. I would not port the model to MLX to get a first result. The weights are a PyTorch state dict , and the math already ran once the device and the two unsupported ops were handled. I would run a full split before quoting quality numbers. Two solved Sudoku boards and one summary only prove the graph executes on Metal. Next-Concept Modeling is from Hermann Kumbong, Ishan S. Khare, Iddah Mlauzi, Andrew Woen, Charles Ding, E. Kelly Buchanan, Simran Arora, Azalia Mirhoseini, and Christopher Ré, arXiv:2610.14162 https://arxiv.org/abs/2610.14162 . The code above is a local change to the MIT-licensed repo https://github.com/HazyResearch/ncm . This article was drafted with AI assistance.