⚡ Apple 실리콘 GPU 가속을 위한 MLX 통합
Lucid 2.0의 가장 큰 변화는 Apple 실리콘을 위한 MLX 기반 Metal GPU 가속을 정식 통합한 것이다. CPU 전용에서 GPU 지원으로 넘어가는 일은 단순히 커널을 옮기는 문제가 아니라, 백엔드 추상화(lucid._backend.core, lucid._backend.metal), 텐서의 디바이스 일관성, lazy evaluation 대응, 학습 루프 관례까지 전면 재검토를 요구했다.
이 글은 MLX라는 라이브러리 소개부터, Lucid 내부에서 GPU 디바이스를 처리하는 구체적인 메커니즘, 그리고 실제 사용 시 주의점까지 테크 리포트 형식으로 정리한다.
🧭 MLX 소개와 GPU 지원의 필요성
MLX는 Apple이 제공하는 Metal 기반 수치 연산 라이브러리로, NumPy와 유사한 API를 가지면서 GPU/ANE 가속을 제공한다. Apple 실리콘의 통합 메모리 아키텍처를 활용해 데이터 이동 비용을 줄이고, lazy execution으로 계산 그래프를 최적화한다. 딥러닝 프레임워크에서 GPU 지원은 다음 이유로 필수적이다.
- 대규모 행렬 연산: 학습 시 대다수 FLOPs가 GEMM/conv에 집중되며, GPU의 병렬성이 필수.
- 실행 시간 단축: 동일 모델을 더 빠르게 학습/추론 → 실험 주기 단축.
- 모바일/온디바이스 최적화: Apple 실리콘은 Mac뿐 아니라 향후 온디바이스 시나리오도 고려 대상.
Lucid는 MLX를 통해 기존 NumPy 코드 경로를 크게 바꾸지 않으면서도 GPU 가속을 선택적으로 사용할 수 있도록 설계했다.
🧱 백엔드 추상화: lucid._backend.core와 metal
- 경로:
lucid/_backend/core.py,lucid/_backend/metal.py - 핵심 역할: device-aware 연산 래핑, CPU/GPU 분기, 디바이스 충돌 검출.
core.py는 operation 추상 클래스와 @func_op 데코레이터 팩토리를 제공하며, metal.py은 MLX(Metal) 환경에서 디바이스 판단, 인덱싱 변환, 가용성 체크를 담당한다. GPU를 쓸지 말지는 연산 호출 시점에 결정되며, Tensor.device 필드와 is_gpu_op 판단을 조합해 올바른 경로를 선택한다.
🧮 operation 클래스와 CPU/GPU 디스패치
class operation(ABC):
@abstractmethod
def cpu(self, *args, **kwargs): ...
@abstractmethod
def gpu(self, *args, **kwargs): ...
def __call__(self, *args, **kwargs):
if is_gpu_op(*args):
return self.gpu(*args, **kwargs)
return self.cpu(*args, **kwargs)- 이중 구현: 각 연산은
cpu/gpu메서드를 구현한다. MLX 가속은gpu경로에 넣고, CPU는 NumPy를 사용. - 런타임 선택: 입력 텐서의
device가 하나라도gpu이면gpu경로를 호출. 모든 입력이cpu면 CPU 경로. - fallback:
@fallback데코레이터로 GPU 미지원 연산을 표시해 CPU 강제 사용 가능.
이 구조 덕분에 동일 연산 클래스가 두 백엔드를 자연스럽게 지원하며, 호출자는 별도 분기 없이 동일 API를 사용한다.
🪄 @func_op 데코레이터의 디바이스 처리
- 경로:
core.func_op
주요 흐름 요약:
- Tensor 정제:
_check_is_tensor(arg, device=device)로 입력을 Tensor로 승격하며 디바이스를 강제한다. - 디바이스 충돌 검증: 이미 점유된 Tensor가 다른 디바이스일 때 GPU 연산을 요구하면 예외를 던진다.
- 자동
.to(device):tensor.is_free이면 연산에 맞춰 자동으로 이동; 아니면 충돌 검사 후 유지. - 결과 텐서 처리:
result.to(device)로 결과를 대상 디바이스에 배치, free 가능하면result.free()호출. - autograd 연결: backward 시
_match_grad_shape(..., device=device)로 장치에 맞는 grad 정렬.
if tensor.is_free:
tensor.to(device)
else:
if tensor.device != device:
raise RuntimeError("...passed for {device} operation...")
...
result.to(device)
if is_free:
result.free()device 인자는 데코레이터 수준에서 지정(예: @binary_func_op(device="gpu"))해 CPU/GPU 변종을 명시적으로 구분한다.
🧲 lucid._backend.metal의 역할: 가용성/인덱싱/판단
- 가용성:
check_metal_availability()가mx.metal.is_available()을 확인, 미지원 시 경고(MetalNotSupportedWarning) 후 CPU로 폴백. - 디바이스 판별:
is_gpu_op/is_cpu_op가 Tensor.device와 MLX array 여부를 보고 연산 경로를 선택. - 인덱싱 변환: MLX는 NumPy식 일부 인덱싱이 다르기 때문에
parse_mlx_indexing이 bool mask/list를 MLX int32 인덱스로 변환. - 예외 메시지: GPU 텐서에 CPU 배열 인덱싱을 시도하면 명확한 에러를 던져 디버깅을 돕는다.
이 레이어는 MLX 특유의 제약(인덱싱, 가용성)을 중앙에서 처리해 연산 구현부를 간결하게 유지한다.
디바이스 판별 코드
def is_cpu_op(*tensor_or_any) -> bool:
for t in tensor_or_any:
device = getattr(t, "device", None)
if device is None:
if isinstance(t, mx.array):
return False
else:
if device == "gpu":
return False
return True
def is_gpu_op(*tensor_or_any) -> bool:
for t in tensor_or_any:
device = getattr(t, "device", None)
if device is None:
if isinstance(t, mx.array):
return True
else:
if device == "gpu":
return True
return Falsedevice속성이 없더라도 MLX array면 GPU로 간주.- 하나라도 GPU면 GPU 연산, 하나라도 GPU면 CPU 연산은 비활성.
인덱싱 파서 코드
def parse_mlx_indexing(index: Any) -> Any:
if isinstance(index, np.ndarray):
raise TypeError("GPU tensors do not support CPU tensor or NumPy array indexing.")
if isinstance(index, tuple):
parsed = []
for i, idx in enumerate(index):
if isinstance(idx, np.ndarray):
raise ValueError(f"NumPy array indexing found at {i}-th index.")
if isinstance(idx, bool):
parsed.append(1 if idx else 0)
elif isinstance(idx, mx.array) and idx.dtype == mx.bool_:
parsed.append(mx.array(np.flatnonzero(idx.tolist()), dtype=mx.int32))
elif isinstance(idx, list) and all(isinstance(i, bool) for i in idx):
mask = mx.array(idx, dtype=mx.bool_)
parsed.append(mx.array(np.flatnonzero(mask.tolist()), dtype=mx.int32))
elif isinstance(idx, list):
parsed.append(mx.array(idx, dtype=mx.int32))
else:
parsed.append(idx)
return tuple(parsed)
elif isinstance(index, bool):
return 1 if index else 0
elif isinstance(index, mx.array) and index.dtype == mx.bool_:
return mx.array(np.flatnonzero(index.tolist()), dtype=mx.int32)
elif isinstance(index, list) and all(isinstance(i, bool) for i in index):
mask = mx.array(index, dtype=mx.bool_)
return mx.array(np.flatnonzero(mask.tolist()), dtype=mx.int32)
elif isinstance(index, list):
return mx.array(index, dtype=mx.int32)
return index| 입력 타입 | 처리/변환 결과 |
|---|---|
np.ndarray | 예외 발생 (GPU 텐서 인덱싱 불가) |
tuple 내부 np.ndarray | 예외 발생 |
bool | 1 또는 0 |
mx.array bool mask | flatnonzero 후 mx.int32 인덱스로 변환 |
list of bool | mask → mx.bool_ → flatnonzero → mx.int32 |
list of int | mx.array(int32) |
| 기타 | 원본 유지 |
요약: GPU 텐서에 NumPy/CPU 배열로 인덱싱하는 실수를 차단하고, MLX가 요구하는 int32 인덱스/플랫 위치로 변환해 호환성을 보장한다.
🔄 텐서 디바이스 전환: .to()와 자동 이동
Tensor.to(device)는 데이터 복사와 메타데이터 업데이트를 동시에 수행한다. 주요 규칙:
- 이동 대상:
data를 NumPy ↔ MLX array로 변환,device필드 업데이트. - 자동 이동:
@func_op내부에서 free 텐서는 연산 대상 디바이스로 자동 이동한다. - 사용자 호출: 모델/파라미터/버퍼를 GPU로 옮길 때
model.to("gpu")가 각 모듈의 파라미터/버퍼를 재귀적으로 이동.
예시:
model = MyNet()
model.to("gpu") # 파라미터/버퍼 모두 GPU로 이동
x = lucid.Tensor(...).to("gpu")
y = model(x) # 모든 연산이 GPU 경로로 흘러간다🆓 Tensor.free()와 디바이스 자유도
Tensor.free()는 디바이스에 구애받지 않도록 표시하는 기능이다. free 텐서는 연산 시 요구되는 디바이스로 자동 이동할 수 있고, 고정 텐서는 충돌 시 예외가 난다.
- is_free=True:
@func_op내부에서 연산 디바이스로 자유롭게.to(device)이동. - is_free=False: 이미 고정된 텐서는 디바이스가 다르면 명시적
.to()없이는 에러.
예시:
a = lucid.Tensor([1,2]) # cpu, is_free=True (기본)
b = lucid.Tensor([3,4]).to("gpu") # gpu
c = lucid.add(a, b) # a가 자동으로 gpu로 이동 후 연산
c.device # "gpu"역으로 GPU free 텐서가 CPU 텐서와 만나면 CPU로 내려온다. 그래프 분리와는 별개이므로 grad 추적은 requires_grad로 제어한다.
🧮 연산 구현에서의 CPU/GPU 분기 패턴
예: lucid/_func/bfunc.py의 add 연산
class add(operation):
@binary_func_op() # CPU
def cpu(self, a, b):
self.result = Tensor(a.data + b.data)
return self.result, self.__grad__
@binary_func_op(device="gpu") # GPU
def gpu(self, a, b):
self.result = Tensor(mx.add(a.data, b.data))
return self.result, self.__grad__- 동일 클래스에 CPU/GPU 메서드를 정의해
__call__이 자동 분기. @binary_func_op(device="gpu")가 GPU 경로임을 명시하고,Tensor생성 시 MLX array를 사용.- grad 계산은 공통(
__grad__)으로 두어 디바이스 독립 로직을 재사용.
이 패턴은 matmul, conv, activation 등 모든 연산에 적용된다.
⚔️ 디바이스 충돌 처리와 에러 전략
- 입력 텐서가 서로 다른 디바이스일 때 GPU 연산을 호출하면
RuntimeError로 즉시 중단. - free 텐서는 자동 이동하지만, 이미 고정된 텐서가 섞이면 명시적
.to()호출을 요구. - 인덱싱 시 CPU 배열을 GPU 텐서에 사용하면
parse_mlx_indexing에서 타입 에러 발생 → 명확한 가이드 제공.
이러한 강한 제약은 애매한 암시적 이동을 방지해 버그를 줄인다.
💤 MLX의 Lazy Evaluation과 loss.eval()의 필수성
MLX는 계산을 지연시키고 필요 시점에 평가한다. 학습 루프에서 loss.backward() 전에 반드시 loss.eval()(또는 동등한 평가 호출)을 수행해 실제 값을 계산해야 한다. 그렇지 않으면 그래프가 예상과 다르게 구성되어 grad가 비거나, 디바이스 전환 타이밍이 어긋날 수 있다. 또한, 명시적 평가 없이 훈련 iteration을 반복하면 계산 그래프가 계속해서 누적되어 심각한 메모리 낭비가 발생할 수 있다.
예시:
opt.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.eval() # MLX lazy → 명시적 평가
loss.backward()
opt.step()Lucid 문서/예제에서도 MLX 경로에서는 이 호출을 강조한다.
🚦 디바이스 혼합 상황: 안전 가드
- CPU 텐서 + GPU 텐서: GPU 연산 호출 시 즉시 에러, CPU 연산 호출 시 GPU 텐서를 CPU로 자동 이동하지 않는다. 사용자가 명시적으로
.to()해야 한다. - 버퍼/파라미터 혼합:
model.to("gpu")를 호출해 파라미터/버퍼를 일괄 이동, 입력도 GPU로 맞추는 것이 권장된다. - 비텐서 인자: 스칼라/리스트는 그대로 통과하지만, NumPy 배열을 GPU 연산에 넘기면
_check_is_tensor가 Tensor로 승격하며 디바이스를 맞춘다.
이 가드는 디바이스 불일치로 인한 미묘한 성능 저하나 잘못된 계산을 예방한다.
🔧 .to 체인: 텐서, 모듈, 옵티마이저
- Tensor.to: 데이터 변환 + device 업데이트.
- Module.to: 파라미터/버퍼를 재귀 이동, 서브모듈 전파.
- Optimizer와의 호환성: Optimizer는 초기화 시점의 파라미터 참조를 그대로 사용하므로, 모델 이동 후 Optimizer를 재생성하거나 파라미터 참조를 재동기화하는 것이 안전하다.
실제 권장 흐름:
model = MyNet().to("gpu")
opt = MyOptimizer(model.parameters(), defaults={"lr": 1e-3})
for epoch in ...:
...모델 이동 → 옵티마이저 생성 순서로 참조 불일치 문제를 방지한다.
🔬 메모리와 성능 고려
- im2col/컨볼루션: GPU 경로에서 MLX 연산을 사용하면 CPU 대비 큰 속도 향상을 기대. 단, 컬럼 버퍼가 GPU 메모리를 사용하므로 batch/커널 설정에 주의.
- grad 누적: MLX grad도 lazy로 축적될 수 있으므로, 필요 없는 중간 텐서는
free()로 그래프에서 분리. - 인덱싱 비용:
parse_mlx_indexing이 bool mask를 int32 인덱스로 변환하는 과정이 추가되므로, 빈번한 고차원 마스킹은 성능에 영향을 줄 수 있다.
성능 측정 시 MLX의 lazy 특성을 고려해 mx.eval() 또는 Lucid의 loss.eval() 등 평가 호출로 타이밍을 명확히 하는 것이 중요하다.
📚 실사용 예제 요약
import lucid
import lucid.nn as nn
from lucid.optim import SGD
from lucid.optim.lr_scheduler import LRScheduler # 예: Cosine 등
model = MyNet().to("gpu")
opt = SGD(model.parameters(), defaults={"lr": 1e-2})
for epoch in range(10):
for x, y in loader:
x = x.to("gpu"); y = y.to("gpu")
opt.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.eval() # MLX lazy → 평가 필수
loss.backward()
opt.step()
# sched.step() # 스케줄러 사용 시여기서 모든 텐서/파라미터/버퍼가 GPU에 올라가 있으므로 연산은 자동으로 gpu 경로를 탄다. loss 평가 호출을 잊지 않는 것이 핵심이다.
🌅 Lucid 발전과 향후 방향성
MLX 통합은 Lucid가 NumPy 전용 연구용 프레임워크에서 Apple 실리콘을 온전히 활용하는 실전 엔진으로 이동했다는 전환점이다. 백엔드 추상화가 정리되면서
- 새로운 디바이스를 추가해도
operation/@func_op패턴만 확장하면 된다. - 모델·옵티마이저·스케줄러는 디바이스 무관한 동일 코드를 유지하고,
.to()호출만으로 환경 전환이 가능하다. - lazy eval 대응과 디바이스 가드가 안정적인 학습 루프를 보장한다.
Lucid는 애초에 개인 공부·교육용 프레임워크라는 목표를 갖고 시작했지만, GPU 가속 지원으로 학습 속도와 실험 주기를 근본적으로 단축할 수 있게 되었다. 이는 “미니 PyTorch”라는 경험적 목표를 한층 현실화하며, 교육용 코드와 실전형 코드의 간극을 줄여준다. 학습자가 동일한 API로 CPU와 GPU를 넘나들며 실험할 수 있다는 점에서, 프레임워크의 학습 가치도 크게 높아졌다.
향후에는 MLX의 추가 최적화(ANE 활용, fused kernel), 다중 디바이스 지원, CPU/GPU 혼합 실행 등에 같은 패턴을 적용해 확장할 수 있다. 이번 통합은 Lucid 2.x 이후 성능·유연성 로드맵의 기반이 된다.
✅ 결론
Lucid의 MLX 통합은 백엔드 추상화(core + metal), 디바이스 일관성 검사(@func_op, operation), lazy evaluation 대응(loss.eval()), 자동 디바이스 이동(.to, free)이라는 네 축을 중심으로 이루어졌다. GPU 가속을 도입하면서도 PyTorch 유사한 사용성을 유지했고, CPU와 GPU 경로를 하나의 연산 클래스 안에서 공존시키는 구조로 확장성을 확보했다. Apple 실리콘에서 Lucid를 사용할 때는 입력/모델/옵티마이저를 GPU로 맞추고, MLX의 lazy 특성에 유의해 명시적 평가를 수행하는 것만 지키면, 기존 코드 변경 없이도 GPU 가속의 이점을 누릴 수 있다.
