https://github.com/jaymody/speculative-sampling

Simple implementation of Speculative Sampling in NumPy for GPT-2.
https://github.com/jaymody/speculative-sampling

Last synced: 2 months ago
JSON representation

Simple implementation of Speculative Sampling in NumPy for GPT-2.

• Host: GitHub
• URL: https://github.com/jaymody/speculative-sampling
• Owner: jaymody
• Created: 2023-02-09T00:51:07.000Z (almost 2 years ago)
• Default Branch: main
• Last Pushed: 2023-08-20T01:23:43.000Z (over 1 year ago)
• Last Synced: 2024-08-02T19:38:49.265Z (6 months ago)
• Language: Python
• Homepage:
• Size: 25.4 KB
• Stars: 84
• Watchers: 3
• Forks: 9
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

• awesome-adaptive-computation - PyTorch code

README

        
# Speculative Sampling

A simple implementation of [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318.pdf) in NumPy for GPT-2. See [`main.py`](https://github.com/jaymody/speculative-sampling/blob/main/main.py). I also wrote a [blog post](https://jaykmody.com/blog/speculative-sampling/) for this implementation.

**Install Dependencies**:

```bash

pip install -r picoGPT/requirements.txt

```

Tested on `Python 3.9.10`.

**Usage**:

```python

python main.py \

    --prompt "Alan Turing theorized that computers would one day become" \

    --n_tokens_to_generate 40 \

    --draft_model_size "124M" \

    --target_model_size "1558M" \

    --K 4 \

    --temperature 0 # 0 for greedy sampling

```

Which outputs:

```text

Autoregressive Decode

---------------------

Time = 60.64s

Text = Alan Turing theorized that computers would one day become so powerful that they would be able to think like humans.

In the 1950s, he proposed a way to build a computer that could think like a human. He called it the "T

Speculative Decode

------------------

Time = 27.15s

Text = Alan Turing theorized that computers would one day become so powerful that they would be able to think like humans.

In the 1950s, he proposed a way to build a computer that could think like a human. He called it the "T

```