Vocab Projection Memory

Problem

Estimate the memory of the vocabulary projection (output head) matrix: vocab_size * d_model * bytes_per_value.

The signature is vocab_projection_memory(vocab_size, d_model, bytes_per_value=4). Return an integer number of bytes.

Hint

The output head is one of the largest single matrices in an LLM.

case 1

in: [32000,1024,2]

expected: 65536000

case 2

in: [50000,768,4]

expected: 153600000

case 3

in: [1,1,1]

expected: 1

case 4

in: [1000,128,4]

expected: 512000

Starter codePython

def vocab_projection_memory(vocab_size, d_model, bytes_per_value=4):
    # Your code here
    pass

4 test cases · run your code against them in the editor