FSDP Shard Size

Problem

Compute the per-GPU view of fully sharded data parallel parameters: params_per_gpu = ceil(num_params / num_gpus), and bytes_per_gpu = params_per_gpu * bytes_per_param.

The signature is fsdp_shard_size(num_params, num_gpus, bytes_per_param=4). Return [params_per_gpu, bytes_per_gpu].

Hint

Each rank stores only its shard plus transient gathered parameters.

Test cases

case 1

in: [1000000,8,4]

expected: [125000,500000]

case 2

in: [1001,8,4]

expected: [126,504]

case 3

in: [1,1,2]

expected: [1,2]

case 4

in: [100,3,4]

expected: [34,136]

Starter codePython

def fsdp_shard_size(num_params, num_gpus, bytes_per_param=4):
    # Your code here
    pass

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