shadow-o-tokens-20b

SHADOW-O OLMo-2 Top-16 Logits (20B) 是一个用于知识蒸馏训练的大型稀疏下一词元教师logits数据集，专门设计用于训练参数量为252.8M的SHADOW-O学生模型。技术规格方面：教师模型采用allenai/OLMo-2-0425-1B；词表ID空间基于AI2 Dolma/OLMo的前100,278个ID；上下文长度为2,048个词元；存储格式为每个词元保存输入token（uint32）、前16个预测token的ID（uint32）以及对应的未归一化logits值（float16）。数据集计划规模为200亿词元，约2.0TB存储空间，按逻辑分片组织，每个分片包含三个同名文件：*.tok.npy（输入词元，形状[N]）、*.ids.npy（教师预测top-16 token ID，形状[N, 16]）和*.val.npy（教师未归一化logits，形状[N, 16]）。数据混合计划包含四个组成部分：基础数据（6.4B词元，教育导向的网络爬取数据、维基百科和Stack-Edu）、密度数据（4.55B词元，数学、代码和阅读理解）、推理数据（4.55B词元，数学、代码、合成QA和元推理）以及退火数据（4.5B词元，推理轨迹、FLAN/Tülu监督微调数据和可验证问题）。底层词元流来源于AI2的公开OLMo/Dolma数据分布，保留了相应的源许可证和归属要求。捕获过程具有可恢复性，仅在所有数组写入并上传后才提交源分片。学生模型在捕获的前16个词元集合上对教师和学生分布进行重归一化处理。SHADOW-O OLMo-2 Top-16 Logits (20B) is a large-scale sparse next-token teacher logits dataset for knowledge distillation training, specifically designed for training the SHADOW-O student model with 252.8M parameters. Technical specifications include: teacher model using allenai/OLMo-2-0425-1B; vocabulary ID space based on the first 100,278 IDs from AI2 Dolma/OLMo; context length of 2,048 tokens; storage format saving per token the input token (uint32), top-16 predicted token IDs (uint32), and corresponding unnormalized logits values (float16). The dataset is planned to scale to 20 billion tokens, approximately 2.0TB in storage, organized into logical shards, each containing three files with the same name: *.tok.npy (input tokens, shape [N]), *.ids.npy (teacher-predicted top-16 token IDs, shape [N, 16]) and *.val.npy (teacher unnormalized logits, shape [N, 16]). The data mixture plan consists of four components: base data (6.4B tokens, education-oriented web crawl data, Wikipedia, and Stack-Edu), density data (4.55B tokens, math, code, and reading comprehension), reasoning data (4.55B tokens, math, code, synthetic QA, and meta-reasoning), and annealing data (4.5B tokens, reasoning traces, FLAN/Tülu supervised fine-tuning data, and verifiable questions). The underlying token stream originates from AI2s public OLMo/Dolma data distribution, retaining corresponding source licenses and attribution requirements. The capture process is recoverable, with source shards committed only after all arrays are written and uploaded. The student model renormalizes teacher and student distributions over the captured set of top-16 tokens.

原始信息汇总

SHADOW-O OLMo-2 Top-16 Logits (20B) 数据集概述

数据集名称：SHADOW-O OLMo-2 top-16 logits（20B tokens）

许可证：ODC-BY（Open Data Commons Attribution License）

任务类别：文本生成

语言：英语

---

数据集用途

该数据集用于训练252.8M参数的SHADOW-O学生模型，提供稀疏的下一词元教师logits。教师模型为 allenai/OLMo-2-0425-1B（1B参数）。

---

技术细节

• 词元ID空间：AI2 Dolma/OLMo 词表，前100,278个ID。
• 上下文长度：2,048个词元。
• 存储格式：每个词元存储以下内容：
  • 输入词元（uint32）
  • 前16个ID（uint32）
  • 前16个logits（float16）
• 预计规模：20B tokens，约2.0 TB。
• 捕获代码：10_distill/capture_logits.py（位于SHADOW-O工作空间）。

数据分片逻辑：每个分片包含三个同名的文件：

• *.tok.npy：输入词元，形状 [N]
• *.ids.npy：教师前16个词元ID，形状 [N, 16]
• *.val.npy：未归一化的教师logits，形状 [N, 16]

捕获过程可中断恢复，仅在所有数组写入并上传完成后提交分片。学生模型在捕获的前16个集合上对教师分布和学生分布进行重新归一化。

---

数据混合计划（20B tokens）

• Foundation（64亿tokens）：教育导向的网络爬虫数据、维基百科、Stack-Edu
• Density（45.5亿tokens）：数学、代码、阅读理解
• Reasoning（45.5亿tokens）：数学、代码、合成问答、元推理
• Anneal（45亿tokens）：推理轨迹、FLAN/Tülu SFT、可验证问题

底层词元流来自AI2公开的OLMo/Dolma数据分布，并保留其适用的源许可证和署名要求。

搜集汇总

数据集介绍

构建方式

该数据集专为训练252.8M参数的SHADOW-O学生模型而设计，通过捕获教师模型（allenai/OLMo-2-0425-1B）对每个输入令牌的稀疏前16个logits来构建。令牌ID空间基于AI2 Dolma/OLMo的前100,278个ID，上下文长度固定为2,048。每个逻辑分片包含三个文件：输入令牌（.tok.npy）、教师前16个令牌ID（.ids.npy）及未归一化的logits（.val.npy）。捕获过程支持断点续传，确保数据完整性，仅在所有数组写入并上传后才提交分片。学生模型在训练时对教师和自身分布在前16个logits集上进行重新归一化。

特点

该数据集的突出特点在于其稀疏性设计，仅存储每令牌前16个最高logits，大幅减少存储与计算开销，同时保留教师模型的关键知识。数据集规模达20B令牌，约2.0 TB，涵盖多样化的训练混合：包括6.4B的面向教育的基础语料（如Wikipedia、Stack-Edu）、4.55B的高密度数学与代码数据、4.55B的推理增强数据（含合成问答与元推理），以及4.5B的退火阶段数据（含推理轨迹与可验证问题）。所有令牌流源自AI2公开的OLMo/Dolma数据分布，并保留原始许可与归属要求。

使用方法

该数据集主要用于知识蒸馏场景，支持学生模型通过教师logits学习生成分布。用户可加载分片中的.tok.npy、.ids.npy和.val.npy文件，结合稀疏logits重新归一化后计算交叉熵损失。建议根据规划混合比例选择分片，按顺序或随机组合训练。由于上下文长度为2,048，需确保输入序列对齐。代码实现参考SHADOW-O工作区的10_distill/capture_logits.py脚本，可复现捕获流程。数据集以ODC-BY许可发布，使用时需遵守基础数据的来源许可与归属要求。

背景与挑战

背景概述

在自然语言处理领域，知识蒸馏作为一种有效的模型压缩技术，通过将大型教师模型的知识迁移至小型学生模型，在保持性能的同时显著降低计算开销。SHADOW-O Tokens-20B数据集由艾伦人工智能研究所（AI2）于近期创建，核心研究问题在于利用稀疏化注意力机制实现高效的知识蒸馏。该数据集以OLMo-2-0425-1B模型作为教师，存储其针对252.8M参数学生模型生成的前16个候选标记的对数概率，共计20B个训练标记，容量约2.0TB。其研究背景植根于大语言模型部署时对高效推理的需求，该数据集为稀疏化蒸馏提供了标准化训练资源，对推动轻量化语言模型的发展具有重要影响力。

当前挑战

数据集面临的挑战涵盖领域问题与构建过程两个层面。领域层面，知识蒸馏需解决教师与学生模型间分布差异的挑战，尤其在仅保留稀疏化对数概率（top-16 logits）时，如何确保学生模型能有效学习完整分布信息成为关键。构建过程中，数据集需管理超大规模数据存储与处理：20B标记的捕获需确保可恢复性，每个逻辑分片包含输入标记、教师候选ID及未归一化对数概率三部分，由于数据量约2.0TB，如何高效分片上传并维持数据完整性成为技术难点。此外，数据混合来源多样（包括教育爬虫、数学推理、代码等），需平衡不同领域数据的分布以支持学生模型的泛化能力。

常用场景

经典使用场景

在语言模型知识蒸馏领域，SHADOW-O Tokens-20B数据集被设计为教师模型OLMo-2-0425-1B向252.8M参数学生模型传递稀疏对数概率的桥梁。该数据集以20B tokens的大规模容量，存储了教师模型对每个输入token预测的前16个候选token的ID及对应的未归一化logits。这一经典用法使得学生模型能够通过局部重归一化，在仅有2048上下文长度的条件下，高效学习教师模型的输出分布，从而在保持计算可行性的同时实现知识迁移。

解决学术问题

该数据集核心解决了大规模语言模型蒸馏中计算成本与知识保真度的矛盾问题。通过仅保留教师模型最顶部的16个logits而非完整词表（100,278个ID），大幅降低了存储与计算开销，同时避免了因忽略尾部概率导致的知识损失。这为探讨稀疏蒸馏理论提供了实证基础，揭示了在有限资源下如何平衡教师信号完整性与学生模型性能的学术难题，推动了轻量级模型继承大模型推理能力的相关研究。

衍生相关工作

该数据集的发布催生了一系列关于稀疏条件蒸馏方法的后续研究。基于其存储格式，衍生工作探索了不同稀疏度（如4位或8位logits量化）对蒸馏效果的影响，以及针对特殊领域（如代码或数学推理）调整top-K候选数量的自适应策略。此外，该数据集的捕获代码与混合数据配比方案（包括基础、密度、推理与退火阶段）为后续建立更高效的蒸馏数据管线提供了可复现的基准框架。