大世界模型 (LWM)

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

方法

当前的语言模型在理解难以用文字描述的世界方面存在不足，并且在处理复杂的长篇任务时面临困难。视频序列提供了语言和静态图像所缺乏的宝贵时序信息，使其成为与语言进行联合建模的理想选择。这类模型可以同时理解人类文本知识和物理世界，从而为辅助人类提供更广泛的人工智能能力。然而，从数百万个视频和语言序列标记中学习面临着内存限制、计算复杂性和有限数据集等挑战。为了应对这些挑战，我们精心策划了一个包含多样化视频和书籍的大型数据集，利用RingAttention技术在长序列上进行可扩展训练，并逐步将上下文大小从4K增加到1M标记。本文作出以下贡献：(a) 最大上下文神经网络：我们在长视频和语言序列上训练了一个具有最大上下文大小的transformer，在困难的检索任务和长视频理解方面树立了新的基准。(b) 解决视觉-语言训练挑战的方案，包括使用掩码序列打包混合不同长度的序列、利用损失加权平衡语言和视觉，以及为长序列对话生成模型生成的问答数据集。(c) 高度优化的实现，包括RingAttention、掩码序列打包和其他关键特性，用于训练百万长度的多模态序列。(d) 完全开源的7B参数模型系列，能够处理超过1M标记的长文本文档(LWM-Text, LWM-Text-Chat)和视频(LWM, LWM-Chat)。这项工作为在大规模长视频和语言数据集上进行训练铺平了道路，以发展对人类知识和多模态世界的理解，并拓展更广泛的能力。

LWM 能力

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/7b9ebe80-79a0-457b-931d-ef94ad71a50a.png"/> <p> LWM能够在1M上下文中高精度检索事实。 </p> </div> <br /> <div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/281a31c1-11d1-4ec6-928c-20ac9adb8e2c.png"/> <p> LWM能够回答关于1小时YouTube视频的问题。 </p> </div> <br /> <div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/582a2496-99ea-4270-892e-90ec8511d0b5.png"/> <p> LWM能够与图像进行对话。 </p> </div> <br /> <div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/b6f93c7e-0eba-4b81-9f06-40861d76905e.png"/> <p> LWM能够根据文本生成视频和图像。 </p> </div>

设置

本代码库支持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 VM：

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 <脚本名称>.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='城市上空的烟花' \
    --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 上提出问题！

引用

如果您使用此代码库，或者发现我们的工作有价值，请引用：

@article{liu2023world,
    title={World Model on Million-Length Video and Language with RingAttention},
    author={Liu, Hao and Yan, Wilson and Zaharia, Matei and Abbeel, Pieter},
    journal={arXiv preprint},
    year={2024},
}
@article{liu2023ring,
    title={Ring Attention with Blockwise Transformers for Near-Infinite Context},
    author={Liu, Hao and Zaharia, Matei and Abbeel, Pieter},
    journal={International Conference on Learning Representations},
    year={2024}
}
@article{liu2023blockwise,
    title={Blockwise Parallel Transformer for Large Context Models},
    author={Liu, Hao and Abbeel, Pieter},
    journal={Advances in neural information processing systems},
    year={2023}
}