PyTorch 는 GPU 가 별도의 대륙인 세상에서 태어났어. NVIDIA 카드는 자체 VRAM 을 가지고, 시스템 RAM 과 물리적으로 분리돼서 PCIe 로 연결돼. GPU 를 만지는 모든 op 은 시민권을 선언해야 해 — tensor.to('cuda') — 그리고 메모리 복사라는 통관세를 내야 해.
Apple Silicon 엔 그 국경이 없어. M 시리즈 칩은 CPU, GPU, Neural Engine 을 한 조각의 실리콘에 통합하고, 모두가 같은 물리 RAM 풀에 손을 뻗어. 별도 VRAM 없음. 프로세서 사이의 PCIe 버스 없음. 그냥 메모리가 있고, 어느 compute unit 이 다음에 거기서 읽을지의 질문만 있어.
MLX 는 그 물리 현실을 위해 박힌 array library 야. MLX 의 마케팅에서 자랑하는 "unified memory" feature 는 소프트웨어 트릭이 아냐 — framework 가 하드웨어가 아닌 척하기를 거부한 것일 뿐.
코드에서 그게 어떻게 보이나
두 스니펫을 옆에 놓으면 대비가 가장 선명해. CUDA 위의 PyTorch 는 일련의 통관 신고서야:
(이 PyTorch 스니펫은 대비용 — mlx env 에서 돌리지 마. 그냥 의식만 읽어.)
그리고 여기 MLX 의 동등한 코드. 뭐가 없는지 봐.
그게 변화 전부야. .to(device) 없음. 돌아오는 길의 .cpu() 없음. 관리할 비동기 복사 스트림 없음. Array 의 저장소가 GPU 와 CPU 가 공유하는 unified memory 그 자체이고, framework 는 그냥 맞는 compute unit 에서 맞는 kernel 을 돌려.
증명 — 복사 없는 cross-device 산술
이게 내가 이걸 완전히 내 것으로 만든 실험이야. Array 두 개를 만드는데, 하나는 CPU 에 선언, 하나는 GPU 에 선언, 그리고 둘을 더해. CUDA 위의 PyTorch 에선 이건 에러 (먼저 .to() 로 하나를 옮겨야 함). MLX 에선 그냥 돼:
MLX 의 두 "device" 는 별도 메모리 풀이 아냐. 다음 op 을 어느 compute unit 이 실행할지의 라벨이야. 데이터는 시작부터 끝까지 한 곳에 있어.
왜 이게 design center 지 feature 가 아닌가
MLX 의 다른 선택 대부분이 이 한 사실에서 흘러나와. Lazy evaluation 이 존재하는 이유 — framework 가 개별 호출 지점이 아니라 전체 그래프 기준으로 어느 compute unit 이 각 op 을 돌릴지 결정할 수 있어야 하니까. 그 결정을 똑똑하게 하려면 그래프를 먼저 손에 잡아야 해. mx.grad, mx.vmap 같은 function transform 도 더 쉬워 — 스케줄할 transfer 비용이 없으니까. Training loop 가 compute 와 복사를 interleave 할 필요 없어 — 그냥 돌아.
이 quest 전체에서 한 가지 mental image 만 가져간다면, 이거 가져가 — Apple Silicon 에서 array 는 모든 compute unit 이 읽을 수 있는 단일 byte 풀에 살아. MLX 는 그 사실을 진지하게 받아들이는 framework 야.
Code
CUDA 위의 PyTorch — 통관 의식 (대비용으로 읽기, mlx env 에서 돌리지 마)·python
# Don't run this in the mlx env — it's a contrast, not a recipe.
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
x = torch.randn(1024, 1024).to(device) # ← copy bytes from RAM to VRAM
y = torch.randn(1024, 1024).to(device) # ← copy bytes from RAM to VRAM
z = (x @ y).cpu() # ← copy result back from VRAM to RAM
# Forget any one of these .to() / .cpu() calls and you get a runtime error
# or a silent slowdown.
MLX 동등한 코드 — 뭐가 없는지 봐·python
import mlx.core as mx
x = mx.random.normal((1024, 1024))
y = mx.random.normal((1024, 1024))
z = x @ y
mx.eval(z) # materialize the lazy graph (more on this in core.lesson4)
print("shape :", z.shape)
print("dtype :", z.dtype)
print("device:", mx.default_device()) # → Device(gpu, 0) on Apple Silicon
# No .to('cuda'). No .cpu(). The unified-memory pool is the storage,
# always, and 'device' is just a label about which compute unit ran it.
Cross-device add — 'device' 가 풀이 아니라 라벨이라는 증명·python
import mlx.core as mx
mx.set_default_device(mx.cpu)
a_cpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device()) # Device(cpu, 0)
print("a_cpu :", a_cpu) # array([1, 2, 3, 4], dtype=float32)
mx.set_default_device(mx.gpu)
a_gpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device()) # Device(gpu, 0)
print("a_gpu :", a_gpu) # array([1, 2, 3, 4], dtype=float32)
# In PyTorch on CUDA, this next line is a runtime error (cross-device op).
# In MLX, it just works. There is one pool of bytes; the labels were a hint
# about scheduling, not a partition of memory.
b = a_cpu + a_gpu
mx.eval(b)
print("a_cpu + a_gpu :", b) # array([2, 4, 6, 8], dtype=float32)
mlx env 에서 cross-device add 코드 블록 돌려. array([2, 4, 6, 8]) 돌아오는 거 확인해. 그 다음 한 줄 바꿔 — a_cpu 를 전처럼 먼저 선언하되, 바로 b = a_cpu + a_cpu (양쪽 같은 device) 하고 b 의 device 봐. 그 다음 mx.set_default_device(mx.gpu) 하고 a_cpu + a_cpu 다시 더해. 결과의 뭐가 바뀌고 뭐가 안 바뀌었어? 두 문장. 운동의 포인트는 라벨이 어디서 중요하고 어디서 안 중요한지 느끼는 것.
Progress
Progress is local-only — sign in to sync across devices.