`bridge.data.megatron_mimo.reorder_buffer`#

Per-sample reorder buffer for MegatronMIMO intra-microbatch reordering.

Each data rank reads only its disjoint scalable-data-parallel shard of the global micro-batch. The per-microbatch cost imbalance between shards is corrected by a cost all-gather plus one ragged per-sample all_to_all_single over the module DP group, optionally overlapped with compute. Paired modules derive the same assignment from the same costs and canonical grid, so the BridgeCommunicator positional routing stays valid without cross-module communication.

Module Contents#

Classes#

`_WorkerException`	Carry an overlap-worker exception across the queue for re-raising on the consumer thread.
`ReorderingBuffer`	Wrap a MegatronMIMO data iterator so each yielded micro-batch is cost-balanced.

Functions#

`_alignment_padding_bytes`	Return the zero-padding bytes that round `nbytes` up to the serialization alignment.
`_dtype_info`	Return `(dtype, byte width)` for a dtype name from sample metadata.
`balanced_index_order`	Permute indices so contiguous equal-size groups have balanced total cost.
`balanced_assignment`	Map each sample of one micro-batch to its owner rank and canonical position.
`intra_route`	Lift a single-slot assignment to the `(dst_slot, owner_rank, canonical_pos)` window route.
`build_window_route`	Build the slot-major window route by balancing each slot independently.
`assert_intra_no_cross_slot`	Raise if any sample’s `dst_slot` differs from its source slot `g // b_global`.
`window_cost_spread`	Compute per-slot per-rank cost max/min before and after balancing and the remote-sample count.
`sample_keys`	Return the sorted tensor keys of a sample’s metadata, identical on sender and receiver.
`tensor_metadata`	Build per-field `{shape, dtype}` / `{non_tensor: value}` / `None` metadata for one sample.
`serialize_sample`	Serialize one sample’s tensors into a 1D aligned `uint8` buffer in `keys` order.
`sample_byte_size`	Return the serialized byte size of a sample from its metadata alone.
`deserialize_sample`	Reconstruct one sample from a `uint8` buffer (inverse of :func:`serialize_sample`).
`prepare_sample_exchange`	Build the kept-local partition, packed send buffer, and splits for one `all_to_all_single`.
`_flatten_nested`	Flatten a nested batch dict to dotted keys, preserving tensors and `None`.
`_unflatten_nested`	Rebuild the nested dict from dotted keys (inverse of :func:`_flatten_nested`).
`_cu_img_from_counts`	Return the cumulative image-offset table `[0, c0, c0+c1, ...]` from per-sample image counts.
`empty_like_vision`	Return empty `[0, d]` hidden states and `[0, 3]` grid matching a reference sample’s dtype and device.
`patches_per_image`	Return the per-image patch count `prod(t, h, w)` for a `[n_images, 3]` grid.
`_reorder_vision_by_images`	Permute `hidden_states` (patch dim) and `grid_thw` (image dim) by whole image.
`_gather_vision_subdict`	Gather a vision `{hidden_states, grid_thw}` sub-dict by sample, supporting 0..N images per sample.
`_apply_sample_dispatch`	Gather `sample_indices` from every tensor, vision sub-dict, and per-sample list in a batch.
`split_microbatch`	Split a micro-batch into `B` per-sample flat (dotted-key) dicts with size-1 batch dims.
`merge_samples`	Rebuild a micro-batch from per-sample flat dicts (inverse of :func:`split_microbatch`).
`_merge_nested`
`_batch_to_cuda`	Move a nested batch dict’s tensors to CUDA (non-blocking).
`_batch_size`	Return the sample count of a micro-batch from one of its LM tensors.
`sample_cost`	Return `encoder_cost_weight * patches + language_cost_weight * real_tokens` for one sample.
`reassemble_window`	Rebuild this rank’s `W` micro-batches from kept-local and received samples.
`exchange_by_route`	Run one ragged `all_to_all_single` for `route` and reassemble `W` micro-batches.
`exchange_window`	Cost-balance a window of `W` micro-batches across a module DP group in one exchange.
`reorder_window_local`	Permute a window into canonical order on a `dp_size == 1` module without any exchange.
`_record_stream_for_batch`	Tie every CUDA tensor in `value` to `stream` before a cross-thread handoff.
`build_module_dp_process_groups`	Create the Gloo and side NCCL process groups for this rank’s module DP group.

Data#

`logger`
`_SERIALIZED_TENSOR_BYTE_ALIGNMENT`
`_DTYPE_INFO`
`_SENTINEL`

API#

bridge.data.megatron_mimo.reorder_buffer.logger#: ‘getLogger(…)’

bridge.data.megatron_mimo.reorder_buffer._SERIALIZED_TENSOR_BYTE_ALIGNMENT#: 8

bridge.data.megatron_mimo.reorder_buffer._alignment_padding_bytes(nbytes: int) → int#: Return the zero-padding bytes that round nbytes up to the serialization alignment.

bridge.data.megatron_mimo.reorder_buffer._DTYPE_INFO#: None

bridge.data.megatron_mimo.reorder_buffer._dtype_info(dtype_name: str) → Tuple[torch.dtype, int]#: Return (dtype, byte width) for a dtype name from sample metadata.

bridge.data.megatron_mimo.reorder_buffer.balanced_index_order( costs: List[float], n_groups: int, ) → List[int]#

Permute indices so contiguous equal-size groups have balanced total cost.

Equal-cardinality LPT with a stable tie-break by group index, so paired modules that feed the same costs produce the same permutation.

Parameters:

costs – Per-item cost.
n_groups – Number of contiguous shards; must divide len(costs).

Returns:

A permutation of range(len(costs)) laid out group-by-group.

Raises:

ValueError – If n_groups <= 0 or does not divide len(costs).

bridge.data.megatron_mimo.reorder_buffer.balanced_assignment( costs: List[float], n_groups: int, dp_size: int, ) → List[Tuple[int, int]]#

Map each sample of one micro-batch to its owner rank and canonical position.

n_groups is the canonical grid (LCM of the module DP sizes); a module of size dp_size owns n_groups // dp_size consecutive groups, so modules with different DP sizes agree.

Parameters:

costs – Per-sample cost indexed by global position in the micro-batch.
n_groups – Canonical group count; must divide len(costs).
dp_size – This module’s DP size; must divide n_groups.

Returns:

(owner_rank, canonical_pos) per global sample index.

Raises:

ValueError – If n_groups does not divide len(costs) or dp_size does not divide it.

bridge.data.megatron_mimo.reorder_buffer.intra_route( assignment: List[Tuple[int, int]], *, src_slot: int = 0, ) → List[Tuple[int, int, int]]#

Lift a single-slot assignment to the (dst_slot, owner_rank, canonical_pos) window route.

Intra routing never moves a sample across slots, so dst_slot == src_slot for every sample.

Parameters:

assignment – Output of :func:balanced_assignment for one slot.
src_slot – Slot the samples were read from.

Returns:

(dst_slot, owner_rank, canonical_pos) per global sample index.

bridge.data.megatron_mimo.reorder_buffer.build_window_route( costs_per_slot: List[List[float]], n_groups: int, dp_size: int, ) → List[Tuple[int, int, int]]#

Build the slot-major window route by balancing each slot independently.

The window global index of slot s, slot-local position j is g = s * B_global + j.

Parameters:

costs_per_slot – Per-slot per-sample costs indexed by slot-local position.
n_groups – Canonical group count; must divide each slot’s sample count.
dp_size – This module’s DP size.

Returns:

(dst_slot, owner_rank, canonical_pos) per window global index.

bridge.data.megatron_mimo.reorder_buffer.assert_intra_no_cross_slot( route: List[Tuple[int, int, int]], b_global: int, ) → None#

Raise if any sample’s dst_slot differs from its source slot g // b_global.

Parameters:

route – (dst_slot, owner_rank, canonical_pos) per window global index.
b_global – Samples per micro-batch slot; 0 skips the check.

Raises:

RuntimeError – If a sample’s dst_slot differs from its source slot.

bridge.data.megatron_mimo.reorder_buffer.window_cost_spread( costs_per_slot: List[List[float]], route: List[Tuple[int, int, int]], dp_size: int, ) → List[Dict[str, float]]#

Compute per-slot per-rank cost max/min before and after balancing and the remote-sample count.

Parameters:

costs_per_slot – Per-slot per-sample costs indexed by slot-local position.
route – Window route from :func:build_window_route.
dp_size – This module’s DP size.

Returns:

One {before_max, before_min, after_max, after_min, remote} dict per slot.

bridge.data.megatron_mimo.reorder_buffer.sample_keys(meta: Dict[str, Any]) → List[str]#

Return the sorted tensor keys of a sample’s metadata, identical on sender and receiver.

Parameters:: meta – Sample metadata from :func:tensor_metadata.
Returns:: Sorted keys whose metadata entry carries a dtype.

bridge.data.megatron_mimo.reorder_buffer.tensor_metadata( flat: Dict[str, Any], ) → Dict[str, Any]#

Build per-field {shape, dtype} / {non_tensor: value} / None metadata for one sample.

Parameters:: flat – One flattened sample dict.
Returns:: Metadata keyed like flat, sufficient to size and decode the sample without its data.
Raises:: ValueError – If a tensor field has a dtype the serializer does not support.

bridge.data.megatron_mimo.reorder_buffer.serialize_sample( flat: Dict[str, Any], keys: List[str], ) → torch.Tensor#

Serialize one sample’s tensors into a 1D aligned uint8 buffer in keys order.

Parameters:

flat – Flattened sample dict; None values are skipped.
keys – Sorted tensor keys from :func:sample_keys.

Returns:

1D uint8 tensor with each field padded to the serialization alignment.

bridge.data.megatron_mimo.reorder_buffer.sample_byte_size( meta: Dict[str, Any], keys: List[str], ) → int#

Return the serialized byte size of a sample from its metadata alone.

Parameters:

meta – Sample metadata from :func:tensor_metadata.
keys – Sorted tensor keys from :func:sample_keys.

Returns:

Byte count matching :func:serialize_sample for the same sample.

bridge.data.megatron_mimo.reorder_buffer.deserialize_sample( buf: torch.Tensor, offset: int, meta: Dict[str, Any], keys: List[str], ) → Tuple[Dict[str, Any], int]#

Reconstruct one sample from a uint8 buffer (inverse of :func:serialize_sample).

Parameters:

buf – Received 1D uint8 buffer.
offset – Byte offset of this sample in buf.
meta – This sample’s metadata from :func:tensor_metadata.
keys – Sorted tensor keys from :func:sample_keys.

Returns:

the rebuilt sample and the byte offset just past it in buf.

Return type:

(flat, new_offset)

bridge.data.megatron_mimo.reorder_buffer.prepare_sample_exchange( local_flats: List[Dict[str, Any]], local_global_indices: List[int], route: List[Tuple[int, int, int]], all_global_indices: List[List[int]], all_tensor_meta: List[List[Dict[str, Any]]], dp_rank: int, dp_size: int, *, window_size: int = 1, ) → Dict[str, Any]#

Build the kept-local partition, packed send buffer, and splits for one all_to_all_single.

Sender and receiver sort moves by (dst_slot, canonical_pos, src_local_idx), so the packed bytes and recv_schedule line up without further communication.

Parameters:

local_flats – This rank’s flattened samples in local order.
local_global_indices – Global index of each entry in local_flats.
route – (dst_slot, owner_rank, canonical_pos) per global sample index.
all_global_indices – Global indices held by each rank, all-gathered.
all_tensor_meta –

func:

tensor_metadata of each rank’s samples, all-gathered.
dp_rank – This rank’s module-local DP rank.
dp_size – This module’s DP size.
window_size – Number of micro-batch slots; every dst_slot must be in [0, window_size).

Returns:

Dict with local_samples ((dst_slot, canonical_pos, local_idx) kept here), send_buf, per-rank byte send_splits / recv_splits, and recv_schedule ({src_rank: [(dst_slot, canonical_pos, src_local_idx)]} in receive order).

Raises:

ValueError – If a dst_slot is outside [0, window_size).

bridge.data.megatron_mimo.reorder_buffer._flatten_nested( d: Dict[str, Any], prefix: str = '', ) → Dict[str, Any]#: Flatten a nested batch dict to dotted keys, preserving tensors and None.

bridge.data.megatron_mimo.reorder_buffer._unflatten_nested( flat: Dict[str, Any], ) → Dict[str, Any]#: Rebuild the nested dict from dotted keys (inverse of :func:_flatten_nested).

bridge.data.megatron_mimo.reorder_buffer._cu_img_from_counts(image_counts: Any) → List[int]#: Return the cumulative image-offset table [0, c0, c0+c1, ...] from per-sample image counts.

bridge.data.megatron_mimo.reorder_buffer.empty_like_vision( ref_hidden_states: torch.Tensor, ref_grid_thw: torch.Tensor, ) → Tuple[torch.Tensor, torch.Tensor]#

Return empty [0, d] hidden states and [0, 3] grid matching a reference sample’s dtype and device.

Text-only samples contribute these placeholders so every sample carries the same vision keys and the merge concat and byte-split metadata stay symmetric across ranks.

Parameters:

ref_hidden_states – A [patches, d] hidden-states tensor to match.
ref_grid_thw – A [n_images, 3] grid tensor to match.

Returns:

(hidden_states, grid_thw) with zero rows.

bridge.data.megatron_mimo.reorder_buffer.patches_per_image(grid_thw: torch.Tensor) → torch.Tensor#

Return the per-image patch count prod(t, h, w) for a [n_images, 3] grid.

Parameters:: grid_thw – Per-image (t, h, w) grid.
Returns:: [n_images] patch counts.

bridge.data.megatron_mimo.reorder_buffer._reorder_vision_by_images( hidden_states: torch.Tensor, grid_thw: torch.Tensor, image_perm: list[int], ) → Tuple[torch.Tensor, torch.Tensor]#: Permute hidden_states (patch dim) and grid_thw (image dim) by whole image.

bridge.data.megatron_mimo.reorder_buffer._gather_vision_subdict( value: Dict[str, Any], sample_indices: list[int], n_samples: int, *, cu_img: list[int] | None = None, ) → Dict[str, Any]#: Gather a vision {hidden_states, grid_thw} sub-dict by sample, supporting 0..N images per sample.

bridge.data.megatron_mimo.reorder_buffer._apply_sample_dispatch( batch: Dict[str, Any], sample_indices: list[int], *, n_samples: int, cu_img: list[int] | None = None, ) → Dict[str, Any]#: Gather sample_indices from every tensor, vision sub-dict, and per-sample list in a batch.

bridge.data.megatron_mimo.reorder_buffer.split_microbatch( batch: Dict[str, Any], *, cu_img: List[int] | None = None, ) → List[Dict[str, Any]]#

Split a micro-batch into B per-sample flat (dotted-key) dicts with size-1 batch dims.

Parameters:

batch – Nested micro-batch.
cu_img – Cumulative per-sample image offsets (length B + 1); None assumes one image per sample.

Returns:

B flat sample dicts in batch order.

bridge.data.megatron_mimo.reorder_buffer.merge_samples( flats: List[Dict[str, Any]], ) → Dict[str, Any]#

Rebuild a micro-batch from per-sample flat dicts (inverse of :func:split_microbatch).

Parameters:: flats – Per-sample flat dicts in the desired batch order.
Returns:: The nested micro-batch, tensors concatenated on their batch, patch, or image dim.
Raises:: ValueError – If flats is empty.

bridge.data.megatron_mimo.reorder_buffer._merge_nested( samples: List[Dict[str, Any]], ) → Dict[str, Any]#

bridge.data.megatron_mimo.reorder_buffer._batch_to_cuda( batch: Dict[str, Any], ) → Dict[str, Any]#: Move a nested batch dict’s tensors to CUDA (non-blocking).

bridge.data.megatron_mimo.reorder_buffer._batch_size(batch: Dict[str, Any]) → int#: Return the sample count of a micro-batch from one of its LM tensors.

bridge.data.megatron_mimo.reorder_buffer.sample_cost( flat: Dict[str, Any], *, encoder_cost_weight: float, language_cost_weight: float = 0.0, image_token_id: int | None = None, square_merge_size: int = 1, ) → float#

Return encoder_cost_weight * patches + language_cost_weight * real_tokens for one sample.

Both terms must be collation-independent so vision and language derive the same cost without communicating: with image_token_id the patch count comes from input_ids, which every module holds (grid_thw is nulled on language shards), and tokens come from attention_mask.

Parameters:

flat – One per-sample flat (dotted-key) dict.
encoder_cost_weight – Weight per vision patch.
language_cost_weight – Weight per real token; 0.0 disables the term.
image_token_id – Placeholder token id counted in input_ids; None sums grid_thw instead.
square_merge_size – Patches represented by one placeholder token (spatial_merge_size**2).

Returns:

The weighted per-sample cost.

bridge.data.megatron_mimo.reorder_buffer.reassemble_window( local_flats: List[Dict[str, Any]], plan: Dict[str, Any], recv_buf: torch.Tensor, all_tensor_meta: List[List[Dict[str, Any]]], *, dp_size: int, window_size: int, local: int, ) → List[Dict[str, Any]]#

Rebuild this rank’s W micro-batches from kept-local and received samples.

Parameters:

local_flats – This rank’s input samples; kept-local ones are indexed by local_idx.
plan – Output of :func:prepare_sample_exchange.
recv_buf – The all_to_all_single output buffer.
all_tensor_meta – All-gathered per-rank tensor metadata used to decode recv_buf.
dp_size – This module’s DP size.
window_size – Number of micro-batch slots W.
local – Expected samples per slot on this rank.

Returns:

W micro-batches, slot s at index s, each in canonical order.

Raises:

RuntimeError – If the received byte count, a slot’s sample count, or the canonical positions do not match plan.

bridge.data.megatron_mimo.reorder_buffer.exchange_by_route( local_flats: List[Dict[str, Any]], route: List[Tuple[int, int, int]], all_global_indices: List[List[int]], all_tensor_meta: List[List[Dict[str, Any]]], *, dp_rank: int, dp_size: int, window_size: int, local: int, dp_group_nccl: Any, nccl_stream: Any = None, ) → List[Dict[str, Any]]#

Run one ragged all_to_all_single for route and reassemble W micro-batches.

Parameters:

local_flats – This rank’s input samples.
route – (dst_slot, owner_rank, canonical_pos) per global sample index.
all_global_indices – Global indices held by each rank, all-gathered.
all_tensor_meta – All-gathered per-rank tensor metadata.
dp_rank – This rank’s module-local DP rank.
dp_size – This module’s DP size.
window_size – Number of micro-batch slots W.
local – Samples per slot on this rank.
dp_group_nccl – NCCL process group for the sample all-to-all.
nccl_stream – Optional CUDA stream to run the all-to-all on.

Returns:

W micro-batches for this rank, slot s at index s.

bridge.data.megatron_mimo.reorder_buffer.exchange_window( batches: List[Dict[str, Any]], *, dp_rank: int, dp_size: int, n_groups: int, cost_of: Any, dp_group_gloo: Any, dp_group_nccl: Any, nccl_stream: Any = None, image_count_of: Any = None, probe: bool = False, ) → List[Dict[str, Any]]#

Cost-balance a window of W micro-batches across a module DP group in one exchange.

Costs and metadata are all-gathered over Gloo once per window and the samples move in one all_to_all_single; paired modules must pass the same n_groups and cost_of.

Parameters:

batches – This rank’s W micro-batch shards, each with the same local sample count.
dp_rank – This rank’s module-local DP rank.
dp_size – This module’s DP size.
n_groups – Canonical group count (LCM of the module DP sizes).
cost_of – callable(flat) -> float per-sample cost (see :func:sample_cost).
dp_group_gloo – Gloo process group for the metadata all-gather.
dp_group_nccl – NCCL process group for the sample all-to-all.
nccl_stream – Optional CUDA stream to run the all-to-all on.
image_count_of – Optional callable(batch) -> LongTensor per-sample image counts.
probe – When True, log the per-slot cost spread on dp_rank == 0.

Returns:

W rebalanced micro-batches for this rank, slot s at index s.

Raises:

ValueError – If batches is empty or the shards differ in local size.

bridge.data.megatron_mimo.reorder_buffer.reorder_window_local( batches: List[Dict[str, Any]], *, n_groups: int, cost_of: Any, image_count_of: Any = None, ) → List[Dict[str, Any]]#

Permute a window into canonical order on a dp_size == 1 module without any exchange.

Parameters:

batches – This rank’s W whole micro-batches.
n_groups – Canonical group count (LCM of the module DP sizes), same as the paired module.
cost_of – callable(flat) -> float per-sample cost, same as the paired module.
image_count_of – Optional callable(batch) -> LongTensor per-sample image counts.

Returns:

W micro-batches permuted into canonical order.

bridge.data.megatron_mimo.reorder_buffer._SENTINEL#: ‘object(…)’

class bridge.data.megatron_mimo.reorder_buffer._WorkerException(exc: BaseException)#

Carry an overlap-worker exception across the queue for re-raising on the consumer thread.

Initialization

bridge.data.megatron_mimo.reorder_buffer._record_stream_for_batch( value: Any, stream: torch.cuda.Stream, ) → None#: Tie every CUDA tensor in value to stream before a cross-thread handoff.

class bridge.data.megatron_mimo.reorder_buffer.ReorderingBuffer( data_iterator: Any, *, dp_rank: int, dp_size: int, n_groups: int, cost_of: Any, dp_group_gloo: Any, dp_group_nccl: Any, overlap: bool = True, image_count_of: Any = None, window_size: int = 1, )#

Wrap a MegatronMIMO data iterator so each yielded micro-batch is cost-balanced.

window_size micro-batches are balanced in one :func:exchange_window call and served one at a time; at dp_size == 1 they are permuted locally. With overlap=True a background thread exchanges the next window on a dedicated CUDA stream and the side dp_group_nccl while the main thread computes the current one, so up to three windows are resident at once. CUDA_DEVICE_MAX_CONNECTIONS=1 is rejected with overlap=True because the side all-to-all and the DDP all-reduce then share one hardware queue and can deadlock.

Parameters:

data_iterator – Source iterator yielding this rank’s micro-batch shards.
dp_rank – This rank’s module-local DP rank.
dp_size – This module’s DP size.
n_groups – Canonical group count (LCM of the module DP sizes).
cost_of – callable(flat) -> float per-sample cost.
dp_group_gloo – Gloo process group for the metadata all-gather.
dp_group_nccl – NCCL process group for the sample all-to-all.
overlap – Exchange the next window on a background thread.
image_count_of – Optional callable(batch) -> LongTensor per-sample image counts.
window_size – Micro-batches balanced per exchange.

Raises:

ValueError – If window_size < 1 or overlap=True under CUDA_DEVICE_MAX_CONNECTIONS=1.

Initialization

__iter__()#

shutdown() → None#: Stop the overlap worker and wait briefly for it to exit.

__del__() → None#: Shut down the worker for runs that drop the iterator without closing it.

_exchange_window( batches: List[Dict[str, Any]], stream: Any, ) → List[Dict[str, Any]]#

_collect_window() → List[Dict[str, Any]]#: Pull up to window_size micro-batches from the source iterator.

_refill_window() → None#: Collect and exchange the next window synchronously; leave _active empty when exhausted.

_put_worker_item(item: Any) → None#: Put an item on the bounded queue unless shutdown is requested while it is full.

_worker() → None#

__next__() → Dict[str, Any]#

bridge.data.megatron_mimo.reorder_buffer.build_module_dp_process_groups( dp_group_nccl_main: Any, *, overlap: bool, ) → Tuple[int, int, Any, Any]#

Create the Gloo and side NCCL process groups for this rank’s module DP group.

dist.new_group is collective over the world, so every rank creates every module’s groups in the same sorted order and keeps the ones it belongs to.

Parameters:

dp_group_nccl_main – This rank’s main DP NCCL group (grid.get_pg(["dp"])).
overlap – Create a separate NCCL group for the side-stream exchange; otherwise reuse the main one.

Returns:

(dp_rank, dp_size, dp_group_gloo, dp_group_nccl).

bridge.data.megatron_mimo.reorder_buffer#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.data.megatron_mimo.reorder_buffer`#