Activation Memory Per Layer

Problem

Estimate activation memory stored per transformer block during training: 4 tensors of width d_model (attention input, attention output and two MLP projections) plus 2 tensors of width d_ff, each of shape (batch_size, seq_len) and bytes_per_value bytes wide.

The signature is activation_memory_per_layer(batch_size, seq_len, d_model, d_ff, bytes_per_value=2). Return an integer number of bytes.

Hint

Total elements per block is batch times sequence length times the summed widths.

Test cases

case 1

in: [8,1024,768,3072,2]

expected: 150994944

case 2

in: [1,128,64,256,4]

expected: 393216

case 3

in: [2,512,512,2048,2]

expected: 12582912

case 4

in: [4,256,128,512,2]

expected: 3145728

Starter codePython

def activation_memory_per_layer(batch_size, seq_len, d_model, d_ff, bytes_per_value=2):
    # Your code here
    pass

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