Training FLOPs Per Token

Problem

Estimate training FLOPs per token with the 6N rule: the forward pass costs 2N, the backward pass costs roughly twice the forward pass, and N is the parameter count.

The signature is training_flops_per_token(num_params). Return an integer number of FLOPs.

Hint

Forward plus backward is about three forward passes, and each forward pass is 2N.

Test cases

case 1

in: [7000000000]

expected: 42000000000

case 2

in: [1]

expected: 6

case 3

in: [125000000]

expected: 750000000

case 4

in: [0]

expected: 0

Starter codePython

def training_flops_per_token(num_params):
    # Your code here
    pass

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