TokenSwift:90分钟生成10万Token!文本生成提速3倍,无损加速黑科技

本文涉及的产品
NLP 自学习平台,3个模型定制额度 1个月
NLP自然语言处理_基础版,每接口每天50万次
NLP自然语言处理_高级版,每接口累计50万次
简介: TokenSwift 是北京通用人工智能研究院团队推出的超长文本生成加速框架,能在90分钟内生成10万Token的文本,速度提升3倍,生成质量无损,支持多种模型架构。

❤️ 如果你也关注 AI 的发展现状,且对 AI 应用开发感兴趣,我会每日分享大模型与 AI 领域的开源项目和应用,提供运行实例和实用教程,帮助你快速上手AI技术!

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜:蚝油菜花 🥦


🚀 "长文本生成进入高铁时代!北通院黑科技让AI写作效率飙升300%"

大家好,我是蚝油菜花。你是否也经历过这些抓狂时刻——

  • 👉 写行业白皮书卡在生成环节,5小时只能产出10万Token
  • 👉 加速时遭遇"鬼畜循环",相同段落反复出现
  • 👉 换不同模型就要重构代码,开发效率断崖下跌

今天带来的 TokenSwift 正在改写游戏规则!这个由北京通用人工智能研究院打造的文本引擎,凭借多Token并行生成+动态KV缓存两大绝技,实现:

  • ✅ 10万Token文本90分钟极速生成(传统方案需5小时)
  • ✅ 独创上下文惩罚机制,重复率降低76%
  • ✅ 无缝适配1.5B到14B不同规模模型

从万字小说到百页行业报告,现在连科研团队都在用它冲刺论文——你的键盘准备好迎接这场生产力核爆了吗?

🚀 快速阅读

TokenSwift 是一个超长文本生成加速框架,能够在 90 分钟内生成 10 万 Token 的文本。

  1. 核心功能:显著缩短超长文本生成时间,支持多种模型架构,生成质量无损。
  2. 技术原理:通过多 Token 并行生成与 Token 复用、动态 KV 缓存更新策略、树结构的多候选 Token 验证以及上下文惩罚策略,提升生成效率。

TokenSwift 是什么

TokenSwift

TokenSwift 是北京通用人工智能研究院团队推出的超长文本生成加速框架,能够在 90 分钟内生成 10 万 Token 的文本,相比传统自回归模型的近 5 小时,速度提升了 3 倍,生成质量无损。

TokenSwift 通过多 Token 生成与 Token 重用、动态 KV 缓存更新以及上下文惩罚机制等技术,减少模型加载延迟、优化缓存更新时间并确保生成多样性。支持多种不同规模和架构的模型,如 1.5B、7B、8B、14B 的 MHA 和 GQA 模型。

TokenSwift 的主要功能

  • 超长文本生成加速:TokenSwift 能显著缩短超长文本的生成时间。例如,传统自回归模型生成 10 万 Token 的文本可能需要近 5 个小时,TokenSwift 仅需 90 分钟,速度提升了 3 倍,极大地提高了生成效率。
  • 无损加速:在加速的同时,TokenSwift 能保持原始模型的输出质量,确保生成文本的质量和多样性。
  • 支持多种模型:TokenSwift 支持多种不同规模和架构的模型,包括 1.5B、7B、8B、14B 的模型,以及多头注意力(MHA)和分组查询注意力(GQA)架构。

TokenSwift 的技术原理

  • 多 Token 并行生成与 Token 复用:TokenSwift 借鉴了 Medusa 等方法,通过引入额外的线性层,使模型在一次前向传播中能同时生成多个草稿 Token。基于生成文本中的 n-gram 频率信息,系统会自动检索并复用高频短语,进一步减少模型重新加载的次数,提升整体效率。
  • 动态 KV 缓存更新策略:在 KV 缓存管理上,TokenSwift 采用动态更新策略。系统在生成过程中将初始 KV 缓存保留,同时根据 Token 的重要性对后续缓存进行有序替换。
  • 基于树结构的多候选 Token 验证:为保证生成结果与目标模型预测的一致性,TokenSwift 引入了树形注意力机制。通过构建包含多个候选 Token 组合的树形结构,采用并行验证的方式,从中随机选择最长且有效的 n-gram 作为最终输出,确保生成过程无损且多样性得到提升。
  • 上下文惩罚策略:为了进一步抑制重复生成问题,TokenSwift 设计了一种上下文惩罚方法。在生成过程中为近期生成的 Token 施加惩罚,使模型在选择下一 Token 时更倾向于多样化输出,有效减少重复现象。

如何运行 TokenSwift

1. 安装

方法 1:使用 pip 安装

pip install tokenswift

方法 2:从源码安装(推荐)

git clone https://github.com/bigai-nlco/TokenSwift.git
cd TokenSwift
conda create -n tokenswift python=3.11
conda activate tokenswift
conda install nvidia::cuda-nvcc
pip install -r requirements.txt
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

2. 模型下载

3. 推理

以 LLaMA3.1-8B 为例:

torchrun  --master-port 1111 --nproc_per_node=1 main.py \
    --model_type llama3_1 \
    --ckpt_path your_checkpoint_path \
    --prefill_len 4096 \
    --retrival_max_budget 4096 \
    --gen_len 102400 \
    --gamma 4 \
    --min_p 0.1 \
    --temperature 1.0 \
    --tree_decoding \
    --ngram_topk 20 \
    --penalty 1.2 \
    --penalty_length 1024 \
    --prompt_id 0

  <NOTE: 修改数据和模型路径>

对于其他模型,可以运行 infer_scripts/ 文件夹中的脚本。例如:

bash infer_scripts/r1_qwen_32b.sh

4. 训练指南(可选)

数据集下载

PG-19训练集中,根据不同的 tokenizer 过滤出大于 8K 的数据。

或从以下链接下载处理后的训练数据集:

如何训练

以 LLaMA3.1-8B 为例:

torchrun --master-port 1111 --nproc_per_node=4 train/train_legacy.py \
    --model_name_or_path /your_model_path/Meta-Llama-3.1-8B \
    --llama_type llama3_1 \
    --data_path /your_data_path/llama3_1_pg19_8k_data \
    --output_dir /your_checkpoint_path/adapter_ckpts_llama3_1 \
    --max_steps 200 \
    --per_device_train_batch_size 3 \
    --gradient_accumulation_steps 10 \
    --save_steps 200 \
    --learning_rate 5e-3 \
    --weight_decay 0.1 \
    --warmup_steps 50 \
    --lr_scheduler_type cosine \
    --logging_steps 5 \
    --report_to tensorboard \
    --bf16 True \
    --medusa_heads 3 \
    --remove-unused-columns false

  <NOTE: 修改数据和模型路径>

对于其他模型,可以运行 train/scripts/ 文件夹中的脚本。例如:

cd train
bash scripts/train_R1_qwen2_5_32b.sh

资源


❤️ 如果你也关注 AI 的发展现状,且对 AI 应用开发感兴趣,我会每日分享大模型与 AI 领域的开源项目和应用,提供运行实例和实用教程,帮助你快速上手AI技术!

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜:蚝油菜花 🥦

相关文章
|
6月前
|
存储 人工智能 JSON
Evolving Agents:开源Agent革命!智能体动态进化框架上线,复杂任务一键协同搞定
Evolving Agents 是一个开源的AI Agent管理与进化框架,支持智能代理之间的通信与协作,能够根据语义理解需求动态进化,适用于文档处理、医疗保健、金融分析等多个领域。
288 26
Evolving Agents:开源Agent革命!智能体动态进化框架上线,复杂任务一键协同搞定
|
6月前
|
机器学习/深度学习 人工智能 算法
MIDI-3D:单图秒变3D场景!40秒生成360度空间,多实例扩散黑科技
MIDI-3D 是一种先进的 AI 3D 场景生成技术,能够将单张图像快速转化为高保真度的 360 度 3D 场景,具有强大的全局感知能力和细节表现力,适用于游戏开发、虚拟现实、室内设计等多个领域。
223 18
MIDI-3D:单图秒变3D场景!40秒生成360度空间,多实例扩散黑科技
|
6月前
|
存储 人工智能 前端开发
Botgroup.chat:超有趣的开源 AI 聊天室!多个 AI 在线互怼,一键搭建你的专属 AI 社群
Botgroup.chat 是一款基于 React 和 Cloudflare Pages 的开源 AI 聊天应用,支持多个 AI 角色同时参与对话,提供类似群聊的交互体验。
1720 28
|
6月前
|
人工智能 自然语言处理 并行计算
Chitu:清华核弹级开源!推理引擎3倍提速+50%省卡,国产芯片告别英伟达绑架
Chitu(赤兔)是清华大学与清程极智联合开源的高性能大模型推理引擎,支持多硬件适配,显著提升推理效率,适用于金融、医疗、交通等多个领域。
649 10
Chitu:清华核弹级开源!推理引擎3倍提速+50%省卡,国产芯片告别英伟达绑架
|
6月前
|
人工智能 监控 搜索推荐
MedRAG:医学AI革命!知识图谱+四层诊断,临床准确率飙升11.32%
MedRAG是南洋理工大学推出的医学诊断模型,结合知识图谱与大语言模型,提升诊断准确率11.32%,支持多模态输入与智能提问,适用于急诊、慢性病管理等多种场景。
333 32
MedRAG:医学AI革命!知识图谱+四层诊断,临床准确率飙升11.32%
|
6月前
|
人工智能 自然语言处理
TxGemma:谷歌DeepMind革命药物研发!270亿参数AI药理学家24小时在线
谷歌推出专为药物研发设计的TxGemma大模型,具备药物特性预测、生物文献筛选、多步推理等核心能力,提供20亿至270亿参数版本,显著提升治疗开发效率。
221 7
TxGemma:谷歌DeepMind革命药物研发!270亿参数AI药理学家24小时在线
|
5月前
|
人工智能 自然语言处理 JavaScript
测试工程师要失业?Magnitude:开源AI Agent驱动的端到端测试框架,让Web测试更智能,自动完善测试用例!
Magnitude是一个基于视觉AI代理的开源端到端测试框架,通过自然语言构建测试用例,结合推理代理和视觉代理实现智能化的Web应用测试,支持本地运行和CI/CD集成。
729 15
测试工程师要失业?Magnitude:开源AI Agent驱动的端到端测试框架,让Web测试更智能,自动完善测试用例!
|
6月前
|
前端开发 搜索推荐
使用DeepSeek快速创建的个人网站
这是一份使用DeepSeek快速创建个人网站的10分钟指南。内容分为四个步骤:搭建基础架构(HTML框架)、设计核心内容区块(关于我、作品展示等)、快速配置样式(CSS美化页面)以及添加联系表单并部署到GitHub Pages。通过简单的代码和DeepSeek的智能辅助功能,用户可以轻松实现个性化调整,如更换主题色、增加模块或优化响应式设计。虽然整体流程简单高效,但可能因功能有限或美观度不足而需进一步扩展与改进。
579 11
|
6月前
|
传感器 存储 安全
鸿蒙开发:权限管理之权限声明
本文,主要简单概述了为什么要有权限管理,以及权限管理的声明原则,这些都是基本的概念内容,大家做为了解即可,重要的是怎么声明权限,在什么位置声明权限,这一点需要掌握。
224 16
鸿蒙开发:权限管理之权限声明

热门文章

最新文章