Source-linked AI summary

CodeGemma: Open Code Models Based on Gemma

CodeGemma Team, Heri Zhao, Jeffrey Hui, Joshua Howland, Nam Nguyen, Siqi Zuo, Andrea Hu, Christopher A. Choquette-Choo, Jingyue Shen, Joe Kelley, Kshitij Bansal, Luke Vilnis, Mateo Wirth, Paul Michel, Peter Choy, Pratik Joshi, Ravin Kumar, Sarmad Hashmi, Shubham Agrawal, Zhitao Gong, Jane Fine, Tris Warkentin, Ale Jakse Hartman, Bin Ni, Kathy Korevec, Kelly Schaefer, Scott Huffman

arXiv:2406.11409v2cs.CLcs.AI

TL;DR

CodeGemma addresses the need for capable open code models across completion, generation, language understanding, and reasoning tasks. It extends Gemma with large-scale code pretraining, fill-in-the-middle training, repository-aware packing, and instruction tuning, producing strong coding and language results while making the 2B model suitable for latency-constrained deployment.

  • Problem

    Open code models need to support diverse coding tasks while retaining natural-language and reasoning capabilities for practical deployment.

  • Method

    CodeGemma extends Gemma with large-scale code pretraining, fill-in-the-middle training, repository-aware multi-file packing, and instruction tuning using mathematical and synthetic coding data.

  • Results

    CodeGemma retains most base-Gemma natural-language capabilities, excels at mathematical reasoning, and delivers high-quality code completion, while the 2B model is considerably faster than comparison models.

  • Takeaways & Limitations

    The released models are designed for effective real-world use in latency-constrained settings, including IDEs, local environments, and memory-constrained applications.

  • Takeaways & Limitations

    GSM8K highlights limitations in models’ multi-step reasoning despite the simplicity of its problems.

Abstract

from arXiv · show

This paper introduces CodeGemma, a collection of specialized open code models built on top of Gemma, capable of a variety of code and natural language generation tasks. We release three model variants. CodeGemma 7B pretrained (PT) and instruction-tuned (IT) variants have remarkably resilient natural language understanding, excel in mathematical reasoning, and match code capabilities of other open models. CodeGemma 2B is a state-of-the-art code completion model designed for fast code infilling and open-ended generation in latency-sensitive settings.

Introduction

CodeGemma is a collection of open code models built on Gemma, spanning 7B pretrained and instruction-tuned variants plus a specialized 2B model. The models target code completion and generation while retaining natural-language and reasoning capabilities.

  • Model family: CodeGemma extends Gemma pretrained models with more than 500 to 1000 billion tokens of primarily code using the Gemma architectures.The released collection includes 7B pretrained and instruction-tuned models, plus a specialized 2B model.
  • Goals: CodeGemma models achieve state-of-the-art code performance in completion and generation while maintaining strong understanding and reasoning skills at scale.
  • Releases: Version 1.1 updates the pretrained 2B and instruction-tuned 7B models with a recommended quality improvement unless speed is critically important.
  • Report scope: The report covers training additions, broad academic and real-world evaluations, strengths, limitations, deployment recommendations, and differences between versions.

Pretraining

CodeGemma pretraining combines large-scale filtered code and natural-language data with fill-in-the-middle training and repository-aware multi-file packing. The models support both PSM and SPM infilling formats.

  • Data: CodeGemma v1.0 models use 500 billion tokens, while the 2B v1.1 model uses 1 trillion tokens; 2B training is all code and 7B training mixes 80% code with 20% natural language.The corpus uses publicly available repositories and is deduplicated and filtered for evaluation contamination and sensitive data.
  • Fill-in-the-Middle: Pretrained models use fill-in-the-middle training with an 80% FIM rate in most models and 90% for pretrained 2B v1.1.
  • Fill-in-the-Middle: The models support both Prefix-Suffix-Middle and Suffix-Prefix-Middle modes for code infilling.
  • Multi-file Packing: Repository-level training examples co-locate relevant source files using dependency graph-based packing and unit test-based lexical packing.
  • Multi-file Packing: Dependency-based packing groups repository files, estimates relationships through imports and suffix matching, removes unimportant cyclic edges, and linearizes files topologically.
  • Multi-file Packing: Files outside the dependency graph are alphabetized, with unit tests packed beside their implementations.

Instruction Tuning

Instruction tuning combines mathematical datasets, synthetic coding instructions, supervised fine-tuning, and reinforcement learning to improve coding and reasoning behavior. The training examples span algebra, competition mathematics, grade-school problems, and filtered code question-answer pairs.

  • Instruction Tuning: The instruction-tuned 7B v1.1 model changes the reinforcement-learning algorithm and synthetic-data generation details relative to version 1.0.
  • Mathematics Datasets: Supervised fine-tuning uses diverse mathematics datasets to enhance coding models’ mathematical reasoning capabilities.
  • Mathematics Datasets: The MATH dataset contains 12,500 competition problems with step-by-step solutions for answer derivation and explanation generation.
  • Mathematics Datasets: GSM8K contains 8,500 grade-school math problems intended to test multi-step reasoning.The paper notes that these problems highlight model limitations despite their simplicity.
  • Mathematics Datasets: Synthetic Mathematical Data A is a programmatically generated algebra dataset designed to improve solving long algebra problems.
  • Mathematics Datasets: Training experiments indicate that the diverse mathematics datasets significantly boost code-generation performance.
  • Coding Dataset: Synthetic code instruction generation creates question-answer pairs for supervised fine-tuning and reinforcement learning from human feedback.
  • Coding Dataset: Generated coding pairs are post-filtered by an LLM for helpfulness and correctness.

Evaluation

CodeGemma is evaluated across code completion, code generation, natural language understanding, and mathematical reasoning. The results show strong coding performance, retained language capability, and competitive mathematical reasoning, with the 2B model offering substantially faster inference.

  • Evaluation scope: The evaluation covers code completion and generation, natural language understanding, question answering, and mathematical reasoning using automated benchmarks.
  • Code Completion: HumanEval Infilling single-line and multi-line metrics are used to evaluate code completion against other fill-in-the-middle-aware models.
  • Code Completion: The 2B pretrained model performs on par with comparison models for code completion while often running nearly twice as fast during inference.The paper attributes the speedup to base Gemma architectural decisions.
  • Code Completion: CodeGemma validates infilling with random masked snippets involving cross-file dependencies and retests generated files using an internal benchmark with the same methodology.
  • Python Coding: CodeGemma models perform significantly better than base Gemma models on coding-domain tasks.
  • Language Capability: CodeGemma retains most natural-language capabilities of base Gemma, with PT and IT outperforming Mistral 7B by 7.2% and Llama-2 13B by 19.1%.
  • Language Capability: CodeGemma excels at mathematical reasoning compared with similarly sized 7B code models on GSM8K and MATH.

Practical Considerations

CodeGemma offers model-size choices for deployment: the 2B model targets speed and memory-constrained use, while the 7B models prioritize quality and broader coding capability.

  • Deployment trade-offs: The 2B model is considerably faster than comparison models for latency-sensitive applications such as code completion.Evaluations found no significant measured quality compromise, with performance at least matching other open models in its class on code infilling.
  • Deployment trade-offs: CodeGemma 2B is suited to IDEs, local environments, and applications with memory constraints.
  • Deployment trade-offs: The 7B models surpass baseline Gemma on coding tasks while maintaining strong natural-language comprehension.Their larger inference memory requirement makes them particularly suitable for hosted environments where model quality is prioritized.
  • Scope: The report discusses the limitations and benefits of using an open model in Gemma's Responsible Deployment section.

Inference Recommendations

Inference recommendations distinguish pretrained code-completion prompting from instruction-tuned prompting, emphasizing exact formatting and deployment-specific stopping behavior.

  • Pretrained models: Pretrained models should receive prompts formatted for function completion, docstring generation, or import suggestion.An optional file path is recommended in the prompt format.
  • Output handling: Stopping should match the deployment setting, with truncation at a generated FIM sentinel token as the straightforward option.
  • Pretrained models: Figure 4 illustrates PSM prompting, including a carriage return and no spaces after the suffix.
  • Instruction-tuned models: Instruction-tuned models can also be prompted using Gemma's <start_of_turn> and <end_of_turn> formatting.

Conclusion

CodeGemma is an open family of Gemma-based models specialized for coding, combining code completion and generation with retained language capabilities and practical deployment targets.

  • Contribution: CodeGemma is a collection of open coding models built on top of the openly available Gemma family.
  • Contribution: The models target state-of-the-art code completion and generation while retaining natural-language capabilities from the base models.
  • Deployment: CodeGemma is optimized for latency-constrained real-world deployment while delivering high-quality code completion across tasks and languages.
  • Evaluation: Table 2 evaluates single-line and multi-line completion against other FIM-aware code models using HumanEval Infilling tasks and measured generation time.
  • Evaluation: Table 4 reports multilingual coding capability on BabelCode-translated HumanEval and MBPP datasets, distinguishing instruction-tuned models.

Contributions and Acknowledgments

The report acknowledges broad contributions spanning core research, safety, evaluation, infrastructure, partnerships, and program support.

  • Core Contributors: Core contributors are acknowledged for the central work on CodeGemma.
  • Additional Support: The report specifically thanks contributors in other specialty areas, including central support and checkpoint conversions.
  • Ethics and Safety: Separate acknowledgments recognize ethics and safety contributors.
  • Evaluations: Evaluation contributors are credited for supporting the reported assessments.
  • Model and Partnerships: The Gemma Model team is acknowledged alongside go-to-market and partnering teams.
  • Additional Support: Additional acknowledgments cover reinforcement learning, team support, contributors, product management, program management, executive sponsorship, and specialty areas.Named specialty areas include central support and checkpoint conversions.
Loading 2406.11409v2…