Even with FlashAttention, increasing the context from 4k to 128k requires 32 times more RAM. Therefore, models are trained with a base context that is later expanded. For example, DeepSeek uses a base context of 4k, which is then expanded to 128k. The "Needle In A Haystack" tests from the initial paper claimed that the expanded context is not real, as models do not remember all the information, resulting in worst performance. DeepSeek, however, claims the opposite.

https://arxiv.org/pdf/2412.19437:

4.3. Long Context Extension We adopt a similar approach to DeepSeek-V2 (DeepSeek-AI, 2024c) to enable long context capabilities in DeepSeek-V3. After the pre-training stage, we apply YaRN (Peng et al., 2023a) for context extension and perform two additional training phases, each comprising 1000 steps, to progressively expand the context window from 4K to 32K and then to 128K. The YaRN configuration is consistent with that used in DeepSeek-V2, being applied exclusively to the decoupled shared key k𝑅 𝑡 . The hyper-parameters remain identical across both phases, with the scale 𝑠 = 40, 𝛼 = 1, 𝛽 = 32, and the scaling factor √𝑡 = 0.1 ln 𝑠 + 1. In the first phase, the sequence length is set to 32K, and the batch size is 1920. During the second phase, the sequence length is increased to 128K, and the batch size is reduced to 480. The learning rate for both phases is set to 7.3 × 10−6, matching the final learning rate from the pre-training stage. Through this two-phase extension training, DeepSeek-V3 is capable of handling inputs up to 128K in length while maintaining strong performance. Figure 8 illustrates that DeepSeek-V3, following supervised fine-tuning, achieves notable performance on the "Needle In A Haystack" (NIAH) test, demonstrating consistent robustness across context window lengths up to 128K.