# TinyLlama

小身板,大能量。

当大家都在研究大模型(LLM)参数规模达到百亿甚至千亿级别的同时,小巧且兼具高性能的小模型开始受到研究者的关注。

小模型在边缘设备上有着广泛的应用,如智能手机、物联网设备和嵌入式系统,这些边缘设备通常具有有限的计算能力和存储空间,它们无法有效地运行大型语言模型。因此,深入探究小型模型显得尤为重要。

接下来我们要介绍的这两项研究,可能满足你对小模型的需求。

TinyLlama-1.1B

来自新加坡科技设计大学(SUTD)的研究者近日推出了 TinyLlama,该语言模型的参数量为 11 亿,在大约 3 万亿个 token 上预训练而成。

  • 论文地址:https://arxiv.org/pdf/2401.02385.pdf
  • 项目地址:https://github.com/jzhang38/TinyLlama/blob/main/README_zh-CN.md

TinyLlama 以 Llama 2 架构和分词器(tokenizer)为基础,这意味着 TinyLlama 可以在许多基于 Llama 的开源项目中即插即用。此外,TinyLlama 只有 11 亿的参数,体积小巧,适用于需要限制计算和内存占用的多种应用。

该研究表示仅需 16 块 A100-40G 的 GPU,便可在 90 天内完成 TinyLlama 的训练。

该项目从上线开始,持续受到关注,目前星标量达到 4.7K。

轻量级模型-合集1_数据

TinyLlama 模型架构详细信息如下所示:

轻量级模型-合集1_人工智能_02

训练细节如下:

轻量级模型-合集1_卷积_03

研究者表示,这项研究旨在挖掘使用较大数据集训练较小模型的潜力。他们重点探究在用远大于扩展定律(scaling law)建议的 token 数量进行训练时,较小模型的行为表现。

具体来说,该研究使用大约 3 万亿个 token 训练具有 1.1B 个参数的 Transformer (仅解码器)模型。据了解,这是第一次尝试使用如此大量的数据来训练具有 1B 参数的模型。

尽管规模相对较小,但 TinyLlama 在一系列下游任务中表现相当出色,它的性能显著优于同等大小的现有开源语言模型。具体来说,TinyLlama 在各种下游任务中都超越了 OPT-1.3B 和 Pythia1.4B 。

此外,TinyLlama 还用到了各种优化方法,如 flash attention 2、FSDP( Fully Sharded Data Parallel )、 xFormers 等。

在这些技术的加持下,TinyLlama 训练吞吐量达到了每 A100-40G GPU 每秒 24000 个 token。例如,TinyLlama-1.1B 模型对于 300B token 仅需要 3,456 A100 GPU 小时,而 Pythia 为 4,830 小时,MPT 为 7,920 小时。这显示了该研究优化的有效性以及在大规模模型训练中节省大量时间和资源的潜力。

TinyLlama 实现了 24k tokens / 秒 / A100 的训练速度,这个速度好比用户可以在 8 个 A100 上用 32 小时训练一个具有 11 亿参数、220 亿 token 的 chinchilla-optimial 的模型。同时,这些优化也大大减少了显存占用,用户可以把 11 亿参数的模型塞入 40GB 的 GPU 里面还能同时维持 16k tokens 的 per-gpu batch size。只需要把 batch size 改小一点, 你就可以在 RTX 3090/4090 上面训练 TinyLlama。 

轻量级模型-合集1_人工智能_04

实验中,该研究主要关注具有纯解码器架构的语言模型,包含大约 10 亿个参数。具体来说,该研究将 TinyLlama 与 OPT-1.3B、Pythia-1.0B 和 Pythia-1.4B 进行了比较。

TinyLlama 在常识推理任务上的性能如下所示,可以看出 TinyLlama 在许多任务上都优于基线,并获得了最高的平均分数。

轻量级模型-合集1_卷积_05

此外,研究者在预训练期间跟踪了 TinyLlama 在常识推理基准上的准确率,如图 2 所示,TinyLlama 的性能随着计算资源的增加而提高,在大多数基准中超过了 Pythia-1.4B 的准确率。

轻量级模型-合集1_卷积_06

表 3 表明,与现有模型相比,TinyLlama 表现出了更好的问题解决能力。

轻量级模型-合集1_卷积_07

手快的网友已经开始整活了:运行效果出奇得好,在 GTX3060 上运行,能以 136 tok / 秒的速度运行。 

「确实是快!」

小模型 LiteLlama

由于 TinyLlama 的发布,SLM(小型语言模型)开始引起广泛关注。德克萨斯工农大学的 Xiaotian Han 发布了 SLM-LiteLlama。它有 460M 参数,由 1T token 进行训练。这是对 Meta AI 的 LLaMa 2 的开源复刻版本,但模型规模显著缩小。

项目地址:https://huggingface.co/ahxt/LiteLlama-460M-1T

LiteLlama-460M-1T 在 RedPajama 数据集上进行训练,并使用 GPT2Tokenizer 对文本进行 token 化。作者在 MMLU 任务上对该模型进行评估,结果如下图所示,在参数量大幅减少的情况下,LiteLlama-460M-1T 仍能取得与其他模型相媲美或更好的成绩。

轻量级模型-合集1_英伟达_08

以下为该模型的性能表现,更详细内容请参阅:

https://huggingface.co/datasets/open-llm-leaderboard/details_ahxt__llama2_xs_460M_experimental

轻量级模型-合集1_卷积_09

面对规模大幅缩小的 LiteLlama,有网友好奇,它是否能够在 4GB 的内存上运行。如果你也想知道,不如亲自试试看吧。

参考链接:

https://twitter.com/_akhaliq/status/1744009616562819526

https://twitter.com/XiaotianHan1/status/1743824496916656275

https://twitter.com/abacaj/status/1743303507594097136

 # 轻量级模型设计与部署总结

本文在一些经典的轻量级网络论文阅读理解的基础上,对不同硬件平台部署轻量级模型相关的经验进行了总结。

轻量级网络的核心是在尽量保持精度的前提下,从体积和速度两方面对网络进行轻量化改造。关于如何手动设计轻量级网络的研究,目前还没有广泛通用的准则,只有一些指导思想,和针对不同芯片平台(不同芯片架构)的一些设计总结,建议大家从经典论文中吸取指导思想和建议,然后自己实际做各个硬件平台的部署和模型性能测试。

一些关键字定义及理解

计算量 FLOPs

  • FLOPsfloating point operations 指的是浮点运算次数,理解为计算量,可以用来衡量算法/模型时间的复杂度
  • FLOPS:(全部大写),Floating-point Operations Per Second,每秒所执行的浮点运算次数,理解为计算速度, 是一个衡量硬件性能/模型速度的指标,即一个芯片的算力。
  • MACCsmultiply-accumulate operations,乘-加操作次数,MACCs 大约是 FLOPs 的一半。将 w[0]∗x[0]+... 视为一个乘法累加或 1 个 MACC

内存访问代价 MAC

MACMemory Access Cost 内存访问代价。指的是输入单个样本(一张图像),模型/卷积层完成一次前向传播所发生的内存交换总量,即模型的空间复杂度,单位是 Byte

FLOPs 和 MAC 的计算方式,请参考我之前写的文章神经网络模型复杂度分析:https://github.com/HarleysZhang/cv_note/blob/79740428b6162630eb80ed3d39052cac52f60c32/9-model_deploy/B-%E7%A5%9E%E7%BB%8F%E7%BD%91%E7%BB%9C%E6%A8%A1%E5%9E%8B%E5%A4%8D%E6%9D%82%E5%BA%A6%E5%88%86%E6%9E%90.md。

GPU 内存带宽

  • GPU 的内存带宽决定了它将数据从内存 (vRAM) 移动到计算核心的速度,是比 GPU 内存速度更具代表性的指标。
  • GPU 的内存带宽的值取决于内存和计算核心之间的数据传输速度,以及这两个部分之间总线中单独并行链路的数量

NVIDIA RTX A4000 建立在 NVIDIA Ampere 架构之上,其芯片规格如下所示:

轻量级模型-合集1_人工智能_10

英伟达 RTX A4000 芯片规格

A4000 芯片配备 16 GB 的 GDDR6 显存、256 位显存接口(GPU 和 VRAM 之间总线上的独立链路数量),因为这些与显存相关的特性,所以A4000 内存带宽可以达到 448 GB/s

Latency and Throughput

参考英伟达-Ashu RegeDirector of Developer Technology 的 ppt 文档 An Introduction to Modern GPU Architecture:https://download.nvidia.com/developer/cuda/seminar/TDCI_Arch.pdf。

深度学习领域延迟 Latency 和吞吐量 Throughput的一般解释:

  • 延迟 (Latency): 人和机器做决策或采取行动时都需要反应时间。延迟是指提出请求与收到反应之间经过的时间。大部分人性化软件系统(不只是 AI 系统),延迟都是以毫秒来计量的。
  • 吞吐量 (Throughput): 在给定创建或部署的深度学习网络规模的情况下,可以传递多少推断结果。简单理解就是在一个时间单元(如:一秒)内网络能处理的最大输入样例数

CPU 是低延迟低吞吐量处理器;GPU 是高延迟高吞吐量处理器

Volatile GPU Util

一般,很多人通过 nvidia-smi 命令查看 Volatile GPU Util 数据来得出 GPU 利用率,但是!关于这个利用率(GPU Util),容易产生两个误区

  • 误区一: GPU 的利用率 = GPU 内计算单元干活的比例。利用率越高,算力就必然发挥得越充分。
  • 误区二: 同条件下,利用率越高,耗时一定越短。

但实际上,GPU Util 的本质只是反应了,在采样时间段内,一个或多个内核(kernel)在 GPU 上执行的时间百分比,采样时间段取值 1/6s~1s

原文为 Percent of time over the past sample period during which one or more kernels was executing on the GPU. The sample period may be between 1 second and 1/6 second depending on the product. 来源文档 nvidia-smi.txt(https://developer.download.nvidia.com/compute/DCGM/docs/nvidia-smi-367.38.pdf)

通俗来讲,就是,在一段时间范围内, GPU 内核运行的时间占总时间的比例。比如 GPU Util 是 69%,时间段是 1s,那么在过去的 1s 中内,GPU 内核运行的时间是 0.69s。如果 GPU Util 是 0%,则说明 GPU 没有被使用,处于空闲中。

也就是说它并没有告诉我们使用了多少个 SM 做计算,或者程序有多“忙”,或者内存使用方式是什么样的,简而言之即不能体现出算力的发挥情况

GPU Util 的本质参考知乎文章-教你如何继续压榨GPU的算力(https://zhuanlan.zhihu.com/p/346389176) 和 stackoverflow 问答(https://link.zhihu.com/?target=https%3A//stackoverflow.com/questions/40937894/nvidia-smi-volatile-gpu-utilization-explanation)。

英伟达 GPU 架构

GPU 设计了更多的晶体管(transistors)用于数据处理(data process)而不是数据缓冲(data caching)和流控(flow control),因此 GPU 很适合做高度并行计算highly parallel computations)。同时,GPU 提供比 CPU 更高的指令吞吐量内存带宽instruction throughput and memory bandwidth)。

CPU 和 GPU 的直观对比图如下所示

轻量级模型-合集1_英伟达_11

distribution of chip resources for a CPU versus a GPU

图片来源 CUDA C++ Programming Guide(https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html)

最后简单总结下英伟达 GPU 架构的一些特点:

  • SIMT (Single Instruction Multiple Threads) 模式,即多个 Core 同一时刻只能执行同样的指令。虽然看起来与现代 CPU 的 SIMD(单指令多数据)有些相似,但实际上有着根本差别。
  • 更适合计算密集数据并行的程序,原因是缺少 Cache 和 Control

2008-2020 英伟达 GPU 架构进化史如下图所示:

轻量级模型-合集1_英伟达_12

2008-2020 英伟达 GPU 架构进化史

另外,英伟达 GPU 架构从 2010 年开始到 2020 年这十年间的架构演进历史概述,可以参考知乎的文章-英伟达GPU架构演进近十年,从费米到安培(https://zhuanlan.zhihu.com/p/413145211)。

GPU 架构的深入理解可以参考博客园的文章-深入GPU硬件架构及运行机制()。

CNN 架构的理解

在一定的程度上,网络越深越宽,性能越好。宽度,即通道(channel)的数量,网络深度,及 layer 的层数,如 resnet18 有 18 层网络。注意我们这里说的和宽度学习一类的模型没有关系,而是特指深度卷积神经网络的(通道)宽度。

  • 网络深度的意义CNN 的网络层能够对输入图像数据进行逐层抽象,比如第一层学习到了图像边缘特征,第二层学习到了简单形状特征,第三层学习到了目标形状的特征,网络深度增加也提高了模型的抽象能力。
  • 网络宽度的意义:网络的宽度(通道数)代表了滤波器(3 维)的数量,滤波器越多,对目标特征的提取能力越强,即让每一层网络学习到更加丰富的特征,比如不同方向、不同频率的纹理特征等。

手动设计高效 CNN 架构建议

一些结论

  1. 分析模型的推理性能得结合具体的推理平台(常见如:英伟达 GPU、移动端 ARM CPU、端侧 NPU 芯片等);目前已知影响 CNN 模型推理性能的因素包括: 算子计算量 FLOPs(参数量 Params)、卷积 block 的内存访问代价(访存带宽)、网络并行度等。但相同硬件平台、相同网络架构条件下, FLOPs 加速比与推理时间加速比成正比。
  2. 建议对于轻量级网络设计应该考虑直接 metric(例如速度 speed),而不是间接 metric(例如 FLOPs)。
  3. FLOPs 低不等于 latency 低,尤其是在有加速功能的硬体 (GPUDSP 与 TPU)上不成立,得结合具硬件架构具体分析。
  4. 不同网络架构的 CNN 模型,即使是 FLOPs 相同,但其 MAC 也可能差异巨大。
  5. 大部分时候,对于 GPU 芯片,Depthwise 卷积算子实际上是使用了大量的低 FLOPs、高数据读写量的操作。这些具有高数据读写量的操作,加上多数时候 GPU芯片算力的瓶颈在于访存带宽,使得模型把大量的时间浪费在了从显存中读写数据上,导致 GPU 的算力没有得到“充分利用”。结论来源知乎文章-FLOPs与模型推理速度:https://zhuanlan.zhihu.com/p/122943688

一些建议

  1. 在大多数的硬件上,channel 数为 16 的倍数比较有利高效计算。如海思 351x 系列芯片,当输入通道为 4 倍数和输出通道数为 16 倍数时,时间加速比会近似等于 FLOPs 加速比,有利于提供 NNIE 硬件计算利用率。
  2. 低 channel 数的情况下 (如网路的前几层),在有加速功能的硬件(NPU 芯片)上使用普通 convolution 通常会比 separable convolution 有效率。(来源 MobileDets 论文(https://medium.com/ai-blog-tw/mobiledets-flops%E4%B8%8D%E7%AD%89%E6%96%BClatency-%E8%80%83%E9%87%8F%E4%B8%8D%E5%90%8C%E7%A1%AC%E9%AB%94%E7%9A%84%E9%AB%98%E6%95%88%E6%9E%B6%E6%A7%8B-5bfc27d4c2c8)
  3. shufflenetv2 论文(https://arxiv.org/pdf/1807.11164.pdf) 提出的四个高效网络设计的实用指导思想: G1同样大小的通道数可以最小化 MAC、G2-分组数太多的卷积会增加 MAC、G3-网络碎片化会降低并行度、G4-逐元素的操作不可忽视。
  4. GPU 芯片上 3\times 33\times 3 卷积非常快,其计算密度(理论运算量除以所用时间)可达 1\times 11\times 1 和 5\times 55\times 5 卷积的四倍。(来源 RepVGG 论文(https://zhuanlan.zhihu.com/p/344324470))
  5. 从解决梯度信息冗余问题入手,提高模型推理效率。比如 CSPNet(https://arxiv.org/pdf/1911.11929.pdf) 网络。
  6. 从解决 DenseNet 的密集连接带来的高内存访问成本和能耗问题入手,如 VoVNet(https://arxiv.org/pdf/1904.09730.pdf) 网络,其由 OSAOne-Shot Aggregation,一次聚合)模块组成。

轻量级网络模型部署总结

在阅读和理解经典的轻量级网络 mobilenet 系列、MobileDetsshufflenet 系列、cspnetvovnetrepvgg 等论文的基础上,做了以下总结:

  1. 低算力设备-手机移动端 cpu 硬件,考虑 mobilenetv1(深度可分离卷机架构-低 FLOPs)、低 FLOPs 和 低MACshuffletnetv2channel_shuffle 算子在推理框架上可能不支持)。
  2. 专用 asic 硬件设备-npu 芯片(地平线 x3/x4 等、海思 3519、安霸cv22 等),目标检测问题考虑 cspnet 网络(减少重复梯度信息)、repvgg(直连架构-部署简单,网络并行度高有利于发挥 GPU 算力,量化后有掉点风险) 。
  3. 英伟达 gpu 硬件-t4 芯片,考虑 repvgg 网络(类 vgg 卷积架构-高并行度带来高速度、单路架构省显存/内存)。

MobileNet block (深度可分离卷积 blockdepthwise separable convolution block)在有加速功能的硬件(专用硬件设计-NPU 芯片)上比较没有效率。

这个结论在 CSPNet(https://arxiv.org/pdf/1911.11929.pdf) 和 MobileDets(https://arxiv.org/pdf/2004.14525.pdf) 论文中都有提到。

除非芯片厂商做了定制优化来提高深度可分离卷积 block 的计算效率,比如地平线机器人 x3 芯片对深度可分离卷积 block 做了定制优化。

下表是 MobileNetv2 和 ResNet50 在一些常见 NPU 芯片平台上做的性能测试结果。

轻量级模型-合集1_人工智能_13

深度可分离卷积和常规卷积模型在不同NPU芯片平台上的性能测试结果

以上,均是看了轻量级网络论文总结出来的一些不同硬件平台部署轻量级模型的经验,实际结果还需要自己手动运行测试。

参考资料

------