大世界模型 (LWM)

[项目] [论文] [模型]

大世界模型 (LWM) 是一种通用的大上下文多模态自回归模型。它基于 RingAttention 技术，使用包含多样化长视频和书籍的大型数据集进行训练，能够执行语言、图像和视频的理解与生成任务。

方法

当前的语言模型在理解那些难以用文字描述的世界方面存在不足，并且在处理复杂、长篇的任务时表现欠佳。视频序列提供了语言和静态图像所不具备的宝贵时间信息，因此非常适合与语言联合建模。这样的模型可以同时掌握人类的文本知识和物理世界的运行规律，从而赋予人工智能更广泛的能力来辅助人类。然而，从数百万个标记的视频和语言序列中学习面临着内存限制、计算复杂度高以及数据集有限等挑战。为应对这些难题，我们精心构建了一个包含多样化视频和书籍的大型数据集，采用 RingAttention 技术以可扩展的方式对长序列进行训练，并逐步将上下文长度从 4K 扩展到 1M 个标记。本文的主要贡献包括：(a) 拥有最大上下文长度的神经网络：我们在长视频和语言序列上训练了目前上下文长度最大的 Transformer 之一，在困难的检索任务和长视频理解方面树立了新的标杆。(b) 提出了克服视觉-语言训练挑战的解决方案，包括使用掩码序列打包技术混合不同长度的序列、通过损失加权平衡语言和视觉任务，以及利用模型生成的问答数据集进行长序列对话。(c) 开发了一套高度优化的实现方案，结合 RingAttention、掩码序列打包等关键技术，支持对百万级标记的多模态序列进行训练。(d) 完全开源了一组参数规模为 7B 的模型系列，这些模型能够处理超过 1M 个标记的长文本文档（LWM-Text、LWM-Text-Chat）和视频（LWM、LWM-Chat）。

这项工作为基于海量长视频和语言数据集的训练铺平了道路，有助于开发对人类知识和多模态世界的理解能力，进而推动更广泛的人工智能应用。

LWM 的能力

环境搭建

本代码库支持 Ubuntu 系统，尚未在 Windows 或 macOS 上测试过。我们建议使用 TPU 进行训练和推理，当然也可以使用 GPU。在 TPU 上，代码经过 Jax 的 Pallas 高度优化，能够在超大上下文长度下借助 RingAttention 实现极高的 MFU 值。而在 GPU 上，代码基于 XLA 架构，优化程度不及 TPU 版本。

安装依赖项的命令如下：

conda create -n lwm python=3.10
conda activate lwm
pip install -r gpu_requirements.txt

或者通过以下脚本设置 TPU 虚拟机环境：

sh tpu_requirements.sh

可用模型

现有仅处理语言的版本和语言-视觉联合处理的版本，上下文长度分别为 32K、128K、256K 和 1M 个标记。其中，视觉-语言模型仅提供 Jax 版本，而仅语言模型则同时支持 PyTorch 和 Jax。以下是可用模型的名称、对应的上下文长度及功能：

模型名称	上下文长度	语言或视觉-语言	对话版或基础版	链接
LWM-Text-Chat-128K	128K	语言	对话版	[Pytorch][Jax]
LWM-Text-Chat-256K	256K	语言	对话版	[Pytorch][Jax]
LWM-Text-Chat-512K	512K	语言	对话版	[Pytorch][Jax]
LWM-Text-Chat-1M	1M	语言	对话版	[Pytorch][Jax]
LWM-Text-128K	128K	语言	基础版	[Pytorch][Jax]
LWM-Text-256K	256K	语言	基础版	[Pytorch][Jax]
LWM-Text-512K	512K	语言	基础版	[Pytorch][Jax]
LWM-Text-1M	1M	语言	基础版	[Pytorch][Jax]
LWM-Chat-32K	32K	视觉-语言	对话版	[Jax]
LWM-Chat-128K	128K	视觉-语言	对话版	[Jax]
LWM-Chat-1M	1M	视觉-语言	对话版	[Jax]

代码结构

使用 scan_query_chunk_size 和 scan_key_chunk_size 来控制自注意力机制分块计算中的块大小。使用 scan_mlp_chunk_size 来控制前馈网络分块计算中的块大小。通过设置 scan_attention=True 和 scan_mlp=True，可以启用或禁用自注意力机制和前馈网络中的分块计算。

您可以使用 mesh_dim=dp, fsdp, tp, sp 来控制并行度和 RingAttention 的配置。它是一个由逗号分隔的 4 个整数组成的字符串，分别表示数据并行、全切片数据并行、张量并行和序列并行的数量。

例如，mesh_dim='1,64,4,1' 表示 1 个数据并行、64 个全切片数据并行、4 个张量并行和 1 个序列并行。而 mesh_dim='1,1,4,64' 则表示 1 个数据并行、1 个全切片数据并行、4 个张量并行以及 64 个序列并行，用于 RingAttention。

运行 Jax 模型

在这一部分，我们提供了如何运行每个提供的脚本的说明。对于每个脚本，您可能需要在脚本开头描述的变量中填写您自己的路径和值。

要运行以下脚本，请使用 bash <script_name>.sh：

• 语言模型训练：bash scripts/run_train_text.sh
• 视觉-语言模型训练：bash scripts/run_train_vision_text.sh
• 单针评估（语言模型）：bash scripts/run_eval_needle.sh
• 多针评估（语言模型）：bash scripts/run_eval_needle_multi.sh
• 采样图像（视觉-语言模型）：bash scripts/run_sample_image.sh
• 采样视频（视觉-语言模型）：bash scripts/run_sample_video.sh
• 图像/视频理解（视觉-语言模型）：bash scripts/run_vision_chat.sh

默认情况下，mesh_dim 参数会将所有设备置于 tp（张量并行）模式。对于较长的序列，您可能希望加入 sp，即 mesh_dim 中的最后一个维度。

在运行针刺评估时，您可能需要根据模型调整脚本中的 theta 和 max_sequence_length 参数。以下是各模型的正确取值：

	LWM-Text-128K / LWM-Text-Chat-128K	LWM-Text-256K / LWM-Text-Chat-256K	LWM-Text-512K / LWM-Text-Chat-512K	LWM-Text-1M / LWM-Text-Chat-1M
theta	10000000	10000000	25000000	50000000
max_sequence_length	131072	262144	524288	1048576

以下是填写脚本（run_sample_video.sh）的示例：

#! /bin/bash

export SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
export PROJECT_DIR="$( cd -- "$( dirname -- "$SCRIPT_DIR" )" &> /dev/null && pwd )"
cd $PROJECT_DIR
export PYTHONPATH="$PYTHONPATH:$PROJECT_DIR"

export llama_tokenizer_path="LargeWorldModel/LWM-Text-1M"
export vqgan_checkpoint="/path/to/ckpt/folder/vqgan"
export lwm_checkpoint="params::/path/to/ckpt/folder/params"

python3 -u -m lwm.vision_generation \
    --prompt='Fireworks over the city' \
    --output_file='fireworks.mp4' \
    --temperature_image=1.0 \
    --temperature_video=1.0 \
    --top_k_image=8192 \
    --top_k_video=1000 \
    --cfg_scale_image=5.0 \
    --cfg_scale_video=1.0 \
    --vqgan_checkpoint="$vqgan_checkpoint" \
    --n_frames=8 \
    --mesh_dim='!1,1,-1,1' \
    --dtype='fp32' \
    --load_llama_config='7b' \
    --update_llama_config="dict(sample_mode='vision',theta=50000000,max_sequence_length=32768,scan_attention=False,scan_query_chunk_size=128,scan_key_chunk_size=128,scan_mlp=False,scan_mlp_chunk_size=8192,scan_layers=True)" \
    --load_checkpoint="$lwm_checkpoint" \
    --tokenizer="$llama_tokenizer_path"
read

针尖藏草堆数据

运行 python scripts/create_needle_data.py

运行 PyTorch 模型

目前，PyTorch 推理仅支持文本和文本聊天模型。PyTorch 模型可以作为 Hugging Face 的 LlamaForCausalLM 模型加载。运行 python scripts/sample_pyt.py 进行采样。您可能需要单独安装 torch。

文档

有关代码库的更多详细信息，请参阅 data.md 和 sharding.md。data.md 提供了关于数据处理的详细信息，而 sharding.md 则介绍了分片和并行化方面的细节。

如遇问题

本项目基于 RingAttention 的 代码库，并添加了视觉-语言训练所需的必要功能。训练和推理已在 TPUv3 和 TPUv4 上进行了测试。

如果您遇到任何错误，请在 GitHub 上提交 issue！

引用

如果您使用了本代码库，或者以其他方式认为我们的工作有价值，请引用以下文献：

@article{liu2023world,
    title={World Model on Million-Length Video and Language with RingAttention},
    author={刘浩、颜威尔逊、扎哈里亚·马泰伊、阿贝尔·皮特},
    journal={arXiv 预印本},
    year={2024},
}
@article{liu2023ring,
    title={Ring Attention with Blockwise Transformers for Near-Infinite Context},
    author={刘浩、扎哈里亚·马泰伊、阿贝尔·皮特},
    journal={国际学习表征会议},
    year={2024}
}
@article{liu2023blockwise,
    title={Blockwise Parallel Transformer for Large Context Models},
    author={刘浩、阿贝尔·皮特},
    journal={神经信息处理系统进展},
    year={2023}
}

许可证

LWM 的代码采用 Apache 2.0 许可证发布。更多详情请参阅 LICENSE。模型则采用 Llama-2 许可证发布。

LWM (Large World Model) 快速上手指南

LWM 是一个通用的大上下文多模态自回归模型，支持语言、图像和视频的理解与生成。它基于 RingAttention 技术训练，能够处理长达 100 万 token 的上下文（如长视频或书籍）。

环境准备

• 操作系统: 推荐 Ubuntu。代码库尚未在 Windows 或 macOS 上进行测试。
• 硬件加速:
  • TPU: 强烈推荐。代码针对 Jax 的 Pallas 进行了高度优化，在使用 RingAttention 处理超大上下文时可实现极高的 MFU（模型浮点利用率）。
  • GPU: 支持。基于 XLA 实现，但优化程度不如 TPU 版本。
• Python 版本: 3.10

安装步骤

1. 创建虚拟环境

conda create -n lwm python=3.10
conda activate lwm

2. 安装依赖

根据使用的硬件选择以下一种方式：

方案 A：使用 GPU

pip install -r gpu_requirements.txt

(注：如果原文未提供国内镜像源配置，建议在 pip 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 以加速下载)

方案 B：使用 TPU

sh tpu_requirements.sh

基本使用

LWM 提供多种预训练模型，包括纯文本模型（支持 PyTorch 和 Jax）和视 - 语多模态模型（仅支持 Jax）。

1. 选择模型

从 Hugging Face 获取模型权重。常用模型如下：

• 长文本对话: LWM-Text-Chat-1M (支持 1M 上下文)
• 视频/图像理解: LWM-Chat-1M (仅 Jax, 支持 1M 上下文)

2. 运行示例 (Jax 版本)

以下示例展示如何使用 Jax 版本生成视频。请确保已设置好检查点路径。

创建一个脚本文件（例如 run_demo.sh），内容如下：

#! /bin/bash

export SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
export PROJECT_DIR="$( cd -- "$( dirname -- "$SCRIPT_DIR" )" &> /dev/null && pwd )"
cd $PROJECT_DIR
export PYTHONPATH="$PYTHONPATH:$PROJECT_DIR"

# 请替换为实际的模型路径和检查点路径
export llama_tokenizer_path="LargeWorldModel/LWM-Text-1M"
export vqgan_checkpoint="/path/to/ckpt/folder/vqgan"
export lwm_checkpoint="params::/path/to/ckpt/folder/params"

python3 -u -m lwm.vision_generation \
    --prompt='Fireworks over the city' \
    --output_file='fireworks.mp4' \
    --temperature_image=1.0 \
    --temperature_video=1.0 \
    --top_k_image=8192 \
    --top_k_video=1000 \
    --cfg_scale_image=5.0 \
    --cfg_scale_video=1.0 \
    --vqgan_checkpoint="$vqgan_checkpoint" \
    --n_frames=8 \
    --mesh_dim='!1,1,-1,1' \
    --dtype='fp32' \
    --load_llama_config='7b' \
    --update_llama_config="dict(sample_mode='vision',theta=50000000,max_sequence_length=32768,scan_attention=False,scan_query_chunk_size=128,scan_key_chunk_size=128,scan_mlp=False,scan_mlp_chunk_size=8192,scan_layers=True)" \
    --load_checkpoint="$lwm_checkpoint" \
    --tokenizer="$llama_tokenizer_path"

运行脚本：

bash run_demo.sh

3. 运行示例 (PyTorch 版本 - 仅限文本)

如果您只需进行长文本推理，可以使用 PyTorch 版本，直接作为 Hugging Face 模型加载。

确保已安装 torch，然后运行采样脚本：

python scripts/sample_pyt.py

(注意：运行前需在脚本内部配置具体的模型路径和提示词)

关键参数说明

• mesh_dim: 控制并行策略。格式为 dp,fsdp,tp,sp。对于长序列推理，建议利用最后一个维度 sp (Sequence Parallelism) 来启用 RingAttention。例如：mesh_dim='1,1,4,64' 表示使用 64 路序列并行。
• theta 和 max_sequence_length: 在进行长上下文评估（如 Needle In A Haystack）时，需根据模型上下文大小调整这两个参数（详见原仓库表格）。