Checkpoint Shard Merge

Problem

Merge parameter shards saved by different ranks into one flat parameter vector and truncate to the true full_length, since the last shard is often padded.

The signature is merge_checkpoint_shards(shards, full_length). Return the flat merged vector.

Hint

Shard order is the rank order used when the checkpoint was written.

Test cases

case 1

in: [[[1,2],[3],[4,5,6]],5]

expected: [1,2,3,4,5]

case 2

in: [[[1],[2],[3]],6]

expected: [1,2,3]

case 3

in: [[[1,2],[3]],3]

expected: [1,2,3]

case 4

in: [[[],[1,2]],2]

expected: [1,2]

Starter codePython

def merge_checkpoint_shards(shards, full_length):
    # Your code here
    pass

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