Function split_packed_batch

src/sequence_packer.py:534–571

(batch: Any, microbatch_size: Union[int, float], padding_tolerance=1.0)

def split_packed_batch(batch: Any, microbatch_size: Union[int, float], padding_tolerance=1.0) -> Sequence:
    # NOTE: Packed sequences are already packed into a microbatch size worth of tokens.
    # So to correctly return a microbatch worth of data, we will simply return each item (i.e. microbatch_size 1)

    num_items = batch["input_ids"].shape[0]
    split_inputs = [x.squeeze() for x in batch["input_ids"].split(1)]
    split_labels = [x.squeeze() for x in batch["labels"].split(1)]
    split_attention_masks = [x.squeeze() for x in batch["attention_mask"].split(1)]
    split_cu_seqlens = batch["cu_seqlens"]

    result = []
    for i in range(num_items):
        attention_mask = split_attention_masks[i]
        padding_amount = 1 - (attention_mask.sum() / len(attention_mask))

        if padding_amount > padding_tolerance:
            last_non_pad = attention_mask.nonzero().max()
            input_ids = split_inputs[i][: last_non_pad + 1]
            labels = split_labels[i][: last_non_pad + 1]
            cu_seqlens = split_cu_seqlens[i][:-1]
            attention_mask = attention_mask[: last_non_pad + 1]
        else:
            input_ids = split_inputs[i]
            labels = split_labels[i]
            cu_seqlens = split_cu_seqlens[i]

        result.append(
            {
                "input_ids": input_ids,
                "labels": labels,
                "cu_seqlens": cu_seqlens,
                "max_seqlen": batch["max_seqlen"][i],
                "attention_mask": attention_mask,
            }
        )

    assert all([x["input_ids"].shape[-1] == y["cu_seqlens"][-1] for x, y in zip(result, result)])
    return result


def get_num_samples_in_packed_batch(batch: Batch) -> int: