Jamba前生今世:1.5开源来袭
AI21服务于企业,为企业构建基础模型和AI系统以加速GenAI在生产中的使用。AI21 成立于2017年,已从NVIDIA、Intel、Google等公司共筹集了3.36亿美元。它是最早将生成式AI推向大众的公司之一,借助AI21平台,企业可以构建自己的生成式AI的应用程序和服务。
Jamba算是世界上第一款基于Mamba的生产级别的模型。近日AI21 Labs发布Jamba 1.5系列开放模型,包括Jamba 1.5 Mini和Jamba 1.5 Large,这些模型建立在新颖的SSM-Transformer架构之上,某种意义上也算是AI技术的突破。
1.Jamba
Jamba 1.5创新的核心是其令人印象深刻的256K上下文窗口,这是目前最长上下文的开放模型。在单个实例中处理和理解如此大量信息的能力允许在各种任务中实现更细致、更全面和更准确的输出。Jamba 1.5 Mini可以在单个GPU上处理高达140K令牌的上下文长度。
图片
AI21 Labs的原话是:“大多数其他模型声称具有较长的上下文窗口,但无法在其上下文窗口的上限处保持相同的性能质量,而Jamba 1.5系列在其 256K上下文窗口的整个跨度内保持其长上下文处理。”
能够有效处理很长的上下文模型对于几乎企业级生成式的应用程序都至关重要。除了彻底、精确地总结和分析冗长的文档外,长上下文模型还通过消除连续分块和重复检索的需要,大大提高了RAG和代理工作流程的质量,并降低了它们的成本。
图片
上图为Jamba 1.5 Mini在2片A100 80GB GPU进行延迟评估, 而Jamba 1.5 Large则是在8块A100 80GB GPU上完成的,整体表现不错。下图为
图片
上表是Jamba和其他模型整体参数规模和实际运行激活参数规模的对比。Jamba 1.5 Large算是巅峰之作,它虽然具备398B的总参数,然而运行时仅仅有94B的活跃参数。而Jamba 1.5 Mini具有52B的总参数和12B的活跃参数。<MoE!!>
2.SSM-Transformer
Jamba 1.5基于混合SSM-Transformer构建,它结合了两种强大的AI 范式的优势:Mamba和Transformer。读者是否还记得Mamba系列?!Mamba为模型带来了卓越的效率和速度,而Transformer架构则在各种 NLP任务中保持高质量输出和稳健的性能。
图片
在Jamba的这个版本中,AI21试验了Mamba-2,这是Mamba的更快和改进版本,它的性能分别优于Mamba和Transformers。在实验中如上图所示在混合架构中,Mamba-1+Attention组合比Mamba-2+Attention效果更好,因此本次在Jamba-1.5-Large中使用Mamba-1。
当然在这个过程中还发现混合架构的性能优于纯Mamba-2。通过推测可能是因为在 Mamba层之间交错有完整的Tranformer层时,Mamba-2相对于Mamba-1 的优势就不那么明显了,Mamba-2虽然可以具备更大状态存储能力,然而Mamba-1通过Transformer可以汇集来自整个上下文的信息。
图片
上图为最早期的Jamba论文中Jamba Block的架构图。它为八层的架构,Mamba层和Transformer层的比例为7:1,每两层运用了MoE。
而Jamba-1.5-Large遵循相同的Jamba结构,但容量更大。它具有94B活动参数和398B总参数。它有9个区块,每个区块都有以下规格:
- l=8,每块具有8层
- a:m=1:7,这个比例是研究人员在对Jamba研究中发现的最佳的比例
- e=2,每两层都使用MoE代替单个MLP
- n=16,有16个专家,在每个Token上选择头部的两个,即K=2。
- 隐藏维度为8192
- Q为64头,KV为8头
3.ExpertsInt8
Jamba 1.5 Large 背后的关键创新之一是引入ExpertsInt8,这是一种专为专家混合 (MoE) 模型设计的新型量化技术。这种技术可以在A100 GPU 上使用,而FP8仅在H100上可以使用。当然这种技术最重要的特点如下:
- 快速的量化过程,只需几分钟,而不是几小时或几天
- 不依赖校准,增强了稳定性和易用性
- 能够将 BF16 用于大型激活,在最重要的地方保持精度
- 在 MoE 模型的 vLLM 量化技术中延迟最低,且不会牺牲质量
因为超过85%的模型权重在MoE层中,超过90%在 MoE或MLP层中,那么如何量化这些权重且同时仍然享受快速BF16内核的好处?研究人员将 MoE和MLP权重量化为INT8,将它们保存在INT8中,并在实际计算之前将它们返回BF16。
最重要的是,逆量化步骤直接发生在vLLM的 fused_moe内核内部。通过这种方式,逆量化过程增加的开销可以忽略不计,甚至导致比BF16更高的延迟。这个修改的fused_moe内核目前已经贡献给vLLM。
图片
上图为不同量化技术比较,显示1024个令牌上下文和128个令牌解码的端到端延迟。ExpertsInt8的性能与FP8相似,但应用起来快速而简单,并且依旧享受BF16激活,最重要的是能够适用于FP8不可用的A100 GPU。
图片
至于在一些通用基准的评测上,Jamba-1.5也获得不俗的表现。与同级别的模型相比,性能相似然而拥有更好的吞吐量和低延迟的能力。