t al. (2023b)
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.
Llama 2: Open foundation and fine-tuned chat models.
arXiv preprint arXiv:2307.09288, 2023b.
Tworkowski et al. (2023)
Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Miłoś.
Focused transformer: Contrastive training for context scaling.
arXiv preprint arXiv:2307.03170, 2023.
Villegas et al. (2022)
Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan.
Phenaki: Variable length video generation from open domain textual description.
arXiv preprint arXiv:2210.02399, 2022.
Wang et al. (2023)
Yi Wang, Yinan He, Yizhuo Li, Kunchang Li, Jiashuo Yu, Xin Ma, Xinhao Li, Guo Chen, Xinyuan Chen, Yaohui Wang, et al.
Internvid: A large-scale video-text dataset for multimodal understanding and generation.
arXiv preprint arXiv:2307.06942, 2023.
Yang et al. (2018)
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning.
Hotpotqa: A dataset for diverse, explainable multi-hop question answering.
arXiv preprint arXiv:1809.09600, 2018.
Yu et al. (2022)
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al.
Scaling autoregressive models for content-rich text-to-image generation.
arXiv preprint arXiv:2206.10789, 2(3):5, 2022.
Zhang et al. (2023)
Renrui Zhang, Jiaming Han, Aojun Zhou, Xiangfei Hu, Shilin Yan, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao.
Llama-adapter: Efficient fine-tuning of language models with zero-init attention.
arXiv preprint arXiv:2303.16199, 2023.
Zheng et al. (2023)
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al.
Judging llm-as-a-judge with mt-bench and chatbot arena.
arXiv preprint arXiv:2306.05685, 2023.
Appendix A Further Details
Model Flops Utilization.
We trained our models using TPUv4-1024, which is approximately equivalent to 450 A100s, with a batch size of 8M using FSDP (Facebook, 2023) and BlockwiseRingAttention (Liu et al., 2024) for large contexts. Figure 8 shows the model FLOPS utilization (MFU) for each training stage. Blue color bars show language training and orange color bars show vision-language training. Our training achieves good MFUs even for very large context sizes.
Figure 8: High MFU training across sequence lengths. Model flops utilization (MFU) of each training stage for LWM-Text (top), and LWM / LWM-Chat (bottom)
Training Loss Curves.
Figure 9 and Figure 10 show the training loss curves for each stage of training the language and vision-language models respectively.
Figure 9: Training progress over multiple days for LWM-Text. Train loss curve for each training stage for LWM-Text models.
Figure 10: Training progress over multiple days for LWM. Train loss curve for each training stage for LWM and LWM-Chat models. Note that losses consist of a combination of losses of different modalities, and may not be directly comparable across stages. The sharp peak in the middle of 1K training is due to newly incporating EOF and EOV tokens into the vision codebook.
Training Hyperparameters.
See Appendix J
Scaling Inference.
We additionally scale our inference code to support million-length sequences by implementing RingAttention for decoding. Inference for such long sequences requires a minimum of v4-128 with a TPU mesh sharding of 32 tensor parallelism, and 4 sequence parallelism (ring dimension). We perform inference in pure single precision, where additional improvements can be made through techniques in scalability such as quantization.
Table 6: LWM-Text Training Stages
32K
128K
256K
512K
1M
Parameters
7B
7B
7B
7B
7B
Sequence Length
2152^{15}
2172^{