Optimal Bucket Batching on GPUs | Dataford Interview Questions

Implement bucket batching: given K documents and G GPUs, how would you find the optimal batching that minimizes padding, with 0 <= K < G?

Represent each document by its sequence length. Implement optimal_bucket_batching(lengths, g), returning exactly g buckets of document indices. Padding for a bucket is max_length * document_count - sum(lengths). Minimize total padding. Empty buckets are allowed, and document order should be preserved.

Examples: lengths=[8, 3], g=4 returns [[0], [1], [], []] with zero padding. lengths=[], g=3 returns [[], [], []].

Constraints

0 <= len(lengths) < g
1 <= g <= 10^5
1 <= lengths[i] <= 10^9
Return exactly g buckets
Each document index must appear exactly once

Function Signature

def optimal_bucket_batching(lengths, g):

Your question is Optimal Bucket Batching on GPUs. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

Problem

Implement bucket batching: given K documents and G GPUs, how would you find the optimal batching that minimizes padding, with 0 <= K < G?

Examples: lengths=[8, 3], g=4 returns [[0], [1], [], []] with zero padding. lengths=[], g=3 returns [[], [], []].

Constraints

0 <= len(lengths) < g
1 <= g <= 10^5
1 <= lengths[i] <= 10^9
Return exactly g buckets
Each document index must appear exactly once

Function Signature

def optimal_bucket_batching(lengths, g):

Your question is Optimal Bucket Batching on GPUs. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.