了解 Jamba 模型

对于 Transformers 系列模型，由于这些模型在自然语言处理（NLP）领域内非常流行，很多深度学习框架都提供了对它们的支持和实现，比如 TensorFlow 的 transformers 库和 PyTorch 的 transformers 库（由 Hugging Face 提供）。这些库中的模型通常包括 BERT、GPT、RoBERTa、Transformer-XL、T5 等。

这里使用 ai21labs/Jamba-v0.1 用作测试模型

1. 下载模型，查看模型的大小

模型大概有96GB大小，单个GPU可能无法加载，需要使用多个GPU进行加载

2. 使用 AutoConfig 获取模型参数

这个时候可以放心运行，不加载模型，只加载模型的配置文件，运行输出结果

从结果来看，其实是模型目录下的 config.json 文件里面的内容

从提供的输出来看，test01.py 脚本很可能是调用了 Hugging Face transformers 库中的 AutoConfig.from_pretrained 方法，用于加载一个名为 ai21labs/Jamba-v0.1 的模型配置。这个配置属于 JambaConfig 类，这是一个定制化的模型配置，可能是为 Jamba 模型特别设计的。

输出中包含了大量的配置信息，其中一些关键字段说明如下：

• "architectures": 显示了模型的架构，这里是 JambaForCausalLM，表明这个配置是为因果语言模型设计的。
• "hidden_size": 隐藏层的大小是 4096。
• "num_hidden_layers": 表明模型有 32 层隐藏层。
• "num_attention_heads": 模型使用了 32 个注意力头。
• "vocab_size": 词汇表的大小是 65536。
• "transformers_version": 指明了 transformers 库的版本是 4.39.2。

此外，还包含了一些针对 Jamba 模型特有的配置项，如 mamba_conv_bias、mamba_d_conv、mamba_d_state 等，这些配置项可能与模型中特殊的卷积或状态处理机制有关。

num_hidden_layers 指的是该模型架构中的隐藏层总数。这个参数就代表了模型中相应部分的层数。

了解这个参数 num_hidden_layers 对后面的模型加载有用。

3. 使用 CPU 加载模型

当你的 GPU 内存不够使用的时候，使用 CPU 加载模型，如果内存不够(此模型大概需要 212GB)，可能需要增大交换区（linux),可以获取模型参数

运行结果如下：

4. 使用多个 GPU 加载模型

当你的单个 GPU 内存不够使用的时候，可能需要多个GPU，这里我们使用了 8 个 RTX 4090 的卡, 实际上对于这个模型，不需要这么多。

从CPU的参数来看，可以得知，有如下权重参数：

1. model.embed_tokens
2. model.norm
3. model.final_layernorm
4. lm_head
5. model.layers (总共32层)

我们可以将Transformer模型的每一层分配给一个GPU，每个GPU可能是多层

代码如下：

输出的结果：

5. 多 GPU int8 和 int4 量化部署

量化后的 GPU 内存,需要的更少。但加载的时间会更长，下面的代码是int4的

如果需要int8 的，只需要修改：

运行结果：