Katkat
Back

Implement SFT Sample Packing

MicrosoftMachine Learning Engineer · HR Screen
MediumMachine Learning EngineerHR ScreenCoding & Algorithms

Last updated

This problem tests your ability to efficiently pack variable-length sequences into fixed-length blocks for training, ensuring padding tokens do not influence the loss calculation.

Problem Statement

When training Large Language Models, it is computationally expensive to train on short sequences padded with zeros. "Sample Packing" concatenates multiple input sequences into a single fixed-length block of size max_seq_length.

Implement a function pack_sequences(sequences, max_seq_length) that takes a list of integer lists (sequences) and packs them into blocks. Each block must be exactly max_seq_length. If a sequence is too long, truncate it. If it fits, add it; if not, move to the next block. Use a special padding_token_id = 0 to fill remaining space in a block.

Additionally, return a loss_mask for each block. A loss_mask is a list of 1s and 0s where 1 indicates an actual token and 0 indicates a padding token.

Function Signature

def pack_sequences(sequences, max_seq_length):

Example

# Input
sequences = [[1, 2], [3, 4, 5], [6]]
max_seq_length = 4

# Output
# Block 1: [1, 2, 3, 4] -> mask: [1, 1, 1, 1]
# Block 2: [5, 6, 0, 0] -> mask: [1, 1, 0, 0]

Constraints

  • max_seq_length >= 1
  • If a single sequence is longer than max_seq_length, truncate it to the first max_seq_length tokens.
  • If after all sequences are processed, the last block is partially full, pad it to max_seq_length.
  • Ignore empty sequences.

Hints

Hint 1

Maintain a buffer representing the 'current block' and keep track of its current length as you iterate through the input sequences.

Hint 2

When a sequence exceeds the remaining space in the current block, truncate the sequence to fit the remaining space, close the block, and start a new one.

Hint 3

Use a boolean or integer array for the mask that is initialized to zeros and updated as you append tokens to the block.

Follow-up questions

  • How would you modify the packing logic if sequences must not be split across blocks (i.e., we cannot allow a single sequence to be 'interrupted' by a block boundary)?
  • How does sample packing affect the model's ability to learn cross-document dependencies?
  • If your sequences have an 'instruction' part and a 'response' part, how would you adjust the `loss_mask` to ensure the model only calculates loss on the response tokens?