저는 Karpathy의 llm.c를 Mojo로 포팅하고 Metal 백엔드를 추가하여, PyTorch나 CPython 없이도 Apple Silicon에서 GPT-2 124M을 학습할 수 있도록 했습니다. 이는 Mojo 25.5에서 CPU 전용이었던 dorjeduck의 llm.mojo를 확장한 것입니다. 이 버전은 직접 작성한 CUDA 및 Metal GPU 커널과 함께 Mojo 1.0.0b3 nightly에서 실행됩니다. 제 M4 Max 환경(B=4, T=1024, GPT-2 124M, 공식 실행일 2026-07-13, Cold GPU, 각 arm 사이에 30초의 냉각 시간, 6개 arm 모두를 인터리빙하여 실행)에서의 결과는 다음과 같습니다: 설정별 평균 ms/step 및 tok/s 결과: MLX bf16 406.5 10077 (가장 빠른 arm), MLX fp32 475.7 8610, llm.mojo bf16 503.3 8138 (MPS bf16 대비 1.71배 빠름), llm.mojo fp32 665.2 6157 (MPS fp32 대비 1.25배 빠름), PyTorch MPS fp32 830.8 4930 (기준점), PyTorch MPS bf16 861.8 4753 (기준점). 네, MLX가 승리했습니다. Apple의 자체 프레임워크가 동일한 하네스(harness)에서 벤치마크했을 때 제 bf16 경로보다 1.24배 더 빠릅니다. 이 격차는 거의 전적으로 행렬 곱셈(matmul, 한 스텝의 약 70%)에서 발생합니다. 제가 사용하는 Metal bf16 matmul은 fp32 속도의 약 1.1배로 작동하는 반면, MLX의 bf16은 텐서 코어(tensor cores)를 사용하여 약 2배의 속도를 냅니다. llm.c는 Metal 포트가 없으므로, Apple Silicon에서는 PyTorch MPS와 MLX가 대조군(baseline) 역할을 합니다. 냉각 시간(cooldowns)은 필수적입니다. M4 Max는 지속적인 GPU 부하가 약 8초간 지속되면 스로틀링(throttling)이 발생합니다(MPS 스텝 시간이 몇 스텝 이내에 ~877 ms에서 1500 ms, 2500 ms로 치솟는 것을 확인했습니다). make benchmark-metal 명령어로 재현할 수 있으며, 냉각 시간이 포함된 상태로 6개 arm을 한 번에 실행합니다. 처음 작동했던 Metal 포트는 MPS보다 약 4.1배 느렸으나(3627 ms/step), 최종 bf16 수치는 그 시작점보다 7.2배 더 빠릅니다. 격차의 대부분은 Metal 특유의 문제였습니다. threadgroup 포인터를 일반 주소 공간(generic address space)으로 캐스팅하면 장치 메모리(device memory)를 조용히 읽게 되며(모든 값이 0인 점수에 대해 softmax를 적용하여 균일한 가중치를 생성함), NVIDIA에 맞춰 튜닝된 스칼라 Flash-Attention 커널은 Apple GPU에서 Peak FLOP의 1% 미만으로 작동했기 때문에, GEMM 분해 어텐션(GEMM-decomposed attention)이 810배 더 빨랐습니다. 정확성 검증: make test를 통해 PyTorch와 비교하여 16개의 그래디언트 텐서(gradient tensors) 및 10단계 손실 궤적(loss trajectory)을 확인하며, 235개의 테스트로 구성된 등가성 스위트(equivalence suite)를 실행합니다.

HuggingFace에는 HellaSwag에서 29.53%를 기록한 학습된 124M FineWeb 체크포인트(ulmentflam/gpt2-124m-fineweb-mojo)가 있으며, 이는 Karpathy의 llm.c 재현 결과인 29.9%와 통계적으로 구별할 수 없는 수준입니다. AI에 관하여 말씀드리자면, 모든 커널(kernel)과 트레이너(trainer) 코드는 수작업으로 작성되었습니다. LSP(Language Server Protocol)나 LLM(Large Language Model)의 도움을 받지 않았습니다. 그것이 이 프로젝트의 원래 목적이었습니다. 테스트와 이후의 최적화 작업에는 AI의 도움을 받았으며, 모델별 명시와 사용된 경우의 완전한 공개(출처 표기 포함)를 포함하여 리포지토리(repo)에 모두 공개되어 있습니다. Mojo 자체에 대해서는, 컴파일러가 하드웨어 간 이식성(portability)에 대한 힘든 작업을 대신해 줄 것이라고 가정하고 시작했습니다. 하지만 그렇지 않았습니다. 결국 벤더별로 장치 특화 로직(device-specific logic)을 분기해야 했고, CUDA와 비교했을 때 예상보다 더 많은 상용구 코드(boilerplate)가 필요했습니다 (Karpathy의 커널은 훨씬 더 압축적입니다). Modular 엔지니어가 그들의 포럼에서 이 포트(port)를 검토했으며, 그들의 답변은 장치별 특화(per-device specialization)는 예상된 결과이며, 그들의 전략은 컴파일러의 마법이 아닌 라이브러리 중심의 구조화된 커널(TileTensor)이라는 것입니다. CPU는 정반대의 이야기입니다. 아주 적은 노력만으로도 llm.c의 20스레드 OpenMP 경로보다 4.0배 더 빨랐습니다. 한계점: GPT-2 124M 모델만 지원하며, 발표된 GPT-3 결과는 없습니다 (트레이너에 설정값은 존재하며, 혼합 정밀도(mixed precision)는 NVIDIA에서 FP8 및 NVFP4까지 지원합니다). 툴체인(toolchain)은 Mojo 1.0 beta nightly 버전이므로 변경 사항이 잦을 수 있습니다. Multi-GPU ZeRO (stage 0~3)는 월드 사이즈(world size) 2 및 8에서 단일 GPU와 등가성 검증을 거쳤으나, 해당 실행은 NVIDIA 기준이며, Apple Silicon은 여기에서 단일 GPU로 동작합니다. GB10에서 bf16은 llm.c CUDA와 대등한 수준(0.999x)이며, fp32는 현재 약간 앞서 있습니다(1.07x, TF32 vs TF32). 리포지토리: https://github.com/ulmentflam/llm.mojo /u/ulmentflam에 의해 제출됨 [link] [comments]