Scaling Granite Code Models to 128K Context

We continue pretrain the full attention Granite code base models using sequence parallelism¹¹1https://github.com/jzhang38/EasyContext (Li et al., 2021) by gradually increasing its RoPE base frequency without using any sparse or linear attention. Specifically, we continue pretrain Granite Code 3B/8B base models using the original pretraining data used in Mishra et al. (2024) but with repository-level file packing and per-language context length upsampling, that we found to be critical for long-context continual pretraining. This continued training stage focused on a curated selection of programming languages, such as Python, C, C++, Go, Java, JavaScript, and TypeScript, as in Pinnaparaju et al. (2024).

To create long-context data, we develop a new approach that packs files from the same repository together, arranging them to prioritize semantic dependencies. We identify these dependencies by analyzing file imports and create a directed acyclic graph, where each file is a node and edges represent API imports between files. After breaking any cycles in the graph, we perform a topological sort to establish an ordering of files based on their semantic dependencies. We then organize the files in a repository by placing documentation and build files first, followed by the ordered set of files with semantic dependencies, and finally the remaining non-connected files. These non-connected files are arranged according to their folder structure, using a depth-first search to traverse the repository. Finally, we determine the dominant programming language of a repository based on file extensions and presence of build files, to organise repo-ordered files by programming languages.

The documents’ lengths and their source domains/languages are two closely related confounding factors in data engineering because long data usually come from particular sources. Thus, in addition to repository-level file packing, we artificially oversampled longer document sequences on a per-language basis to ensure the quantity of long sequences, thereby improving the overall quality of our training data corpus, as in  Fu et al. (2024); Yu (2023). In particular, we downsample documents under 4096 tokens to a rate of 10%, which we find to ensure a sufficient number of total tokens and documents. The total number of documents within the training corpus after processing is 173,336 with a mean length of 73,451.

We adjust the RoPE base frequency, introduced in Xiong et al. (2023), to support long context windows up to 128K where the base model itself is trained on 2K/4K context length. For training, we adopt a progressive approach where we doubled the context window until it reached the desired length of 128K. We train for 500 steps with a batch size of 32 and search for the optimal RoPE theta and learning rate for each iteration. For RoPE theta, we finf optimal values of 100K, 250K, 500K, 2M, and 10M for context windows of 8K, 16K, 32K, 64K, and 128K, respectively. We train with data parallelism and Flash Attention 2 until 64K tokens and then used Ring Attention  (Liu et al., 2023a) to reach 128K tokens. The final models are trained for an extra 4B tokens which is only 0.1% of original pretraining data.

Our training data for long context instruct models consists of a combination of permissively licensed data used in training the original Granite code instruct models (Mishra et al., 2024), in addition to synthetically generated code instruction datasets tailored for solving long context problems. Specifically, the 128K long context instruct models are derived by further finetuning the long context base models on a mix of short and long context data as follows.

Short-Context Instruction Data. Our short context instruction data consists of a combination of CommitPackFT (Muennighoff et al., 2023), MathInstruct²²2We removed GSM8K-RFT and Camel-Math from MathInstruct due to unknown or NC license. (Yue et al., 2023), MetaMathQA (Yu et al., 2023), Glaive-Code-Assistant-v3³³3https://huggingface.co/datasets/glaiveai/glaive-code-assistant-v3, Self-OSS-Instruct-SC2⁴⁴4https://huggingface.co/datasets/bigcode/self-oss-instruct-sc2-exec-filter-50k, Glaive-Function-Calling-v2⁵⁵5https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2, NL2SQL⁶⁶6https://huggingface.co/datasets/bugdaryan/sql-create-context-instruction, HelpSteer (Wang et al., 2023b), OpenPlatypus⁷⁷7https://huggingface.co/datasets/garage-bAInd/Open-Platypus (Lee et al., 2023), and a few synthetically generated datasets for API calling  (Basu et al., 2024), and multi-turn code interactions with execution feedback.

Long-Context Instruction Data. The long context instruction data was synthetically generated by bootstrapping the pretraining data. For each repository-level file-packed document, we created a multi-turn dataset where the instructions within each sample were human-designed for the purpose of enhancing the long-context performance in specific tasks like generation, retrieval and translation. The responses were either parsed semantically from the original document or generated using Granite-8b-Code-Instruct-4K. The dataset first parses the document into classes, methods, and stand-alone functions. It then requests and extracts the implementations of a random subset of the extracted functions/methods (up to 5 per file in the document) and then asks for an explanation of that implementation using available documentation. Additionally, it generates instructions for implementing the sampled functions (methods) based on the remaining documentation and code with the function excluded. These questions and instructions were repeated for different functions until the desired length was achieved.

By exposing the model to both short and long context data, we aim to enhance its long context capability without sacrificing code generation performance at short input context. For finetuning, we use a multiturn loss mask for each sample, as in Wang et al. (2023a). This is particularly important as our finetuning data corpus consists of instruction-response pairs with multiple turns. However, when composing a sequence, we append an EOS token after each response from the model to prevent runaway generation during inference. We followed the same training parameters that produced our previous short-context instruct models (Mishra et al., 2024): 128 global batch size, 2e-5 learning rate, a noise multiplier of 5 for input embeddings, and padding-free transformers.

Long Code Completion. Long Code Completion (LCC) (Guo et al., 2023) tests a model’s ability to predict the next line of code from long repository-based context for Python, Java, and C#. While the benchmark’s context length spans 1/2K through 8K+ tokens, it is heavily weighted around 2K tokens. Thus, following Bai et al. (2024) and Rozière et al. (2023), we rebalance this dataset for equal representation with each context length bucket (<4K, 2 – 4K, 4 – 8K, 8K+), where each bucket has 100 samples when possible.

RepoBench-P. Like LCC, RepoBench-P (Liu et al., 2023c) tests the model’s next line code completion ability for long-context input. We follow the methodology in (Bai et al., 2024) by selecting the Cross-File-First data but then we rebalance the buckets based on the Starcoder tokenizer used for training out Granite code models.

RepoQA. RepoQA (Liu et al., 2024) is an advanced Needle-in-the-Haystack test that focuses on testing LLMs’ capabilities on long-context code understanding and retrieval. Specifically, given a long chunk of source code and a precise function description, and the model is asked to find the function in the context that corresponds to the description. This benchmark focuses on retrieving 10 needle functions from each of 5 languages x 10 repositories (500 sub-tasks/tests) with a set context size of 16K tokens.

Key Retrieval. This is a synthetic benchmark that tests the model’s ability to find and execute a Python function buried within high-quality, syntactically correct Python code. As proposed in Rozière et al. (2023), we took the Code Contest finetuning dataset from Li et al. (2022) and concatenated Python solutions around the key function. We then asked the model to return the output of the key function by emulating a Python interpreter shell. We created sequences of lengths of 512 tokens and key offsets of 512 tokens.

HumanEvalPack. To evaluate model performance at short-context length, we adopt HumanEvalPack (Muennighoff et al., 2023), which extends Python problems of Humaneval Benchmark to five additional commonly used programming languages, namely JavaScript, Java, Go, C++, Rust to test three coding tasks (generation, explanation and fixing). We evaluate our long-context models in a zero-shot manner using greedy decoding with completion format for the base models, and with instruction template for the instruction-tuned models.

Table 1 and Table 2 show the results of Granite 3B/8B code models before and after long-context extension on LCC and RepoBench-P benchmarks respectively. Prior Granite code models with 2K/4K support fail to generate meaningful completions on long sequences. On the other hand, across all the context length (4K to 32K), models scaled to handle long contexts up to 128K achieve significantly higher performance. This demonstrates that long contexts are informative for code completion, and long-context Granite code models are able to effectively leverage this information to improve their generations on both benchmarks.

In Table 3, we compare the performance of Granite code base models to their counterparts prior to long-context extension. Our long-context models exhibit strong retrieval performance across different matching thresholds, while the short context versions mostly fail in finding the needle function successfully. The absolute differences averaged over 5 programming languages are very significant, e.g., $+38.6\%$ for Granite 8B model with a matching threshold of 0.8. By looking at the score distribution across different programming languages, we can see that both models are doing best at Python, with 8B model consistently outperforming the 3B model. This result shows that our long-context Granite code models can better understand natural language description before retrieval, which aligns with the use of advanced code search in many practical situations.

Table 4 compares the performance of long-context instruct models to their short-context counterparts on RepoQA benchmark. As can be seen, our long-context instruct models significantly outperforms short-context versions on all 5 programming languages across different similarity thresholds. As an illustration, figure 1 demonstrates the difference between short and long-context models at similarity threshold of 0.5, where the performance of both 3B and 8B instruct models with 2K/4K context length support fails to achieve a retrieval accuracy of more than 2% across 5 languages (on average 0.6% vs 61.6% for 8B instruct model). We attribute the improvements to the knowledge learned from newly introduced synthetic long data for instruction tuning.

In Figure 2, we investigate key retrieval performance of our long-context instruct models on a synthetic benchmark built on top of Python solutions around a key function from Code Contest finetuning dataset (Li et al., 2022). Note that this retrieval task is analogous to the famous famous Needle-in-a-Haystack test, albeit tailored to code models. As can be seen from Figure 2, our 8B instruct model before long-context extension only exhibit strong retrieval performance up to 4K length, i.e., on the sequence length they were originally trained on. On the other hand, our context scaling demonstrates a perfect-all-green performance though we tend to view that this level of retrieval is relatively easy for long-context code LLMs.

While our long-context models are very effective on long sequences, we observe that our long-context scaling does not significantly change the short-context generic capability on standard code synthesis benchmarks consisting of short sequences. Table 5 summarizes the results on HumanEvalPack, where we find only an average $\sim$1% degradation for the pass@1 metric on 3B and 8B models respectively. We also test the HumanEval-Python performance in Figure 3 and observe that long context extension has any noticeable performance degradation. Interestingly, we notice improvements in HumanEval performance of long-context instruct models, which we attribute to our new long-context synthetic data added to instruction tuning. To summarize, while long-context extension comes at a minimal cost for short sequences, we believe this cost is more than offset by the potential of handling long sequences for many real downstream applications.