DeepSeek 开源周day1: FlashMLA

https://github.com/deepseek-ai/FlashMLA FlashMLA FlashMLA 是针对Hopper架构GPU优化的高效MLA解码内核,专为变长序列服务场景设计。 当前已发布特性: BF16支持 分页式kvcache(块大小为64) 快速入门 安装 python setup.py install 性能测试 python tests/test_flash_mla.py 在H800 SXM5(CUDA 12.6环境)上实现内存受限配置下3000GB/s的带宽吞吐,计算受限配置下580 TFLOPS的算力表现。 使用方法 from flash_mla import get_mla_metadata, flash_mla_with_kvcache tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) for i in range(num_layers): ... o_i, lse_i = flash_mla_with_kvcache( q_i, kvcache_i, block_table, cache_seqlens, dv, tile_scheduler_metadata, num_splits, causal=True, ) ... 环境要求 Hopper架构GPU CUDA 12.3及以上版本 PyTorch 2.0及以上版本 致谢 FlashMLA的设计灵感来源于FlashAttention 2&3以及cutlass项目。

February 24, 2025 · 小茄墩

DeepSeek新作原生稀疏注意力

摘要 长上下文建模对于下一代语言模型至关重要,然而标准注意力机制的高计算成本带来了显著的计算挑战。稀疏注意力为提高效率同时保持模型能力提供了有前景的方向。我们提出了NSA(可原生训练的稀疏注意力机制),通过算法创新与硬件特性对齐的优化相结合,实现了高效的长上下文建模。NSA采用动态分层稀疏策略,结合粗粒度Token压缩与细粒度Token选择,在保持全局上下文感知能力的同时确保局部精度。该方法通过两大关键创新推进稀疏注意力设计:(1)通过算术强度平衡的算法设计实现显著加速,并针对现代硬件进行实现优化;(2)支持端到端训练,在保持模型性能的前提下减少预训练计算量。如图1所示,实验表明采用NSA预训练的模型在通用基准测试、长上下文任务和基于指令的推理中均保持或超越全注意力模型。与此同时,NSA在处理64k长度序列时,在解码、前向传播和后向传播阶段相较全注意力机制均实现了显著加速,有效验证了该方案在整个模型生命周期中的效率优势。 介绍 一种实现高效长文本建模的自然方法是利用 softmax 注意力的固有稀疏性。通过选择性地计算关键的查询-键对,可以在保持性能的同时显著降低计算开销。最近的进展通过各种策略展示了这种潜力,包括:KV-cache 驱逐方法、分块 KV-cache 选择方法,以及基于采样、聚类或哈希的选择方法。尽管这些策略很有前景,但现有的稀疏注意力方法在实际部署中往往表现不足。许多方法未能实现与其理论收益相当的加速;此外,大多数方法主要关注推理阶段,缺乏有效的训练时支持来充分利用注意力的稀疏模式,从而无法充分挖掘注意力的稀疏模式。为了解决这些局限性,有效部署稀疏注意力必须应对两个关键挑战:(1)硬件对齐的推理加速:将理论计算量的减少转化为实际的速度提升,需要在预填充和解码阶段进行硬件友好的算法设计,以缓解内存访问和硬件调度瓶颈;(2)训练感知的算法设计:通过可训练的算子实现端到端计算,以降低训练成本,同时保持模型性能。这些要求对于实际应用至关重要,以实现快速的长文本推理或训练。综合考虑这两个方面,现有方法仍然存在明显的差距。 为了实现更有效和高效的稀疏注意力机制,我们提出了 NSA,一种原生可训练的稀疏注意力架构,它集成了分层 Token 建模。如图2 所示,NSA 通过将键 (key) 和值 (value) 组织成时间块,并通过三个注意力路径处理它们来减少每个查询的计算量:压缩的粗粒度 Token、选择性保留的细粒度 Token 和用于局部上下文信息的滑动窗口。然后,我们实现专门的内核以最大限度地提高其实际效率。NSA 引入了两个与上述关键要求相对应的核心创新:(1)硬件对齐系统:优化块状稀疏注意力,以充分利用 Tensor Core 并优化内存访问,确保平衡的算术强度。(2)训练感知设计:通过高效的算法和反向算子实现稳定的端到端训练。这种优化使 NSA 能够支持高效部署和端到端训练。我们通过对真实世界语言语料库的综合实验来评估 NSA。在使用 260B Token 的 27B 参数 Transformer 主干上进行预训练后,我们评估了 NSA 在通用语言评估、长上下文评估和思维链推理评估中的性能。我们进一步比较了 A100 GPU 上内核速度与优化的 Triton 实现。实验结果表明,NSA 实现了与完整注意力基线相当或更优越的性能,同时优于现有的稀疏注意力方法。此外,与完整注意力相比,NSA 在解码、前向和后向阶段都提供了显著的加速,并且加速比随着序列长度的增加而增加。这些结果验证了我们的分层稀疏注意力设计有效地平衡了模型能力和计算效率。 重新思考稀疏注意力方法 大多数方法主要在推理过程中应用稀疏性,同时保留预训练的完整注意力骨干网络,这可能会引入架构偏差,从而限制它们充分利用稀疏注意力的优势。在介绍我们原生的稀疏架构之前,我们通过两个关键视角系统地分析这些局限性。 高效推理的错觉 尽管在注意力计算中实现了稀疏性,但许多方法未能实现相应的推理延迟降低,这主要是由于两个挑战:阶段性稀疏。诸如 H2O 之类的方法在自回归解码期间应用稀疏性,同时在预填充期间需要计算密集型预处理(例如,注意力图计算、索引构建)。相比之下,像 MInference 这样的方法仅专注于预填充稀疏性。这些方法未能实现跨所有推理阶段的加速,因为至少一个阶段的计算成本与完整注意力机制相当。这种阶段专业化降低了这些方法在以预填充为主的工作负载(如书籍摘要和代码补全)或以解码为主的工作负载(如长链式思维推理)中的加速能力。 与高级注意力架构的不兼容。 一些稀疏注意力方法无法很好地适配现代高效解码架构,例如多查询注意力(MQA)和分组查询注意力(GQA)。这些架构通过在多个查询头之间共享 KV,显著降低了解码过程中的内存访问瓶颈。 例如,像 Quest 这样的方法中,每个注意力头独立地选择其 KV-cache 子集。 虽然这种方法在多头注意力(MHA)模型中表现出一致的计算稀疏性和内存访问稀疏性,但在基于 GQA 等架构的模型中,情况则有所不同。在 GQA 架构中,KV-cache 的内存访问量取决于同一 GQA 组内所有查询头选择的并集。 这种架构特性意味着,虽然这些方法可以减少计算操作,但所需的 KV-cache 内存访问量仍然相对较高。 这种限制迫使我们面临一个关键选择:虽然一些稀疏注意力方法减少了计算量,但它们分散的内存访问模式与高级架构中高效的内存访问设计相冲突。 这些限制的出现是因为许多现有的稀疏注意力方法侧重于 KV-cache 减少或理论计算减少,但难以在高级框架或后端中实现显著的延迟降低。 这促使我们开发算法,将先进的架构和硬件高效的实现相结合,从而充分利用稀疏性来提高模型效率。 ...

February 18, 2025 · 小茄墩

Open R1 项目 第二周总结与展望

摘要 我们现在已经进入了 Open R1 项目 的第二周,该项目旨在重建 DeepSeek R1 缺失的部分——特别是训练管道和合成数据。 分享 OpenR1-Math-220k 的构建:这是我们首个用于数学推理的大规模数据集! 介绍社区在策划用于微调的小型、高质量数据集方面取得的一些令人兴奋的进展,以及关于如何在训练和推理阶段控制推理模型的思维链长度的见解。 OpenR1-Math-220k 数据集 DeepSeek R1 的主要优势之一是它能够通过知识蒸馏将高级推理能力迁移到较小的模型。 DeepSeek 团队通过生成 60 万个推理轨迹并微调一系列 Qwen 和 Llama 模型证明了这一点,表明直接从 R1 进行知识蒸馏可以在无需强化学习的情况下实现具有竞争力的推理性能。 值得注意的是,DeepSeek-R1-Distill-Qwen-7B 在 AIME 2024 上取得了 55.5% 的成绩,超过了像 QwQ-32B-Preview 这样更大的模型。 然而,用于蒸馏的推理轨迹尚未公开,这促使社区独立地重新创建类似的数据集。到目前为止,社区已经发布了多个开放数据集,包括 OpenThoughts-114k、Bespoke-Stratos-17k、Dolphin-R1 和 LIMO。 🐳 隆重推出 OpenR1-Math-220k,这是一个大规模的数学推理数据集,它利用 512 个 H100 在本地生成,且每个问题都对应多个答案。为了创建 OpenR1-Math-220k,我们与 Numina 展开合作,他们开发了广受欢迎的 NuminaMath-CoT 数据集的全新版本。 与现有数据集相比,OpenR1 数据集的新特性:80 万条 R1 推理轨迹:我们使用 DeepSeek R1 为 40 万道问题生成了两个答案。经过筛选的数据集包含 22 万道问题,并带有正确的推理轨迹。 本地运行 512 个 H100: 我们没有依赖 API,而是利用 vLLM 和 SGLang 在我们的科学集群上本地运行生成,每天生成 18 万条推理过程。 基于 NuminaMath 1.5: 我们专注于数学推理过程,并为 NuminaMath 1.5 中的问题生成答案,NuminaMath 1.5 是 NuminaMath-CoT 数据集的改进版本。 自动过滤: 我们应用 Math Verify 来仅保留至少有一个正确答案的问题。我们还利用 Llama3.3-70B-Instruct 作为一个判断器,以检索更多正确的例子(例如,对于答案格式错误,无法使用基于规则的解析器验证的情况)。 我们通过在我们的数据集上微调 Qwen-7B-Math-Instruct 来匹配 DeepSeek-Distill-Qwen-7B 的性能。 通过展示可扩展的、高质量的推理数据生成,我们希望这个流程可以扩展到数学以外的领域,例如代码生成。 ...

February 11, 2025 · 小茄墩

使用 Unsloth 训练您自己的 R1 推理模型

今天,我们很高兴地介绍 Unsloth 中的推理功能!DeepSeek 的 R1 研究揭示了一个“顿悟时刻”,在这个时刻,R1-Zero 通过使用群体相对策略优化(GRPO)自主学习分配更多的思考时间,而无需人类反馈。 我们增强了整个 GRPO 过程,使其使用的 VRAM 比 Hugging Face + FA2 少 80%。这使您能够仅使用 7GB 的 VRAM 和 Qwen2.5(1.5B)重现 R1-Zero 的“顿悟时刻”。 尝试我们的免费 GRPO notebook:Colab 上的 Llama 3.1(8B) 有关其他模型(如 Phi-4)的 GRPO 笔记本,请访问我们的文档 💡 主要细节 使用 15GB VRAM,Unsloth 允许您将任何模型(最多 15B 参数)转换为推理模型,例如 Llama 3.1(8B)、Phi-4(14B)、Mistral(7B)或 Qwen2.5(7B) 最低要求:仅需 7GB VRAM 即可在本地训练您自己的推理模型。 Tiny-Zero 的出色团队展示了您可以使用 Qwen2.5(1.5B)实现自己的“顿悟时刻——但这需要 2xA100 GPU(160GB VRAM)。现在,使用 Unsloth,您只需一台 7GB VRAM 的 GPU 就可以实现同样的“顿悟时刻”。 之前,GRPO 仅支持完全微调,但我们已使其与 QLoRA 和 LoRA 一起工作。 请注意,这并不是微调 DeepSeek 的 R1 蒸馏模型或使用 R1 的蒸馏数据进行微调,Unsloth 已经支持。这是将标准模型转换为完整的推理模型,使用 GRPO。 ...

February 8, 2025 · 小茄墩

R1-Zero类训练中可能没有顿悟时刻 —— 一项初步研究

DeepSeek-R1-Zero最鼓舞人心的结果之一是通过纯强化学习(RL)出现的“顿悟时刻”。在顿悟时刻,模型学习到自我反思等新兴技能,这有助于它进行上下文搜索以解决复杂的推理问题。 在R1-Zero发布后的几天内,几个项目独立地在较小的规模(例如,1B到7B)上“再现”了R1-Zero类训练,并且都观察到了顿悟时刻,这通常伴随着响应长度的增加。我们遵循他们的设置来仔细审查R1-Zero类训练过程,并在本博客中分享以下发现: 💡 a) R1-Zero类训练中可能没有顿悟时刻。 相反,我们发现顿悟时刻(例如自我反思模式)出现在第0个周期,即基础模型。 b) 我们发现基础模型的响应中存在表面自我反思(SSR),在这种情况下,自我反思不一定导致正确的最终答案。 c) 我们通过RL对R1-Zero类训练进行了更深入的观察,发现响应长度增加的现象并不是由于自我反思的出现,而是RL优化良好设计的基于规则的奖励函数的结果。 1. 顿悟时刻出现在第0个周期 1.1 实验设置 基础模型。我们调查了由不同组织开发的广泛的基础模型系列,包括Qwen-2.5, Qwen-2.5-Math, DeepSeek-Math, Rho-Math, 和 Llama-3.x。 提示模板。我们直接使用在R1-Zero和SimpleRL-Zero中应用的模板来提示基础模型: 模板1(与R1-Zero相同) A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>. User: {Question} Assistant: ...

February 7, 2025 · 小茄墩

cuda层面实现kernel的库Liger Kernel

看到一个cuda层面实现kernel的库Liger Kernel,速度极快。 ==直接一行调用GRPO loss:== grpo_loss = LigerFusedLinearGRPOLoss() 以下为Liger Kernel库简介: Liger Kernel 是一个专门为大语言模型(LLM)训练设计的 Triton 内核集合。它可以有效提高多 GPU 训练的吞吐量 20%,并减少 60% 的内存使用。我们已经实现了与 Hugging Face 兼容的 RMSNorm、RoPE、SwiGLU、CrossEntropy、FusedLinearCrossEntropy 等功能,未来还会增加更多。该内核开箱即用,支持 Flash Attention、PyTorch FSDP 和 Microsoft DeepSpeed。 我们还添加了优化的训练后内核,为对齐和蒸馏任务节省高达 80% 的内存。我们支持 DPO、CPO、ORPO、SimPO、JSD 等多种损失函数。 只需一行代码,Liger Kernel 就可以提高 20% 以上的吞吐量,并减少 60% 的内存使用,从而实现更长的上下文长度、更大的批量大小和更大的词汇表。 注意: 基准测试条件:LLaMA 3-8B,批量大小 = 8,数据类型 = bf16,优化器 = AdamW,梯度检查点 = True,分布式策略 = FSDP1,使用 8 个 A100 GPU。 Hugging Face 模型在 4K 上下文长度时开始出现内存不足(Out of Memory, OOM),而 Hugging Face + Liger Kernel 可以扩展到 16K。 ...

February 5, 2025 · 小茄墩