bridge.data.packing.in_batch#

Collate-time construction of Megatron-Core THD packed sequence batches.

Module Contents#

Functions#

_ceil_to_multiple

_token_key

_set_tokens

_right_padded_sequence_lengths

_validate_sequence_tensor

_validate_sequence_row_tensor

build_mcore_thd_sequence_batch_from_rows

Build an MCore THD batch directly from unpadded sequence rows.

pack_right_padded_sequence_batch_to_mcore_thd

Pack a right-padded sequence batch into MCore THD layout.

API#

bridge.data.packing.in_batch._ceil_to_multiple(value: int, multiple: int) → int#
bridge.data.packing.in_batch._token_key(
batch: collections.abc.MutableMapping[str, Any],
) → str#
bridge.data.packing.in_batch._set_tokens(
batch: collections.abc.MutableMapping[str, Any],
token_key: str,
value: torch.Tensor,
) → None#
bridge.data.packing.in_batch._right_padded_sequence_lengths(
tokens: torch.Tensor,
*,
pad_token_id: int,
padding_mask: torch.Tensor | None,
) → list[int]#
bridge.data.packing.in_batch._validate_sequence_tensor(
batch: collections.abc.MutableMapping[str, Any],
key: str,
*,
tokens: torch.Tensor,
) → torch.Tensor | None#
bridge.data.packing.in_batch._validate_sequence_row_tensor(
row: collections.abc.Mapping[str, Any],
key: str,
*,
length: int,
) → torch.Tensor | None#
bridge.data.packing.in_batch.build_mcore_thd_sequence_batch_from_rows(
rows: collections.abc.Sequence[collections.abc.Mapping[str, Any]],
*,
token_key: str = 'input_ids',
sequence_length: int | None = None,
pad_token_id: int = 0,
ignore_index: int = IGNORE_INDEX,
pad_to_multiple_of: int = 1,
sequence_tensor_pad_values: collections.abc.Mapping[str, int | float] | None = None,
) → dict[str, Any]#

Build an MCore THD batch directly from unpadded sequence rows.

Parameters:

  • rows – Per-example mappings containing 1D sequence tensors.

  • token_key – Token tensor key present in each row.

  • sequence_length – Optional maximum length for each unpadded row.

  • pad_token_id – Token value for per-sequence alignment padding.

  • ignore_index – Label value for per-sequence alignment padding.

  • pad_to_multiple_of – Per-sequence alignment multiple for CP/SP.

  • sequence_tensor_pad_values – Additional sequence-aligned tensor keys and the value used for alignment padding.

Returns:

A single-row THD batch with current MCore packed-sequence metadata.

Raises:

ValueError – If rows are empty, inconsistent, padded, or overlength.

bridge.data.packing.in_batch.pack_right_padded_sequence_batch_to_mcore_thd(
batch: collections.abc.MutableMapping[str, Any],
*,
token_key: str | None = None,
sequence_length: int | None = None,
pad_token_id: int = 0,
ignore_index: int = IGNORE_INDEX,
pad_to_multiple_of: int = 1,
sequence_tensor_pad_values: collections.abc.Mapping[str, int | float] | None = None,
) → None#

Pack a right-padded sequence batch into MCore THD layout.

This helper emits the current PackedSeqParams field names consumed by Megatron-Core and Transformer Engine: cu_seqlens_q, cu_seqlens_kv, optional padded cu-seqlens, and max_seqlen_q / max_seqlen_kv.

Parameters:

  • batch – Mutable collate batch with 2D token, position, and optional label/loss-mask tensors.

  • token_key – Token key to use. If unset, detected from tokens or input_ids.

  • sequence_length – Optional maximum length for each unpadded row.

  • pad_token_id – Token value for padding inserted by pad_to_multiple_of.

  • ignore_index – Label value for inserted padding.

  • pad_to_multiple_of – Optional per-sequence packed length multiple.

  • sequence_tensor_pad_values – Additional sequence-aligned tensor keys and their alignment padding values.

Raises:

ValueError – If required tensors are missing or the batch contains no non-padding tokens.