My fork of karpathy/nanoGPT

Andrej Karpathy

Refactoring the code

1. I created a config class with some minimal functionality that mimics Karpathy's config.

2. I modularized the code, so a part of a module can be reused in other modules.

Perplexity

def compute_perplexity(model: GPT, config: Config, ctx: nullcontext | torch.amp.autocast) -> float:
    """Compute perplexity of the model"""
    X, y = get_validation(config)

    # Calculate number of batches
    n_samples = X.shape[0]
    n_batches = (n_samples + config.batch_size - 1) // config.batch_size

    total_log_probs = []

    with torch.no_grad():
        with ctx:
            for i in tqdm(range(n_batches)):
                # Get batch indices
                start_idx = i * config.batch_size
                end_idx = min((i + 1) * config.batch_size, n_samples)

                # Get batch data
                X_batch = X[start_idx:end_idx]
                y_batch = y[start_idx:end_idx]

                # Forward pass
                logits, _ = model(X_batch)
                logits = logits[:, -1, :]
                # convert to probabilities
                probs = torch.softmax(logits, dim=-1)
                # get the predicted logits for the y index
                pred_probs = probs[torch.arange(len(y_batch)), y_batch.int()]
                # apply log
                log_probs = torch.log(pred_probs)
                total_log_probs.append(log_probs)

    # Concatenate all log probabilities and compute mean
    all_log_probs = torch.cat(total_log_probs)
    perplexity = torch.exp(-all_log_probs.mean()).item()

    return perplexity

Experiments

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

First Citizen:
First, you know Caius Marcius is chief enemy to the people.

All:
We know't, we know't.

• train.bin has 1,003,854 tokens

• val.bin has 111,540 tokens

• block_size: 64 (context length)

• batch_size: 12 (for training)

• n_layer: 4 (number of layers)

• n_head: 4 (number of attention heads)

• n_embd: 128 (embedding dimension)

• max_iters: 2000 (total number of training iterations)

• lr_decay_iters: 2000 (learning rate decay iterations)

• dropout: 0.0 (dropout rate)

KING RICHARD II:
And these that he is a blassed man
She be my noble winders.

ISABELLA:
Methom; in the subjurit us,
No be as throat quied, and fear I not be,
I was brother in quiet and thee.

Scaling up

PARIS:
Ay, ay, because the flatterer, and leave your hearts
Where you that seems the wanton of a word.

POMPEY:
And be the more of what commons to him.

PAULINA:
If you must prove what now you may not here,
Have your mother is grace, it am beggar
A woman may be made him not only me.

PETER:
Yes, sir; I do not so? What may the vengeance?

LADY CAPULET:
Yes, my lord, in your grace to
---------------

Norfolk, let you have the friar of the bosom,
Takes the depender of my son whiles are shall.

JULIET:
I heard thee a pretting and true.

ROMEO:
O shame! O thou dost be confession,
See that for thy will speed to be beats;
Known though they do not sand I am soldier.

ROMEO:
O, nurse! I'll find thee to falls of a lights
As death dinner and the wars stopp'd, and love
To name when he hath head is not that slow it.

Inference speed

1. CPU (float32): Standard inference on the CPU without any optimizations.

2. CPU with int8 quantization: Dynamically quantized weights to int8 for faster inference. Here, we used the Pytorch quantization API and its method torch.ao.quantization.quantize_dynamic.

3. CPU with float16 inference: Converted weights to float16 using model.half() with the promise of faster inference speed and reduced model size.

4. MPS (Metal Performance Shaders, float32): GPU acceleration on macOS using the MPS backend.

5. MPS with float16 inference: Reduced precision inference on the MPS backend.

    param_size = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    buffer_size = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
    size_all_mb = (param_size + buffer_size) / 1024**2

• Model size behaves as expected. float16 reduces the model size significantly with respect to float32, and int8 brings the 6M parameter model size to under 1Mb.

• MPS performs significantly faster inference than the default CPU. I didn’t know about my Macbook having this. The MPS backend in PyTorch is specifically designed for macOS devices and leverages the Metal programming framework to enable high-performance computation on Apple's GPUs.

• Perplexity behaves as expected: reducing to float16 precision comes with a slight perplexity increase: 4.111 to 4.113. Finally, the int8 quantization provides the smallest model but at the highest perplexity cost: 4.116.

• The tokens-per-second versus precision trade-off does not behave as expected. Typically, we would expect that a lower precision model can use kernels that exploit reduced memory bandwidth and faster arithmetic operations. This should theoretically lead to improved throughput (tokens processed per second). My hypothesis on why this is not the case is the fact that we do not use NVIDIA GPUs but simply my Macbook CPU and the very recent MPS backend, which does not yet support the native Pytorch quantizations.