사용법 (Usage)

use luminal::prelude::*;
// 계산 그래프 (compute graph) 생성
let mut cx = Graph::new();
...

시작하기 (Getting Started)

Llama 3 8B

다음은 CUDA 환경에서 Luminal을 사용하여 Llama 3 8B를 로컬에서 실행하는 간단한 예시입니다:

cd ./examples/llama
cargo run --release

특징 (Features)

속도 (Speed)

Luminal은 H100에서 Q8 Llama 3 8B를 이론적 최대 성능의 약 80% 수준으로 실행할 수 있습니다. 목표는 어떤 장치에서든 어떤 모델이라도 가장 빠르게 실행할 수 있는 ML 프레임워크가 되는 것입니다.

단순함 (Simplicity)

Luminal의 핵심은 항상 최소한의 상태를 유지할 것입니다. 전체 핵심 라이브러리를 단 하루 만에 이해할 수 있어야 합니다.

PyTorch 네이티브 (PyTorch-native)

Luminal은 컴파일러 백엔드 (compiler backend)로서 PyTorch와 직접 통합됩니다. 단순히 torch.compile(model, backend=luminal_cuda)를 호출하여 PyTorch 모델을 컴파일하면 됩니다. 또한 Rust에서 사용할 수 있는 훌륭한 텐서 API (tensor API)를 제공합니다.

RISC 스타일 아키텍처 (RISC-style architecture)

Luminal의 모든 것은 15개의 기본 연산 (primitive ops)으로 요약됩니다:

• 단항 연산 (Unary) - Log2, Exp2, Sin, Sqrt, Recip
• 이항 연산 (Binary) - Add, Mul, Mod, LessThan
• 기타 (Other) - SumReduce, MaxReduce, Iota, Gather, Scatter, Cast

이 연산들은 트랜스포머 (transformers), 컨볼루션 신경망 (convnets), 그리고 전 세계의 거의 모든 인기 있는 모델을 지원하기에 충분합니다.

검색 (Search)

가장 좋은 휴리스틱 (heuristic)은 휴리스틱을 사용하지 않는 것입니다. Luminal은 컴파일러가 복잡한 최적화를 발견할 수 있는 유연성을 가질 수 있도록 가능한 모든 결정을 검색 (search)하려고 시도합니다. 이를 통해 수동으로 작성된 연산이나 휴리스틱에 의존하지 않고도 Flash Attention 및 이와 유사하게 복잡한 최적화들을 자동으로 발견할 수 있습니다. 또한, 이를 통해 우리는 미래에도 매우 작고 단순한 상태를 유지하면서 훨씬 더 큰 프레임워크들의 성능을 능가할 수 있습니다.

네이티브 (Native)

현재의 머신러닝 (ML) 생태계는 너무 파편화되어 있으며, 그 해결책은 또 다른 추상화 계층 (layer of abstraction)을 만드는 것이 아닙니다. Luminal은 Rust로 작성되었으며, 가속기 API (CUDA, Metal 등)와 직접 상호작용합니다. 간접 참조 (indirection)나 추상화, 호환성 계층 (compatibility layers), Docker 컨테이너, 또는 가상 환경 (virtual environments)이 필요 없습니다. 오직 정적 링크된 (statically-linked) Rust 크레이트 (crate) 하나면 충분합니다.

PyTorch를 통한 검증

정확성 (correctness)은 중요합니다. 우리는 모든 연산 (ops)을 커버하고 이들이 동일한 PyTorch 구현과 동일하게 작동하는지 확인하기 위해 가능한 한 많은 테스트를 작성합니다. (개선 필요!)

이데올로기 (Ideology)

왜 다른 딥러닝 (DL) 라이브러리와 이렇게 달라 보이나요?

대부분의 딥러닝 라이브러리는 즉시 실행 (eager-first) 방식입니다. 즉, 각 연산 (op) 호출이 데이터에 직접 작동함을 의미합니다. PyTorch에서 x + y를 볼 때, 덧셈은 실제로 그 자리에서 즉시 일어납니다. 이는 대부분의 개발자가 기대하는 방식과 정확히 일치하기 때문에 디버깅에 매우 유용합니다.

하지만 이는 성능 측면에서 좋지 않습니다. 개발자에게 합리적인 방식이 기계에게도 잘 작동하는 것은 아니며, 이는 아무도 어셈블리 (Assembly) 코드를 직접 손으로 작성하지 않는 것과 같은 이치입니다. 대부분의 라이브러리는 연산자 융합 (Operator Fusion)이나 JIT 컴파일 (JIT compilation)을 덧붙여 컴파일 흐름을 기계에 더 유리한 방식으로 바꾸려 시도함으로써 이 문제를 해결하려고 합니다. 하지만 이는 PyTorch에게조차 매우 어렵고 심지어 힘든 일임이 드러났습니다!

XLA는 어떤가요?

XLA, torch.compile, TVM 및 기타 전통적인 컴파일러 스택은 복잡성 폭발 (Complexity explosion) 문제를 겪습니다. 이들은 그래프를 고수준 표현 (High-level representation)에서 저수준 기계어 (Low-level machine code)로 낮추고 최적화하기 위해, 매우 방대한 양의 파괴적 (Destructive, 단방향) 재작성 규칙 (Rewrite rules)들로 구성되어 있습니다. 하지만 이러한 규칙들은 파괴적이기 때문에, 성능 이점이 확실할 때만 실행되도록 강제됩니다. 이로 인해 규칙들은 매우 복잡해지고, 특수 사례 (Special-cased)가 많아지며, 그 수가 방대해집니다. 여기에 추가적인 하드웨어 백엔드 (Hardware backends), 모델 아키텍처 (Model architectures), 그리고 새로운 데이터 타입 (Dtypes)이 투입되면, 이들은 복잡성의 무게로 인해 고통받게 되며 종종 매우 최적화되지 않은 코드를 생성합니다. 결국 성능을 다시 확보하기 위해 Pallas나 Triton과 같은 DSL (Domain-Specific Languages)이 필요하게 됩니다.

모든 것을 컴파일하세요

Luminal의 핵심 원칙은 사전 컴파일 (ahead-of-time compilation)입니다. 가능한 한 모든 것을 컴파일 시간 (compile time)으로 밀어내고, 실행 시간 (run time)에는 아무것도 남기지 마세요. Luminal은 XLA 및 tinygrad와 더 유사한 접근 방식을 취합니다. 여기서는 모든 것이 정적 (static)입니다. x + y와 같은 식을 작성할 때 실제 계산은 일어나지 않습니다. 해당 연산은 나중에 실행하기 위해 유향 비순환 계산 그래프 (directed acyclic computation graph)에 기록됩니다. 오직 graph.execute()가 실행될 때만 계산이 수행됩니다. 하지만 이것은 그저 지연 실행 (lazy execution) 아닌가요? 맞습니다! 하지만 Luminal에서는 모든 것이 이런 방식으로 이루어집니다. 모든 신경망 (neural networks)은 정적 계산 그래프로 구축되고, 컴파일된 후 나중에 실행됩니다.

일급 동적 특성 (First-class dynamism)

완전 정적 (fully-static)인 세상이라면 좋겠지만, 우리는 필연적인 동적 특성 (dynamism)이 존재하는 세상에 살고 있습니다. 따라서 우리는 동적 형상 (dynamic shapes)을 기호적 차원 (symbolic dimensions)으로서 네이티브하게 모델링합니다. Luminal은 복잡한 식을 포함한 임의의 기호적 차원을 지원하여 (s, 4096), (b, h, w + 3) 등과 같은 형상 (shapes)을 제공합니다. 이러한 풍부한 표현 방식은 컴파일러가 형상에 대한 완전한 가시성을 갖게 하며, 여전히 공격적인 특수화 (specialization)를 수행할 수 있게 합니다.

하지만 왜 그런가요?

이로 인한 결과로, 실제로 실행되는 실제 계산은 작성된 코드와 근본적으로 다를 수 있습니다. 전체 신경망이 계산 그래프 내에 완전히 표현되어 있기 때문에, Luminal은 전역적 지식 (global knowledge)을 가집니다. 이는 머신러닝 (ML)의 대부분의 복잡성을 컴파일러로 밀어낼 수 있음을 의미합니다. 예를 들어, 디바이스 (devices), 데이터 타입 (datatypes), 심지어 자동 미분 (autograd)까지 사전 모델링되어 컴파일러에 의해 최적화됩니다!

이제 우리는 다음과 같은 것들을 할 수 있습니다:

• 공격적인 커널 퓨전 (kernel fusion)
• 런타임에 컴파일되는 형상 특화 커널 (Shape-specific kernels)
• 저정밀도 데이터 타입 (mxfp4, nvfp4, fp8 등)
• 사전에 탐색된 복잡한 멀티 디바이스 병렬성 토폴로지 (multi-device parallelism topologies)
• 네트워크를 범용 코드로 작성하되, 초특화된 아키텍처에서 빠르게 컴파일 및 실행

현재 단계는 어디인가요?

• 네이티브 PyTorch 지원
• 검색 공간(search space) 내 다수의 커널 라이브러리 지원 (FlashInfer, cuBLASLt 등)
• examples/ 디렉토리 내 Rust 텐서 API로 구현된 다수의 모델
• Transformer를 포함하여 luminal_nn에 구현된 소규모 신경망 (NN) 모듈 라이브러리
• hl_ops에 구현된 상당량의 고수준 연산 (high-level ops). 우리는 PyTorch API 중 가장 많이 사용되는 약 80%를 맞추는 것을 목표로 하고 있습니다.

로드맵 상의 몇 가지 사항:

• TMA 및 tcgen.05와 같은 스레드(thread) 및 워프(warp) 수준의 인트린직 (intrinsics)을 지원하는 더 세밀한 방언 (dialects)
• ROCm 백엔드
• 더 많은 공개 추론 가속기 (inference accelerator) 백엔드 (매우 곧 출시 예정...)
• 공개 벤치마킹 스위트 (benchmarking suite)
• 자동 검색된 모델 병렬화 (model parallelism) (TP, PP, EPS, EPR, SP 등)
• 양자 광학 레트로 인캡슐레이터 (quantum photonic retro encabulator)를 위한 컴파일러 작성
• 다이슨 스웜 (Dyson swarm) 구축

라이선스 (License)

Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) 또는 MIT 라이선스 (http://opensource.org/licenses/MIT) 중 선택하여 라이선스가 부여됩니다. 이 파일은 해당 약관에 따르지 않고는 복제, 수정 또는 배포할 수 없습니다.