Gradient Bucket Size

Problem

Compute training bucket metadata: parameters per bucket = floor(cap_mb * 1024 * 1024 / bytes_per_param), and buckets needed = ceil(num_params / parameters_per_bucket).

The signature is gradient_bucket_size(num_params, bucket_cap_mb, bytes_per_param=4). Return [params_per_bucket, num_buckets].

Hint

Bucketing overlaps all-reduce with backward computation.

Test cases

case 1

in: [100000000,100,4]

expected: [26214400,4]

case 2

in: [1000,1,4]

expected: [262144,1]

case 3

in: [1000000,4,4]

expected: [1048576,1]

case 4

in: [10000000,1,2]

expected: [524288,20]

Starter codePython

def gradient_bucket_size(num_params, bucket_cap_mb, bytes_per_param=4):
    # Your code here
    pass

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