tinygrad/tinygrad

tinygrad – 작고 수정 가능한 딥러닝 스택

소개 – tinygrad는 완전한 엔드투엔드 딥러닝 프레임워크를 구현한 오픈소스 Python 라이브러리이지만, 코드 베이스는 의도적으로 작게 유지됩니다. 다음을 제공합니다:

자동 미분 (autograd) 기능을 갖춘 Tensor 클래스.
연산을 융합하고 장치 커널 (device kernels)로 낮추는 IR / 컴파일러.
JIT / 그래프 실행 (TinyJit 사용).
실제 모델을 학습할 수 있는 경량 nn / optim / dataset 레이어.

PyTorch와 유사한 사용 편의성을 제공하면서도 컴파일러 내부 구조를 공개하여 쉽게 조정할 수 있도록 함으로써, 기능이 풍부한 PyTorch와 극도로 미니멀한 micrograd 사이의 위치를 지향합니다.

주요 프레임워크와의 비교

기능	PyTorch	JAX	tinygrad
Eager Tensor API	✅	–	✅
Autograd	✅	✅ (via JAXPR)	✅
함수형 변환 (vmap, pmap)	–	✅	부분 지원
컴파일러 및 IR 가시성	–	✅	✅
프론트엔드 (nn, optim, datasets)	✅	–	✅
멀티 디바이스 지원	✅	✅	✅ (CPU, CUDA, Metal, OpenCL, AMD, NV, QCOM, WebGPU)

퀵스타트 예제

from tinygrad import Tensor, nn, Context

class LinearNet:
    def __init__(self):
        self.l1 = Tensor.kaiming_uniform(784, 128)
        self.l2 = Tensor.kaiming_uniform(128, 10)
    def __call__(self, x: Tensor) -> Tensor:
        return x.flatten(1).dot(self.l1).relu().dot(self.l2)

model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])
with Context(TRAINING=1):
    for i in range(10):
        optim.zero_grad()
        loss = model(x).sparse_categorical_crossentropy(y).backward()
        optim.step()
        print(i, loss.item())

저장소에는 GPU에서 약 5초 만에 약 98%의 정확도에 도달하는 완전한 MNIST 예제 (examples/beautiful_mnist.py)가 포함되어 있습니다.

장치 지원 (가속기)

tinygrad는 장치당 약 25개의 기본 연산만으로 많은 백엔드에서 실행될 수 있습니다:

CPU
CUDA (NVIDIA GPU)
Metal (Apple GPU)
OpenCL
AMD, NV, QCOM (특수 런타임)
WebGPU (브라우저 / WASM) python3 -c "from tinygrad import Device; print(Device.DEFAULT)"를 실행하여 기본적으로 사용되는 가속기를 확인하세요.

설치

# 소스에서 설치 (권장)
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

# 또는 저장소에서 직접 설치
python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

이 프로젝트는 테스트 스위트를 실행하기 위한 .[testing] 옵션도 제공합니다.

문서 및 커뮤니티

문서 및 퀵스타트: https://docs.tinygrad.org/ (docs/ 폴더에서 생성됨)
Discord: https://discord.gg/ZjZadyC7PK – 사용자와 개발자를 위한 활발한 채팅방.
기여: README에는 엄격한 PR 정책(코드 골프 금지, 속도 관련 주장에 대한 벤치마크, 회귀 테스트 포함 등)이 명시되어 있으며, 채택된 개선 사항에 대해서는 현금 보상이 언급되기도 합니다.

tinygrad를 사용해야 할 때

학습 및 실험 – 컴파일러 파이프라인이 공개되어 있어 고수준 텐서 연산이 어떻게 저수준 커널로 변환되는지 정확히 확인할 수 있습니다.
제한된 하드웨어에서의 프로토타이핑 – 라이브러리가 가볍고 CPU, GPU, 심지어 WebGPU에서도 실행됩니다.
수정 가능한 스택이 필요한 연구 – 몇 개의 파일만 수정하여 새로운 연산을 추가하거나 스케줄링을 변경할 수 있습니다.

거대한 생태계, 프로덕션 수준의 안정성 또는 JAX 변환 기능의 전체 세트가 필요한 경우 PyTorch나 JAX가 더 나은 선택입니다.

요약

tinygrad는 텐서, 자동 미분, JIT 기반 컴파일러, 멀티 디바이스 실행을 작고 읽기 쉬운 코드 베이스로 구현한 실제 기능성 딥러닝 프레임워크입니다. 거대한 라이브러리의 오버헤드 없이 현대적인 ML 스택의 내부 구조를 이해하거나 확장하려는 학생, 취미 개발자 및 연구자에게 이상적입니다.