Sequence Parallel Activation Savings

Problem

Sequence parallelism splits LayerNorm and dropout activations along the token axis, so each rank holds ceil(seq_len / num_gpus) rows instead of all seq_len rows. Compute the per-rank bytes, the bytes saved and the savings fraction for one (seq_len, d_model) activation tensor.

The signature is sequence_parallel_savings(seq_len, d_model, num_gpus, bytes_per_value=2). Return [per_rank_bytes, saved_bytes, saving_fraction].

Hint

Ceiling division means the savings fraction can be slightly below (num_gpus - 1) / num_gpus.

Test cases

case 1

in: [1024,512,8,2]

expected: [131072,917504,0.875]

case 2

in: [10,4,3,2]

expected: [32,48,0.6]

case 3

in: [128,768,4,2]

expected: [49152,147456,0.75]

case 4

in: [1,100,8,4]

expected: [400,0,0]

Starter codePython

def sequence_parallel_savings(seq_len, d_model, num_gpus, bytes_per_value=2):
    # Your code here
    pass

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