Aria

😊 Hugging Face | 📄 论文 | 📚 博客 | 🌐 WebDemo | 🟣 Discord

简介

Aria 是一款多模态原生 MoE 模型。其特点包括：

• 在多种多模态和语言任务上表现处于行业领先水平，尤其在视频和文档理解方面表现出色；
• 具有长达 64K 个 token 的多模态上下文窗口；
• 每个 token 的激活参数量为 3.9B，从而实现快速推理速度和较低的微调成本。

最新消息

• [2025年1月20日] 🚀🚀🚀 Paddle 团队已在 PaddleMIX 中支持 Aria。

• [2024年12月15日] 我们发布了 Aria-Chat! 它针对开放式和多轮对话进行了优化，具有更高的可靠性及多语言支持。

• [2024年12月1日] 我们发布了 Aria 的基础模型（Aria-Base-8K 和 Aria-Base-64K）! 它们与本推理和微调代码库完全兼容。

• [2024年10月10日] 我们正式发布了 Aria!

快速入门

安装

pip install -e .
# 或者如果你打算参与项目贡献，可以安装包含开发依赖的版本
pip install -e .[dev] 

pip install grouped_gemm
pip install flash-attn --no-build-isolation

推理

Aria 总共有 25.3B 个参数，在单张 A100（80GB）GPU 上以 bfloat16 精度即可加载运行。

以下是一个使用 Hugging Face Transformers 调用 Aria 的代码示例。

import requests
import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor

model_id_or_path = "rhymes-ai/Aria"

model = AutoModelForCausalLM.from_pretrained(model_id_or_path, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)

processor = AutoProcessor.from_pretrained(model_id_or_path, trust_remote_code=True)

image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"

image = Image.open(requests.get(image_path, stream=True).raw)

messages = [
    {
        "role": "user",
        "content": [
            {"text": None, "type": "image"},
            {"text": "这是一张什么图片？", "type": "text"},
        ],
    }
]

text = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=text, images=image, return_tensors="pt")
inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
inputs = {k: v.to(model.device) for k, v in inputs.items()}

with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    output = model.generate(
        **inputs,
        max_new_tokens=500,
        stop_strings=["<|im_end|>"],
        tokenizer=processor.tokenizer,
        do_sample=True,
        temperature=0.9,
    )
    output_ids = output[0][inputs["input_ids"].shape[1]:]
    result = processor.decode(output_ids, skip_special_tokens=True)

print(result)

我们还提供了其他推理方法，例如利用 vLLM 来提升性能。更多详细信息请参阅 docs/inference.md。

食谱

请查看这些 推理示例，它们展示了如何在不同应用场景中使用 Aria，比如图表理解、PDF 阅读、视频理解等，同时支持 Hugging Face Transformers 和 vLLM 后端。

微调

⚠️ 关于微调的重要提示：由于 Aria 集成到 transformers 后权重映射发生了变化，训练代码需要特定版本才能正常工作：

• 请使用 transformers 版本 4.45.0
• 使用模型修订版 "4844f0b5ff678e768236889df5accbe4967ec845"

注意：为了获得最佳的微调效果，建议安装可选的 grouped_gemm 依赖：

pip install grouped_gemm

我们提供 LoRA 微调和全参数微调两种方式，并支持多种数据集类型：

• 单张图像数据集
• 多张图像数据集
• 视频数据集
• 代码数据集

若想快速尝试，请访问 examples 文件夹，选择其中一个微调示例。如果希望从基础模型开始微调（推荐用于大型数据集），请在配置文件中将以下模型路径（全参数 或 LoRA）

model_name_or_path: rhymes-ai/Aria
tokenizer_path: rhymes-ai/Aria

修改为对应的基础模型路径：

model_name_or_path: rhymes-ai/Aria-Base-64K # rhymes-ai/Aria-Base-8K
tokenizer_path: rhymes-ai/Aria-Base-64K # rhymes-ai/Aria-Base-8K

准备数据集

请参考 custom_dataset.md，了解如何准备您的数据集。

使用 LoRA 进行微调

准备好数据集后，按照以下步骤使用 LoRA 对 Aria 进行微调：

1. 打开配置文件 recipes/config_lora.yaml。找到 dataset_mixer 部分，并更新为您的数据集路径：

dataset_mixer:
  "path/to/dataset1": 1
  "path/to/dataset2": 0.5
  "path/to/dataset3": 2

关于数据集混合的说明：Aria 支持以不同的采样率组合多个数据集。在上面的例子中：

• dataset1 将被完全使用（权重 1）
• dataset2 将使用其数据的 50%（权重 0.5）
• dataset3 将被使用两次（权重 2）

2. 在一台 A100（80GB）或 H100（80GB）GPU 上运行以下命令开始微调过程：

python aria/train.py --config recipes/config_lora.yaml

3. 如果进行多 GPU 训练，请使用 accelerate 库：

accelerate launch --config_file recipes/accelerate_configs/zero2.yaml aria/train.py --config recipes/config_lora.yaml --num_processes [number_of_gpus]

• 可以从 recipes/accelerate_configs/ 中选择预配置的 accelerate 设置
• 根据您可用的 GPU 数量调整 --num_processes 参数
• 如需自定义配置，请参考 accelerate 文档

4. 使用微调后的模型进行推理：

   请参阅 支持 LoRA 的推理 了解如何使用微调后的模型进行推理。

全参数微调

除配置文件 recipes/config_full.yaml 外，其余步骤与 LoRA 微调流程相同。

全参数微调会占用更多显存，因此需要使用多张 GPU。以下命令已在 8 张 A100（80GB）GPU 上测试通过：

accelerate launch --config_file recipes/accelerate_configs/zero2.yaml aria/train.py --config recipes/config_full.yaml

如果遇到显存不足的错误，请尝试在配置文件中降低 per_device_train_batch_size，并相应调整 gradient_accumulation_steps 以保持有效的训练批次大小。

per_device_train_batch_size: 8
gradient_accumulation_steps: 2

不同数据集的显存消耗差异较大。通常，多图像和视频数据集需要更多的显存。可以通过调整 deepspeed_config 参数来优化显存使用，例如使用 zero_stage 3，并将参数和优化器卸载到 CPU 上。

deepspeed_config:
  gradient_accumulation_steps: auto
  gradient_clipping: auto
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero_stage: 3

使用您训练好的模型进行推理

首先，您需要从 ZeRO 1、2 或 3 的 DeepSpeed 检查点中提取 FP32 整合权重：

cd /path/to/your/output/dir
python zero_to_fp32.py . pytorch_model.bin

有关如何使用微调后的模型进行推理的说明，请参阅 inference.md。

引用

如果您觉得我们的工作对您有所帮助，请考虑引用：

@article{aria,
  title={Aria: 一个开放的多模态原生专家混合模型}, 
  author={李东旭、刘宇东、吴浩宁、王岳、沈志奇、曲博文、牛鑫尧、王国银、陈贝、李俊楠},
  year={2024},
  journal={arXiv 预印本 arXiv:2410.05993},
}

Aria 快速上手指南

Aria 是一款原生多模态混合专家（MoE）模型，在视频和文档理解任务上表现卓越。它支持 64K 长上下文窗口，每 token 仅激活 39 亿参数，兼具高性能与低推理成本。

环境准备

系统要求

• GPU: 推理至少需要 1 张 A100 (80GB) 或同等显存容量的 GPU（使用 bfloat16 精度）。
• 显存:
  • LoRA 微调：推荐 1 张 A100/H100 (80GB)。
  • 全量参数微调：推荐 8 张 A100 (80GB)，需配合 DeepSpeed ZeRO 优化。
• Python: 建议 Python 3.10+。

前置依赖

为确保微调功能正常运作，请严格遵循以下版本要求：

• transformers == 4.45.0
• 模型版本需指定 revision: 4844f0b5ff678e768236889df5accbe4967ec845

提示：国内开发者若遇到网络连接问题，建议在 pip 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源。

安装步骤

1. 克隆项目并安装基础依赖

   git clone https://github.com/rhymes-ai/Aria.git
   cd Aria
   pip install -e .
   # 如需开发贡献，可安装开发依赖
   # pip install -e .[dev]
   

2. 安装关键加速库 必须安装 grouped_gemm 和 flash-attn 以获得最佳性能：

   pip install grouped_gemm
   pip install flash-attn --no-build-isolation
   

基本使用

以下示例展示如何使用 Hugging Face Transformers 加载 Aria 模型并进行简单的图文对话推理。

代码示例

import requests
import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor

# 指定模型路径
model_id_or_path = "rhymes-ai/Aria"

# 加载模型 (自动映射设备，使用 bfloat16)
model = AutoModelForCausalLM.from_pretrained(
    model_id_or_path, 
    device_map="auto", 
    torch_dtype=torch.bfloat16, 
    trust_remote_code=True
)

# 加载处理器
processor = AutoProcessor.from_pretrained(model_id_or_path, trust_remote_code=True)

# 准备图片 (此处使用在线示例图片，也可替换为本地路径)
image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"
image = Image.open(requests.get(image_path, stream=True).raw)

# 构建对话消息
messages = [
    {
        "role": "user",
        "content": [
            {"text": None, "type": "image"},
            {"text": "what is the image?", "type": "text"},
        ],
    }
]

# 处理输入
text = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=text, images=image, return_tensors="pt")
inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
inputs = {k: v.to(model.device) for k, v in inputs.items()}

# 生成回复
with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    output = model.generate(
        **inputs,
        max_new_tokens=500,
        stop_strings=["<|im_end|>"],
        tokenizer=processor.tokenizer,
        do_sample=True,
        temperature=0.9,
    )
    output_ids = output[0][inputs["input_ids"].shape[1]:]
    result = processor.decode(output_ids, skip_special_tokens=True)

print(result)

进阶提示

• 更多后端: 如需更高吞吐量的推理，可参考官方文档使用 vLLM 后端。
• 应用场景: 项目 inference/notebooks 目录下提供了图表理解、PDF 阅读、视频理解等更多示例代码。