Better & Faster LLMs via Multi-token Prediction

Better & Faster Large Language Models via Multi-token Prediction

30 Apr 2024 https://arxiv.org/abs/2404.19737

This paper introduces a new LLM training approach called Multi-token Prediction.

Multi-token prediction shares the same transformer architecture as GPT-style LLMs, but the key difference is the addition of multiple output heads that train the model to predict several tokens simultaneously.

Standard LLMs such as GPT and Llama are trained with a next-token prediction loss, and the paper argues that this approach has several shortcomings:

1. The model learns only local patterns.

2. It tends to overlook "hard" decisions.

The paper concludes that a model trained purely on next-token prediction requires far more data to reach fluency comparable to that of a human child.

This work proposes training language models to predict multiple future tokens at once, with the goal of producing models that are more sample-efficient, higher-performing, and faster.

Method

Looking at the figure below in more detail: at each position in the training corpus, the model is trained to predict the next n tokens using n independent output heads operating on top of a shared model trunk.

This approach has the potential to increase memory usage significantly. To keep that in check, the paper introduces a technique illustrated in the figure and code below: gradients from each head are accumulated into the trunk and then released one at a time, so that adding more heads does not increase the memory footprint.

At inference time, a model trained with Multi-token Prediction uses only the output head for next-token prediction by default, discarding the remaining prediction heads — unlike during training. The implementation also supports using all output heads together via speculative decoding (similar to Medusa: https://arxiv.org/abs/2401.10774) to enable Multi-token Prediction at inference as well.

Results

1) Results by number of parameters

Multi-token Prediction yields larger gains as model size increases.

For smaller models, performance is similar to or slightly below that of the baseline.

2) Results by number of output heads

Performance peaks with 4–8 output heads; increasing beyond 8 heads actually degrades performance.

3) Per-task performance comparison

On CodeContests, the best results were obtained when the number of output heads was set to 4 for both training and inference. Summarization performance also peaked at 4 output heads.

However, for choice tasks and math benchmarks, increasing the number of output heads actually hurts performance.

4) Usage

Four model variants are available on Hugging Face, differentiated by training token count and the number of output tokens at inference. Inference code is fully open-sourced.

https://huggingface.co/facebook/multi-token-prediction

Required libraries for inference:

• torch
• fairscale
• fire
• sentencepiece

Run inference with the following commands:

pip install huggingface_hub
curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt install git-lfs
huggingface-cli login # then enter your Hugging Face API key

git lfs install
git clone https://huggingface.co/facebook/multi-token-prediction

pip install torch fairscale fire sentencepiece

torchrun --nproc_per_node 1 example_completion.py --ckpt_dir 7B_200B_4/ --tokenizer_path tokenizer.model --max_seq_len 128 --max_batch_size 2 # update paths as needed

To run inference with different prompts, edit the prompts section of example_completion.py:

from typing import Optional

import fire

from llama import Llama


def main(
    ckpt_dir: str,
    tokenizer_path: str,
    temperature: float = 0.2,
    top_p: float = 0.9,
    max_seq_len: int = 256,
    max_batch_size: int = 4,
    max_gen_len: Optional[int] = None,
):
    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
    )

    prompts = [
        # Edit this section with your own prompts
        """\
def fizzbuzz(n: int):""",
        """\
import argparse
def main(string: str):
    print(string)
    print(string[::-1])
if __name__ == "__main__":"""
    ]
    results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )
    for prompt, result in zip(prompts, results):
        print(prompt)
        print(f"> {result['generation']}")
        print("\n==================================\n")


if __name__ == "__main__":
    fire.Fire(main)