HackerNews中文版

TLDR: 我修改了 PyTorch 和 Triton 的内部结构。我修改了注意力机制，使其具有线性第一层、中间二次方层和线性最后一层。在测试中，推理速度大大加快，但困惑度略有下降。全注意力机制 O(n²): 17.96 秒 / 5.6 个 token/秒 HybridAttention O(n·W + n·D): 0.35 秒 / 286.6 个 token/秒我一直在 PyTorch 中从头开始构建一个专注于 Rust 的小型语言模型。这并非微调。它基于字节级别，从随机初始化开始，在 Rust 相关的语料库上进行训练，语料库在这里组装：https://codeberg.org/JohannaJuntos/Sisyphus 模型和训练设置该模型有 2560 万个参数，上下文长度为 512。它使用 256 个字节级词汇表，具有 8 层、8 个头和 512 维嵌入。位置嵌入是学习的，嵌入和 LM 头的权重是绑定的。使用单个 RTX 4060 Ti 8GB 在 173.5MB 的 Rust 语料库上训练了 3 万步。最终指标是训练损失 0.5834，验证损失 0.8217，困惑度为 2.15。最佳验证损失出现在大约第 18.5k 步，这表明存在一些晚期过拟合或平台期。架构该模型是 GPT 风格的解码器，但用每个层中的 HybridAttention 块替换了标准的全注意力机制。这结合了局部窗口因果注意力机制和类似 GRU 的循环状态路径，以及一个混合两者的学习门控。局部路径处理短程语法，而循环路径携带压缩的长程状态。门控偏置被初始化为在训练早期倾向于局部注意力。推理使用 Triton 内核和自定义 torch.library ops。语料库最大的收获来自语料库的扩展。运行开始时，大约有 31MB 来自 Rust 官方来源和主要项目，如 rustc、cargo、rust analyzer、tokio、serde、ripgrep、clap 和 axum。通过克隆前 500 个 crate，语料库扩展到 173.5MB，成功克隆了 461 个。这种扩展比任何架构上的改变都更有影响。推理性能全注意力机制的运行速度约为每秒 5.6 个 token，而带有 KV 缓存的 HybridAttention 达到每秒 286.6 个 token。这大约是 51 倍的加速，且没有可见的质量损失。 KV 缓存使用 VRAM 中 64 个 token 的热窗口，而较旧的 token 被压缩为 8 位幅度和角度，并且可以选择性地提升回全精度。对于此设置，这会将有效复杂度从二次方变为接近线性。质量 Rust 表面语法看起来不错，导入和函数签名通常是合理的。语义仍然很弱，重复和递归模式很常见。它看起来像 Rust，但还没有很好地推理。有趣的地方该项目结合了从头开始的字节级 Rust 预训练、混合局部注意力和循环架构、跨 Rust 生态系统的大规模语料库扩展，以及一种实用的 KV 缓存分页策略，可在消费级 GPU 上实现大幅加速。下一步我计划运行消融实验，比较混合注意力与仅局部和仅循环变体，评估大约 18.5k 步的检查点与最终模型，并添加语法级别的验证，例如解析和编译生成的代码。我还想探索将上下文长度从 256 扩展到 2048，并测试切换到 BPE 是否变得值得，因为语料库现在更大了。问题对于小型代码模型，除了困惑度之外，哪些评估最有帮助？有人看到混合局部加循环注意力机制在代码生成方面表现良好吗？考虑到这种设置，您会优先考虑更多 token、更长的上下文还是干净的消融实验？

查看原文

TLDR: Forked pytorch and triton internals . Changed attention so its linear first layer , middle quadratic layer, last linear layer Inference got much faster with a low perplexity hit in tests .Full attention O(n²): 17.96s / 5.6 tok/sHybridAttention O(n·W + n·D): 0.35s / 286.6 tok/sI have been building a small Rust focused language model from scratch in PyTorch. This is not a finetune. It is byte level, trained from random initialization on a Rust heavy corpus assembled here: https://codeberg.org/JohannaJuntos/SisyphusModel and training setupThe model has 25.6M parameters with a 512 context length. It uses a byte level vocabulary of 256, with 8 layers, 8 heads, and 512 dimensional embeddings. Positional embeddings are learned and the embedding and LM head weights are tied.Training ran for 30k steps on a 173.5M byte Rust corpus using a single RTX 4060 Ti 8GB.Final metrics were a train loss of 0.5834, validation loss of 0.8217, and perplexity of 2.15. The best validation loss occurred around step 18.5k, which suggests some late overfitting or plateau.ArchitectureThe model is a GPT style decoder, but replaces standard full attention with a HybridAttention block in each layer. This combines local windowed causal attention with a GRU like recurrent state path, along with a learned gate that mixes the two.The local path handles short range syntax, while the recurrent path carries compressed long range state. The gate bias is initialized to favor local attention early in training.Inference uses Triton kernels and custom torch.library ops.CorpusThe biggest gain came from corpus expansion.The run started with about 31MB from Rust official sources and major projects such as rustc, cargo, rust analyzer, tokio, serde, ripgrep, clap, and axum. The corpus was expanded to 173.5M bytes by cloning the top 500 crates, with 461 successful clones.This expansion had more impact than any architectural change.Inference performanceFull attention runs at about 5.6 tokens per second, while HybridAttention with KV cache reaches 286.6 tokens per second. This is about a 51x speedup with no visible quality loss.The KV cache uses a hot window of 64 tokens in VRAM, while older tokens are compressed to 8 bit magnitude and angle and can be selectively promoted back to full precision. This changes the effective complexity from quadratic to near linear for this setup.QualitySurface Rust syntax looks decent, and imports and function signatures are often plausible. Semantics are still weak, and repetition and recursive patterns are common. It looks like Rust, but does not reason well yet.What seems interestingThis project combines byte level Rust only pretraining from scratch, a hybrid local attention and recurrent architecture, large scale corpus expansion across the Rust ecosystem, and a practical KV cache paging strategy that delivers large speedups on consumer GPUs.Next stepsI plan to run ablations comparing hybrid attention against local only and recurrent only variants, evaluate checkpoints around 18.5k versus the final model, and add syntax level validation such as parsing and compiling generated code. I also want to explore scaling context length from 256 up to 2048 and test whether switching from byte level to BPE becomes worthwhile now that the corpus is larger.QuestionsFor small code models, which evaluations have been most useful beyond perplexity?Has anyone seen hybrid local plus recurrent attention work well for code generation?Given this setup, would you prioritize more tokens, longer context, or clean ablations first?

混合注意力机制