luminal

光速深度学习。

Luminal是一个深度学习库，它使用可组合的编译器来实现高性能。

use luminal::prelude::*;

// Setup graph and tensors
let mut cx = Graph::new();
let a = cx.tensor::<R2<3, 1>>()
    .set([[1.0], [2.0], [3.0]]);
let b = cx.tensor::<R2<1, 4>>()
    .set([[1.0, 2.0, 3.0, 4.0]]);

// Do math...
let mut c = a.matmul(b).retrieve();

// Compile and run graph
cx.compile(<(GenericCompiler, CPUCompiler)>::default(), &mut c);
cx.execute();

// Get result
println!("Result: {:?}", c);

入门指南

Mistral 7B

cd ./examples/mistral
# Download the model
bash ./setup/setup.sh
# Run the model
cargo run --release --features metal    # MacOS (Recommended)
cargo run --release --features cuda     # Nvidia
cargo run --release                     # CPU

功能

速度

Luminal可以在M系列Macbooks上以每秒15-25个token的速度运行Q8 Mistral 7B。目标是成为任何设备上任何模型的最快的ML框架。

简洁性

luminal的核心始终是最小的。应该可以在一个下午内理解整个核心库。

RISC风格架构

luminal中的所有内容都归结为11个原始操作

• 一元 - Log2, Exp2, Sin, Sqrt, Recip
• 二元 - Add, Mul, Mod, LessThan
• 其他 - SumReduce, MaxReduce, Contiguous

这些操作足以支持transformers、convnets等。

本地化

当前的ML生态系统过于碎片化，解决方案不是另一层抽象。Luminal是用rust编写的，直接与CUDA / Metal API交互。没有间接或抽象、docker容器或虚拟环境。只是一个静态链接的rust crate。

经过Pytorch验证

正确性很重要。因此，我们编写了尽可能多的测试来覆盖所有操作，并验证它们与等效的Pytorch实现相同。（需要改进！）

意识形态

为什么这与其他DL库看起来如此不同？

大多数深度学习库都是急切的，这意味着每个操作调用直接作用于数据。在PyTorch中，当你看到 x + y 时，实际上就在那里执行了加法。这对于调试来说很棒，因为它正好符合大多数开发者的预期。

然而，这对性能来说并不是很好。对开发者来说有意义的事情，对机器来说并不好，就像没有人会手动编写汇编语言一样。大多数库试图通过添加操作融合或即时编译来解决这个问题，试图改变编译流程，使其更适合机器。结果发现，这对Pytorch来说是非常困难 甚至 更难 进行！

编译一切

Luminal的一个核心原则是编译时编译。尽可能将所有内容推送到编译时间，而不要留到运行时。Luminal采取的方法更类似于XLA和tinygrad。在这里一切都是静态的。当你编写一个像x + y这样的表达式时，实际上并没有发生计算。操作被记录到一个有向无环计算图中，以便稍后执行。只有在运行graph.execute()之后，计算才会发生。但这不就是延迟执行吗？是的，它是！但在Luminal中，所有事情都是这种方式完成的。所有神经网络都作为单个或少数静态计算图构建，然后编译和执行。

但是为什么？

这种做法的结果是，实际运行的计算可能根本不同于编写的代码。因为我们有一个完整的神经网络，完全表示在计算图中，我们的编译器具有全局知识。这意味着我们可以将大多数机器学习复杂性推给编译器。例如，设备、数据类型和执行计划都由编译器处理。甚至自动微分也会由编译器处理！

现在我们可以这样做

• 积极的内核融合
• 在运行时编译特定形状的内核
• 设备和数据类型通过编译器处理（只需运行CUDA编译器将图转换为使用CUDA内核，然后运行fp16编译器将转换为半精度内核）
• 网络可以用通用代码编写，但可以在超特定架构上快速编译和运行（尝试编写一个同时适用于TF32数据类型和TPU的PyTorch网络；准备好进入if语句的地狱...

编译时形状检查

所有操作都在编译时进行形状检查，因此不再有形状不匹配！这归功于dfdx。

查看图

一旦你编写了所有的计算代码，运行cx.display()以查看整个计算图的所有内容。看起来非常混乱！现在运行cx.compile(GenericCompiler::default())并再次显示图。看起来好多了。

我们在哪里？

• 支持在Mac上运行模型，在Nvidia GPU上运行CUDA，支持全精度和半精度。
• 在M系列Mac上使用LLM的性能与llama.cpp（一个高度优化的库）相比，相差20%以内。
• Mistral 7B和Llama 7B在examples/中实现。有关运行说明，请参阅上面的说明。
• 我们在nn中有一个小的NN模块库，包括Transformer。
• 在hl_ops中实现了大量的高级操作。我们致力于匹配PyTorch API中最常用的约80%。
• 0.3的目标是在M1 Pro上实现SOTA性能（50 tok/s），在单个Nvidia GPU上接近SOTA（>100 tok/s），并支持许多主流模型（Whisper、Stable Diffusion、Yolo v9等）

路线图上的某些事项

• 优化CUDA和Metal矩阵乘法内核
• 细粒度金属和CUDA IR
• 构建基准测试套件以测试其他库
• 自动微分引擎
• 分布式数据、管道和张量并行。
• 在LLM训练中击败PT 2.0性能
• 编写量子光子反封装器编译器
• 构建戴森 swarm

许可证

根据您的选择，许可协议为Apache License，版本2.0 https://apache.ac.cn/licenses/LICENSE2.0 或MIT许可证 http://opensource.org/licenses/MIT。此文件不得复制、修改或分发，除非符合这些条款。