{"slug": "two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu", "title": "Two PyTorch ops that stopped a CUDA text model on an Apple GPU", "summary": "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.", "body_md": "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`.\n\nThe 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.\n\n`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.\n\nA short helper keeps the NVIDIA path and adds the Apple GPU when CUDA is absent:\n\n``` php\ndef accelerator() -> torch.device:\n    if torch.cuda.is_available():\n        return torch.device(\"cuda\")\n    if torch.backends.mps.is_available():\n        return torch.device(\"mps\")\n    return torch.device(\"cpu\")\n```\n\nAn 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.\n\nXSum 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:\n\n```\nTypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.\n```\n\nThe 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.\n\n``` php\ndef as_float64(t: torch.Tensor) -> torch.Tensor:\n    if t.device.type == \"mps\":\n        t = t.detach().to(\"cpu\")\n    return t.double()\n```\n\nThe 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.\n\nGeneration then died inside the residual vector quantizer:\n\n```\nRuntimeError: Adaptive pool MPS: input sizes must be divisible by output sizes.\nNon-divisible input sizes are not implemented on MPS device yet.\n```\n\n`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:\n\n``` python\ndef _interpolate(x, size, mode):\n    out_len = size[-1]\n    if x.device.type == \"mps\" and mode == \"area\" and x.shape[-1] % out_len != 0:\n        return F.interpolate(x.cpu(), size=size, mode=mode).to(\n            device=x.device, dtype=x.dtype\n        )\n    return F.interpolate(x, size=size, mode=mode)\n```\n\nThis 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.\n\nFlexAttention 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.\n\nI 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.\n\nOn this machine `torch.mps.recommended_max_memory()` returned 17.76 GiB.\n\nTwo easy Sudoku boards (`--n 2`) scored 2/2 solved, 100% cell accuracy, and 100% valid grids.\n\nOne 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.\n\nThose 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.\n\nThe 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:\n\n```\nPYTORCH_MPS_HIGH_WATERMARK_RATIO=0.8 PYTORCH_MPS_LOW_WATERMARK_RATIO=0.8 python evaluate.py ...\n```\n\nI 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.\n\nI 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.\n\nI would run a full split before quoting quality numbers. Two solved Sudoku boards and one summary only prove the graph executes on Metal.\n\nNext-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).\n\nThis article was drafted with AI assistance.", "url": "https://wpnews.pro/news/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu", "canonical_source": "https://dev.to/leonewton253/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu-4jb5", "published_at": "2026-10-28 04:09:23+00:00", "updated_at": "2026-10-28 04:11:14.602224+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "ai-infrastructure", "developer-tools"], "entities": ["PyTorch", "Apple", "Metal", "MPS", "Next-Concept Modeling", "HazyResearch", "Kiarina", "GitHub"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu", "markdown": "https://wpnews.pro/news/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu.md", "text": "https://wpnews.pro/news/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu.txt", "jsonld": "https://wpnews.pro/news/two-pytorch-ops-that-stopped-a-cuda-text-model-on-an-apple-gpu.jsonld"}}