`bridge.data.megatron_mimo.sequence_pack`#

Sequence-packing assembly for MegatronMIMO language shards.

Concatenates the real (unpadded) tokens of a group of examples into a single packed sequence — the Transformer Engine thd layout (total-tokens × heads × head-dim) — with cu_seqlens boundaries. Downstream, the cu_seqlens make attention block-diagonal per example (the qwen3-vl attention / flash-attn path already honors this), so packing does not leak attention across examples.

This module packs the whole shard as one group — there is no per-pack length budget here, so the assembled pack is the entire shard.

Module Contents#

Functions#

`assemble_packed_sequence`	Concatenate one group’s real tokens into a single packed sequence (THD layout).
`pack_language_shard`	Pack a per-DP-shard language batch `[bs, S]` into a single packed sequence `[1, T]`.

API#

bridge.data.megatron_mimo.sequence_pack.assemble_packed_sequence( group: List[int], tokens: Optional[torch.Tensor], lengths: List[int], *, labels: Optional[torch.Tensor] = None, loss_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, ) → Dict[str, Any]#

Concatenate one group’s real tokens into a single packed sequence (THD layout).

tokens may be None on a non-first PP stage that has no input_ids but still must pack labels / loss_mask to the same [1, T] shape (F5, PP-consistent packing); in that case input_ids is None in the result and the device is taken from the first available tensor.

Parameters:

group – Row indices (into the source tensors) of the examples in this pack, in pack order.
tokens – Padded [B, S] token ids, or None (last-stage label-only pack).
lengths – Real (unpadded) length per row.
labels – Optional padded [B, S] labels (pad value -100).
loss_mask – Optional padded [B, S] loss mask (pad value 0).
position_ids – Optional padded position ids. Both 2-D [B, S] and Qwen-VL MRoPE 3-D [3, B, S] are handled; each sample’s real slice is concatenated, so the packed positions are per-sample-reset (each sample’s positions already start at 0).

Returns:

Dict with input_ids / labels / loss_mask as packed [1, T] tensors (T = sum of the group’s real lengths; input_ids is None when tokens is None), position_ids as [1, T] (2-D input) or [3, 1, T] (MRoPE input), cu_seqlens (int32, length len(group)+1), and max_seqlen (int).

Raises:

ValueError – If position_ids is neither 2-D [B, S] nor 3-D MRoPE [3, B, S]; if none of tokens / labels / loss_mask / position_ids is a tensor (no device or shape to pack to); or if the group’s segment lengths sum to zero (an empty packed shard cannot be built).

bridge.data.megatron_mimo.sequence_pack.pack_language_shard( data_batch: Dict[str, Any], *, lengths: torch.Tensor | None = None, ) → tuple[Dict[str, Any], Dict[str, Any]] | tuple[Dict[str, Any], None]#

Pack a per-DP-shard language batch [bs, S] into a single packed sequence [1, T].

Concatenates the shard’s real (unpadded) tokens — including image-placeholder tokens, so the downstream MIMO modality splice still fills them in order — into one packed sequence with cu_seqlens block-diagonal boundaries. This removes the per-sample padding compute (the LM no longer processes S-padded sequences) and makes attention cost ∝ real tokens.

The returned packing_kwargs feeds MimoModel.forward(packing_kwargs=...), which builds a THD PackedSeqParams (qkv_format='thd') and threads cu_seqlens to the GPT decoder for block-diagonal attention. Non-LM tensors (modality_inputs, etc.) are carried through unchanged.

The per-sample real length comes only from the caller-supplied lengths (forward_step derives it from the batch’s attention_mask on every language stage, so all stages pack to an identical [1, T]). None returns the batch unchanged.

loss_mask is not used for length (it is a supervision mask, not a padding mask).