Hymba 1.5B Base

🚀 Hymba-1.5B-Base

Hymba-1.5B-Base 是一个基础的文本生成模型，可应用于各种自然语言生成任务。它采用混合架构，具备商业使用许可，能为开发者提供高效且实用的文本处理解决方案。

🚀 快速开始

环境搭建

由于 Hymba-1.5B-Base 使用了 FlexAttention，它依赖于 Pytorch2.5 及其他相关依赖项，我们提供两种方式来搭建环境：

• [本地安装] 使用我们提供的 setup.sh 脚本安装相关包（支持 CUDA 12.1/12.4）：

wget --header="Authorization: Bearer YOUR_HF_TOKEN" https://huggingface.co/nvidia/Hymba-1.5B-Base/resolve/main/setup.sh
bash setup.sh

• [Docker] 我们提供了一个已安装 Hymba 所有依赖项的 Docker 镜像。你可以使用以下命令下载 Docker 镜像并启动容器：

docker pull ghcr.io/tilmto/hymba:v1
docker run --gpus all -v /home/$USER:/home/$USER -it ghcr.io/tilmto/hymba:v1 bash

与 Hymba-1.5B-Base 对话

环境搭建完成后，你可以使用以下脚本与我们的模型进行对话：

from transformers import LlamaTokenizer, AutoModelForCausalLM, AutoTokenizer, AutoModel
import torch

# 加载分词器和模型
repo_name = "nvidia/Hymba-1.5B-Base"

tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)

# 与 Hymba 对话
prompt = input()
inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
outputs = model.generate(**inputs, max_length=64, do_sample=False, temperature=0.7, use_cache=True)
response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)

print(f"模型回复: {response}")

✨ 主要特性

• 混合架构：模型采用 Mamba 和注意力头并行运行的混合架构，结合了标准注意力头和 Mamba 头的优势。
• 元令牌机制：在每个提示前添加一组可学习的元令牌，有助于提高模型的效能。
• KV 缓存共享：模型在两层之间和单层内的头之间共享 KV 缓存，提高了计算效率。
• 滑动窗口注意力：90% 的注意力层采用滑动窗口注意力，优化了注意力计算。
• 商业可用：该模型可用于商业用途。

📦 安装指南

环境搭建

由于 Hymba-1.5B-Base 使用了 FlexAttention，它依赖于 Pytorch2.5 及其他相关依赖项，我们提供两种方式来搭建环境：

• [本地安装] 使用我们提供的 setup.sh 脚本安装相关包（支持 CUDA 12.1/12.4）：

wget --header="Authorization: Bearer YOUR_HF_TOKEN" https://huggingface.co/nvidia/Hymba-1.5B-Base/resolve/main/setup.sh
bash setup.sh

• [Docker] 我们提供了一个已安装 Hymba 所有依赖项的 Docker 镜像。你可以使用以下命令下载 Docker 镜像并启动容器：

docker pull ghcr.io/tilmto/hymba:v1
docker run --gpus all -v /home/$USER:/home/$USER -it ghcr.io/tilmto/hymba:v1 bash

💻 使用示例

基础用法

from transformers import LlamaTokenizer, AutoModelForCausalLM, AutoTokenizer, AutoModel
import torch

# 加载分词器和模型
repo_name = "nvidia/Hymba-1.5B-Base"

tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)

# 与 Hymba 对话
prompt = input()
inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
outputs = model.generate(**inputs, max_length=64, do_sample=False, temperature=0.7, use_cache=True)
response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)

print(f"模型回复: {response}")

📚 详细文档

模型概述

Hymba-1.5B-Base 是一个基础的文本生成模型，可用于各种自然语言生成任务。该模型采用混合架构，结合了 Mamba 和注意力头的优势。元令牌机制有助于提高模型的效能，同时模型在两层之间和单层内的头之间共享 KV 缓存，提高了计算效率。90% 的注意力层采用滑动窗口注意力，优化了注意力计算。该模型可用于商业用途。

模型开发者：NVIDIA

模型训练时间：Hymba-1.5B-Base 于 2024 年 9 月 1 日至 2024 年 11 月 10 日期间进行训练。

许可证：该模型根据 NVIDIA 开放模型许可协议 发布。

模型架构

💡 我们在 GitHub 上发布了 Hymba 的最小实现，以帮助开发者理解并在自己的模型中实现其设计原则。查看链接：barebones-hymba。

Hymba-1.5B-Base 的模型嵌入大小为 1600，有 25 个注意力头，MLP 中间维度为 5504，总共 32 层，16 个 SSM 状态，3 个全注意力层，其余为滑动窗口注意力。与标准 Transformer 不同，Hymba 中的每个注意力层都有标准注意力头和 Mamba 头的并行混合组合。此外，它使用了分组查询注意力（GQA）和旋转位置嵌入（RoPE）。

该架构的特点包括：

• 并行处理：在同一层内融合注意力头和 SSM 头，对相同输入进行并行和互补处理。
• 元令牌机制：在输入序列前添加元令牌，与所有后续令牌交互，存储重要信息并减轻注意力中的“强制关注”负担。
• 高效计算：集成跨层 KV 共享和全局 - 局部注意力，进一步提高内存和计算效率。

性能亮点

• Hymba-1.5B-Base 在性能上优于所有小于 20 亿参数的公开模型。

微调 Hymba

LMFlow 是一个用于微调大语言模型的完整管道。以下步骤提供了一个使用 LMFlow 微调 Hymba-1.5B-Base 模型的示例：

1. 使用 Docker

docker pull ghcr.io/tilmto/hymba:v1
docker run --gpus all -v /home/$USER:/home/$USER -it ghcr.io/tilmto/hymba:v1 bash

2. 安装 LMFlow

git clone https://github.com/OptimalScale/LMFlow.git
cd LMFlow
conda create -n lmflow python=3.9 -y
conda activate lmflow
conda install mpi4py
pip install -e .

3. 使用以下命令微调模型

cd LMFlow
bash ./scripts/run_finetune_hymba.sh

使用 LMFlow，你还可以在自定义数据集上微调模型。你只需要将数据集转换为 LMFlow 数据格式。除了全量微调，你还可以使用 DoRA、LoRA、LISA、Flash Attention 等加速技术高效微调 Hymba。更多详情，请参考 LMFlow for Hymba 文档。

评估

我们使用 LM Evaluation Harness 来评估模型。评估命令如下：

git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness
git fetch --all --tags
git checkout tags/v0.4.4  # squad 完成任务与最新版本不兼容
cd lm-evaluation-harness
pip install -e .

lm_eval --model hf --model_args pretrained=nvidia/Hymba-1.5B-Base,dtype=bfloat16,trust_remote_code=True \
     --tasks mmlu \
     --num_fewshot 5 \
     --batch_size 1 \
     --output_path ./hymba_HF_base_lm-results \
     --log_samples 

lm_eval --model hf --model_args pretrained=nvidia/Hymba-1.5B-Base,dtype=bfloat16,trust_remote_code=True \
     --tasks arc_easy,arc_challenge,piqa,winogrande,hellaswag \
     --num_fewshot 0 \
     --batch_size 1 \
     --output_path ./hymba_HF_base_lm-results \
     --log_samples 

lm_eval --model hf --model_args pretrained=nvidia/Hymba-1.5B-Base,dtype=bfloat16,trust_remote_code=True \
     --tasks squad_completion \
     --num_fewshot 1 \
     --batch_size 1 \
     --output_path ./hymba_HF_base_lm-results \
     --log_samples 

局限性

该模型在包含有毒语言、不安全内容和社会偏见的数据上进行训练，这些数据最初是从互联网上爬取的。因此，模型可能会放大这些偏见并返回有毒回复，尤其是在接收到有毒提示时。即使提示本身不包含任何明确的冒犯性内容，模型也可能生成不准确的答案，遗漏关键信息，或包含无关或冗余的文本，产生社会不可接受或不良的文本。

测试表明，该模型容易受到越狱攻击。如果在 RAG 或代理环境中使用此模型，我们建议采用强大的输出验证控制，以确保用户控制的模型输出的安全风险与预期用例一致。

伦理考量

NVIDIA 认为可信 AI 是一项共同责任，我们已经制定了政策和实践，以支持各种 AI 应用的开发。当按照我们的服务条款下载或使用该模型时，开发者应与内部模型团队合作，确保该模型符合相关行业和用例的要求，并解决不可预见的产品滥用问题。 请 在此 报告安全漏洞或 NVIDIA AI 相关问题。

引用

@misc{dong2024hymbahybridheadarchitecturesmall,
      title={Hymba: A Hybrid-head Architecture for Small Language Models}, 
      author={Xin Dong and Yonggan Fu and Shizhe Diao and Wonmin Byeon and Zijia Chen and Ameya Sunil Mahabaleshwarkar and Shih-Yang Liu and Matthijs Van Keirsbilck and Min-Hung Chen and Yoshi Suhara and Yingyan Lin and Jan Kautz and Pavlo Molchanov},
      year={2024},
      eprint={2411.13676},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2411.13676}, 
}

🔧 技术细节

Hymba-1.5B-Base 的模型嵌入大小为 1600，有 25 个注意力头，MLP 中间维度为 5504，总共 32 层，16 个 SSM 状态，3 个全注意力层，其余为滑动窗口注意力。与标准 Transformer 不同，Hymba 中的每个注意力层都有标准注意力头和 Mamba 头的并行混合组合。此外，它使用了分组查询注意力（GQA）和旋转位置嵌入（RoPE）。该架构在同一层内融合注意力头和 SSM 头，对相同输入进行并行和互补处理；通过在输入序列前添加元令牌，与所有后续令牌交互，存储重要信息并减轻注意力中的“强制关注”负担；集成跨层 KV 共享和全局 - 局部注意力，进一步提高内存和计算效率。

📄 许可证

该模型根据 NVIDIA 开放模型许可协议 发布。