 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^{17}
2182^{18}
2192^{19}
2202^{20}
RoPE θ\theta
1M
10M
10M
25M
50M
Tokens per Batch
4M
4M
4M
4M
4M
Total Tokens
4.8B
12B
12B
3B
1.8B
Wall Clock
8h
45h
83h
47h
58h
Compute (TPU)
v4-512
v4-512
v4-512
v4-512
v4-512
Doc Length
10K-100K
100K-200K
200K-500K
500K-1M
1M+
Table 7: LWM-Text-Chat Training Details
128K
256K
512K
1M
Parameters
7B
7B
7B
7B
Sequence Length
2172^{17}
2182^{18}
2192^{19}
2202^{20}
RoPE θ\theta
10M
10M
25M
50M
Tokens per Batch
4M
4M
4M
4M
Total Tokens
1.2B
1.2B
1.2B
1.2B
Wall Clock
6h
10h
20h
40h
Compute (TPU)
v4-512
v4-512
v4-512
v4-512
Table 8: LWM and LWM-Chat Training Stages
1K
8K
Chat-32K
Chat-128K
Chat-1M
Parameters
7B
7B
7B
7B
7B
Sequence Length
2102^{10}
2132^{13}
2152^{15}
2172^{17}
2202^{20}
RoPE θ\theta
50M
50M
50M
50M
50M
Tokens per Batch
8M
8M
8M
8M
8M
Total Tokens
363B
107B
10B
3.5B
0.4B
Wall Clock
83h
32h
10h
6h
8h
Compute (TPU)
v4-1024
v4-1024
v4-1024
v4-1024
v4-1024
Appendix B Ablation Studies
B.1 Masked Sequence Packing
As mentioned in Section 4.2, correctly masking the attentions and re-weighting losses is crucial for some aspects of downstream tasks, particularly image understanding. Table 9 shows a comparison of our model with and without packing corrections. Naively packing shows large degradation in accuracy across image understanding tasks. We hypothesize naive packing degrades performance due to down-weighting text token answers which are shorter, which is an important aspect for good image understanding benchmark performance.
Table 9: Ablation study comparing standard independent packing and our masked sequence packing mechanisms across three tasks. Results show that masked sequence packing significantly improves performance across all tasks.
VQAv2
SQA
POPE
Standard independent packing
48.3
34.8
62.5
Masked sequence packing (Ours)
55.8
47.7
75.2
B.2 Mixing Synthetic And Chat Data
We additionally evaluate the our model on MT-Bench (Zheng et al., 2023) to test its conversation ability. Table 11 shows the MT-Bench scores of for each of our models. Table 11 illustrates the relationship between the mix of chat and fact retrieval tasks and the performance on MT-Bench score and Needle Retrieval accuracy. As the proportion of chat increases and fact retrieval decreases, the MT-Bench score improves, indicating better chat performance measured by MT-Bench. Conversely, Needle Retrieval accuracy decreases, suggesting a trade-off where increasing chat interaction capabilities may reduce the system’s precision in retrieving specific information or ’needles’ from input context. Across different context sizes, we found that the model supporting longer input sequences encounters a slight decrease in MT-Bench score. We hypothesize that this is because we chose to train with fewer examples on longer sequence training and can be improved by simply training on more data. In addition, this trade-off may be resolved by acquiring higher quality long-context chat data that is closer to the chat distribution of the UltraChat dataset.
Table 10: Results on MT-Bench across different context sizes. Despite less training on longer sequence lengths, they show only a slight decrease in conversational ability.
Model
MT-Bench
LWM-Text-Chat-128k
4.62
LWM-Text-Chat-256k
5
LWM-Text-Chat-512k
4.83
LWM-Text-Chat-1M
4.19
Table 11: Relationship between the mix of chat and fact retrieval tasks and the performance on MT-Bench score and Needle Retrieval accuracy.
Chat / QA Mix
MT-Bench
Needle Acc
0% / 100%
2.42
100%
40% / 60%
4.14
100%
70% / 30%
4.62
96%
90% / 10%
5.1
55%
100% / 0%
5.8
31%
Appendix C More Single-Needle Retrieval Results
F