bridge.training.utils.packed_seq_utils#

Module Contents#

Functions#

get_thd_cp_partition_indices

Return MCore’s context-parallel partition indices for a THD stream.

get_packed_seq_q_cu_seqlens

Return unpadded and physical query cumulative offsets.

get_packed_seq_cp_partition_indices

Return MCore’s partition indices for packed CP.

unpack_mcore_thd_tensor_for_position_ids

Reconstruct logical rows from a single-row MCore THD tensor.

repack_mcore_thd_position_ids

Scatter logical-row MRoPE positions back into a single THD row.

_squeeze_metadata

get_packed_seq_params

Build packed sequence parameters from a batch dictionary.

Data#

PackedMetadataValue

_MIN_MCORE_THD_CP_VERSION

API#

bridge.training.utils.packed_seq_utils.PackedMetadataValue#

None

bridge.training.utils.packed_seq_utils._MIN_MCORE_THD_CP_VERSION#

‘PkgVersion(…)’

bridge.training.utils.packed_seq_utils.get_thd_cp_partition_indices(
cu_seqlens: torch.Tensor,
*,
total_tokens: int,
cp_group: torch.distributed.ProcessGroup,
device: torch.device,
) → torch.Tensor#

Return MCore’s context-parallel partition indices for a THD stream.

Parameters:

  • cu_seqlens – Physical cumulative sequence offsets for the packed stream.

  • total_tokens – Total padded token count before CP partitioning.

  • cp_group – Context-parallel process group.

  • device – Device on which the returned indices will be consumed.

Returns:

Long tensor containing this CP rank’s indices into the full stream.

Raises:

RuntimeError – If the installed Megatron-Core version does not expose THD partitioning through get_batch_on_this_cp_rank.

bridge.training.utils.packed_seq_utils.get_packed_seq_q_cu_seqlens(
packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams,
) → tuple[torch.Tensor | None, torch.Tensor | None]#

Return unpadded and physical query cumulative offsets.

Parameters:

packed_seq_params – MCore THD sequence metadata.

Returns:

Unpadded query offsets and physical offsets. Physical offsets use the padded metadata when available and otherwise fall back to unpadded offsets.

bridge.training.utils.packed_seq_utils.get_packed_seq_cp_partition_indices(
packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams,
*,
total_tokens: int,
cp_size: int,
cp_rank: int,
device: torch.device,
cp_group: torch.distributed.ProcessGroup | None = None,
) → torch.Tensor#

Return MCore’s partition indices for packed CP.

Parameters:

  • packed_seq_params – MCore THD metadata for the full packed stream.

  • total_tokens – Total padded token count before CP partitioning.

  • cp_size – Context-parallel world size.

  • cp_rank – Context-parallel rank.

  • device – Device on which the returned indices will be consumed.

  • cp_group – Context-parallel process group. Uses MCore parallel state when omitted.

Returns:

Long tensor containing this CP rank’s indices into the full stream.

Raises:

ValueError – If packed query boundaries are unavailable or the requested rank and size do not match the context-parallel group.

bridge.training.utils.packed_seq_utils.unpack_mcore_thd_tensor_for_position_ids(
tensor: torch.Tensor,
packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams,
) → tuple[torch.Tensor, torch.Tensor, list[int], list[int]]#

Reconstruct logical rows from a single-row MCore THD tensor.

This is intended for model-specific position-ID builders that require a conventional batch dimension. Attention still consumes the original THD tensor and metadata.

Parameters:

  • tensor – Packed tensor with shape [1, total_padded_tokens].

  • packed_seq_params – Current MCore THD sequence metadata.

Returns:

Padded logical rows, their boolean attention mask, padded row starts, and unpadded row lengths.

Raises:

ValueError – If the tensor or packed metadata is inconsistent.

bridge.training.utils.packed_seq_utils.repack_mcore_thd_position_ids(
position_ids: torch.Tensor,
*,
padded_starts: list[int],
lengths: list[int],
total_length: int,
) → torch.Tensor#

Scatter logical-row MRoPE positions back into a single THD row.

Parameters:

  • position_ids – Position tensor with shape [axes, rows, max_length].

  • padded_starts – Start offset of each row in the padded THD tensor.

  • lengths – Unpadded length of each logical row.

  • total_length – Padded THD tensor length.

Returns:

Position tensor with shape [axes, 1, total_length]. Alignment gaps remain zero because they are excluded by packed metadata and loss masks.

Raises:

ValueError – If row metadata and position IDs are inconsistent.

bridge.training.utils.packed_seq_utils._squeeze_metadata(
value: bridge.training.utils.packed_seq_utils.PackedMetadataValue,
) → bridge.training.utils.packed_seq_utils.PackedMetadataValue#
bridge.training.utils.packed_seq_utils.get_packed_seq_params(
batch: dict[str, bridge.training.utils.packed_seq_utils.PackedMetadataValue],
) → megatron.core.packed_seq_params.PackedSeqParams#

Build packed sequence parameters from a batch dictionary.

Current MCore-style metadata is passed through directly after squeezing possible batch dimensions. Legacy Bridge metadata is still converted by removing any padding marked by -1 values.

Parameters:

batch – A dictionary containing packed-sequence metadata. Current keys are cu_seqlens_q, cu_seqlens_kv, optional padded variants, max_seqlen_q, max_seqlen_kv, and optional total_tokens (required for hybrid SSM/Mamba models to generate seq_idx). Legacy cu_seqlens / cu_seqlens_unpadded batches are also accepted for offline packed SFT compatibility.

Returns:

PackedSeqParams with identical q/kv parameters and qkv_format set to “thd”.