Parameter Shard Range

Problem

Compute the [start, end) index range of shard shard_id when a flat parameter tensor of total_params is split into num_shards contiguous chunks of ceil(total_params / num_shards) elements, stopping at total_params.

The signature is parameter_shard_range(total_params, num_shards, shard_id). Return [start, end).

Hint

Each data-parallel rank owns one contiguous shard of every flattened tensor.

Test cases

case 1

in: [100,4,0]

expected: [0,25]

case 2

in: [100,4,3]

expected: [75,100]

case 3

in: [10,3,1]

expected: [4,8]

case 4

in: [10,3,2]

expected: [8,10]

case 5

in: [7,7,5]

expected: [5,6]

Starter codePython

def parameter_shard_range(total_params, num_shards, shard_id):
    # Your code here
    pass

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