smelte-rs

什么是 smelte-rs ?

Smelt 是一个专注于推理的 ML 库，具有尽可能多的优化，同时仍具有可读性和易于使用的小依赖。

限制不安全的使用，并仅用于性能。

运行模型

尝试在文本分类示例中运行 Bert。

# Download the model + tokenizer + config
# This is a clone of https://hugging-face.cn/ProsusAI/finbert with safetensors support.
curl https://hugging-face.cn/Narsil/finbert/resolve/main/model.safetensors -o model-Narsil-finbert.safetensors -L
curl https://hugging-face.cn/Narsil/finbert/resolve/main/tokenizer.json -o tokenizer-Narsil-finbert.json -L
curl https://hugging-face.cn/Narsil/finbert/resolve/main/config.json -o config-Narsil-finbert.json -L

# Linux
cargo run --example bert --release --features intel-mkl -- "This is a test" -n 3

# M1
cargo run --example bert --release -- "This is a test" -n 3

为什么不使用库 X ?

市面上有很多 ML 库，如 torch 和 tensorflow 都很棒，但现在它们非常庞大，没有静态链接的选项。像 ONNX 这样的库也很有用，但当一个算子缺失时，工作起来真的很难。

对于底层库。 ggml 是一个很棒的库，没有依赖项，二进制文件大小极小。这实际上也是这个项目的灵感来源！但我的 C++ 编程能力还不够好，无法有效地对其进行修改。此外，它很难在预期范围之外使用，例如在编写 web 服务器/API 时，或者如果我们想使用 CUDA 作为后端。

dfdx 是另一个非常好的项目。我也从中得到了灵感。dfdx 的问题是类型系统，虽然它非常强大（编译时大小检查），但它会妨碍工作，并且对其进行优化并不像它那么简单，因为它更难了解正在发生什么。

这个库的架构

• [cpu] 包含所有各种精度的后端操作和张量结构。如果您想从头开始编写代码，这是您的首选。
• [nn] 包含所有基本层和实际模型实现。代码应类似于 torch 实现。
• [traits] 包含使 [nn] 能够独立于 [cpu] 编写的粘合剂，这应该可以使得使用不同的精度（或后端）变得相当容易。

模型看起来像什么样子？

pub struct BertClassifier<T: Tensor + TensorOps<T>> {
    bert: Bert<T>,
    pooler: BertPooler<T>,
    classifier: Linear<T>,
}

impl<T: Tensor + TensorOps<T>> BertClassifier<T> {
    pub fn new(bert: Bert<T>, pooler: BertPooler<T>, classifier: Linear<T>) -> Self {
        Self {
            bert,
            pooler,
            classifier,
        }
    }
    pub fn forward(&self, input_ids: &[usize], type_ids: &[usize]) -> Result<T, SmeltError> {
        let tensor = self.bert.forward(input_ids, type_ids)?;
        let tensor = self.pooler.forward(&tensor)?;
        let mut logits = self.classifier.forward(&tensor)?;
        T::softmax(&mut logits)?;
        Ok(logits)
    }
}

性能如何？

在相对较旧的电脑（i7-4790 CPU）上，这为 GPT-2 在 full f32 精度下提供了 ~40ms/token。相比之下，在同一硬件上，torch 提供 ~47ms/token，ggml 提供 ~37ms。

当前的实现不使用线程，也不使用预计算的 gelu/exp 或 ggml 可以使用的 f16 短路（如 softmax）。

因此，还有很大的改进空间，目前的大部分性能都来自于使用 intel-mkl 库，一旦实现了 ggml 的各种操作（希望获得完整性能），就可以丢弃这个库。