Cobra: 将Mamba扩展为多模态大型语言模型以实现高效推理

(由SDXL生成，提示词：背景为熔岩的可爱眼镜蛇，卡通风格)

发布

[5/23] 已实施对提示格式的修复，解决了生成过程中的异常行为。请确保使用最新的检查点以适应新的提示格式！

[5/18] 评估代码已上传！

[3/27] 一个直观的演示，展示了我们的Cobra模型和LLaVA v1.5的生成速度。 <img src="https://yellow-cdn.veclightyear.com/ab5030c0/816dd9cd-ef8d-49cb-8962-39f4d0c6773b.gif">

[3/24] 试试我们的在线演示！

[3/24] 训练和推理代码已发布。

[3/24] 我们的模型权重现已可用。

[3/21] 仓库已创建。

安装 | 使用 | 预训练模型 | 训练VLMs | 许可证

安装

本仓库使用Python 3.10构建，但应该与任何Python >= 3.8版本向后兼容。我们需要PyTorch 2.1或更高版本，安装说明可在此处找到。本仓库已使用PyTorch 2.1.0和Torchvision 0.16.0进行了全面测试和开发。

一旦正确安装了PyTorch，您可以通过可编辑安装在本地安装此包（或通过pip install git+https://github.com/h-zhao1997/cobra）：

git clone https://github.com/h-zhao1997/cobra
cd cobra
pip install -e .

# 安装mamba和其他包
pip install packaging ninja
pip install "mamba-ssm<2.0.0"
pip install causal-conv1d

# 验证Ninja --> 应返回退出代码"0"
ninja --version; echo $?

如果在安装过程中遇到任何问题，请提交GitHub问题。

使用

安装完成后，加载和运行预训练的cobra模型进行推理很简单：

import requests
import torch

from PIL import Image
from pathlib import Path

from cobra import load

hf_token = Path(".hf_token").read_text().strip()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# 如果您的GPU不支持bf16
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

# 加载预训练的VLM（本地路径或从HF Hub自动下载的ID）
model_id = "cobra+3b"
vlm = load(model_id, hf_token=hf_token)
vlm.to(device, dtype=dtype)

# 下载图像并指定提示
image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png"
image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
user_prompt = "这张图片中发生了什么？"

# 构建提示
prompt_builder = vlm.get_prompt_builder()
prompt_builder.add_turn(role="human", message=user_prompt)
prompt_text = prompt_builder.get_prompt()

# 生成！
generated_text = vlm.generate(
    image,
    prompt_text,
    use_cache=True,
    do_sample=True,
    temperature=0.4,
    max_new_tokens=512,
)

有关与我们的VLMs交互的完整终端CLI，请查看scripts/generate.py。

训练VLMs

预训练数据集

可以通过scripts/preprocess.py中的自动下载脚本下载LLaVa v1.5指令数据集：

# 下载`llava-v1.5-instruct`（指令微调）图像和语言数据（包括额外的后处理）
python scripts/preprocess.py --dataset_id "llava-v1.5-instruct" --root_dir <数据根目录路径>

# （如果您还希望下载显式视觉-语言对齐数据）
python scripts/preprocess.py --dataset_id "llava-laion-cc-sbu-558k" --root_dir <数据根目录路径>

LVIS-Instruct-4V和LRV-Instruct也可以通过scripts/additional-datasets中的脚本下载。

模型配置和训练脚本

以下是如何按照我们论文中的训练方法在单节点的8个GPU上训练Cobra：

# 在仓库根目录下运行
torchrun --standalone --nnodes 1 --nproc-per-node 8 scripts/pretrain.py \
  --model.vision_backbone_id "dinosiglip-vit-so-384px" \
  --model.image_resize_strategy "resize-naive" \
  --model.llm_backbone_id "mamba-2.8b-zephyr" \
  --model.type "cobra+3b" \
  --model.finetune_global_batch_size 128 \
  --model.finetune_per_device_batch_size 8 \
  --dataset.type "llava-lvis4v-lrv"

引用

如果您在工作中发现我们的代码或模型有用，请引用我们的论文：

@article{zhao2024cobra,
      title={Cobra: Extending Mamba to Multi-Modal Large Language Model for Efficient Inference}, 
      author={Han Zhao and Min Zhang and Wei Zhao and Pengxiang Ding and Siteng Huang and Donglin Wang},
      year={2024},
      eprint={2403.14520},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}