Weight-Only Quant Memory

Problem

Estimate weight-only quantization memory in bytes: num_params * bits / 8, plus 4 bytes per group for the scale when group_size > 0 (num_params // group_size groups).

The signature is weight_only_quant_memory(num_params, bits=4, group_size=0). Return a float number of bytes.

Hint

Grouped scales add a small fixed overhead on top of the packed weights.

Test cases

case 1

in: [1000000,4,0]

expected: 500000

case 2

in: [1024,8,0]

expected: 1024

case 3

in: [1024,4,128]

expected: 544

case 4

in: [8000000,4,64]

expected: 4500000

Starter codePython

def weight_only_quant_memory(num_params, bits=4, group_size=0):
    # Your code here
    pass

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