ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMimo V2 Flash

nemo_automodel.components.models.mimo_v2_flash.cp

View as Markdown

MiMo adapters for the framework Transformer Engine THD sharder.

Module Contents

Functions

NameDescription
_chunk_partition_indicesReproduce TE’s data-dependent THD token partition for every chunk.
_flatten_chunksFlatten batch rows into one token stream per pipeline chunk.
_media_maskBuild a global placeholder mask before Transformer Engine shards tokens.
make_mimo_te_cp_sharderCreate MiMo’s thin adapter around the framework TE THD sharder.
shard_batch_for_mimo_teDelegate MiMo packed CP to the framework TE THD sharder.

Data

_MIMO_GLOBAL_IMAGE_MASK

_MIMO_GLOBAL_VIDEO_MASK

_MIMO_THD_LOCAL_INDICES

_SEQ_LENS_PADDING_VALUE

_VLM_PP_MEDIA_KEY

__all__

API

nemo_automodel.components.models.mimo_v2_flash.cp._chunk_partition_indices(
    batch: dict[str, typing.Any],
    cp_mesh,
    num_chunks: int,
    seq_lens_padding_value: int
) -> torch.Tensor

Reproduce TE’s data-dependent THD token partition for every chunk.

Parameters:

batch
dict[str, Any]

Unsharded packed batch. input_ids has shape [batch, sequence] and seq_lens_padded has shape [batch, documents].

cp_mesh

Optional one-dimensional context-parallel mesh.

num_chunks
int

Number of equal pipeline chunks along the batch axis.

seq_lens_padding_value
int

Sentinel used in ragged length rows.

Returns: torch.Tensor

Global token indices with shape [local_tokens] for one chunk or

nemo_automodel.components.models.mimo_v2_flash.cp._flatten_chunks(
    tensor: torch.Tensor,
    num_chunks: int
) -> torch.Tensor

Flatten batch rows into one token stream per pipeline chunk.

Parameters:

tensor
torch.Tensor

Token-aligned tensor of shape [batch, sequence].

num_chunks
int

Number of equal groups along the batch axis.

Returns: torch.Tensor

Tensor of shape [num_chunks, batch * sequence / num_chunks], or a

nemo_automodel.components.models.mimo_v2_flash.cp._media_mask(
    input_ids: torch.Tensor,
    token_id: int | None,
    num_chunks: int
) -> torch.Tensor | None

Build a global placeholder mask before Transformer Engine shards tokens.

Parameters:

input_ids
torch.Tensor

Unsharded token IDs of shape [batch, sequence].

token_id
int | None

Image or video placeholder ID, or None when absent.

num_chunks
int

Number of pipeline token streams.

Returns: torch.Tensor | None

Boolean mask in global THD stream order, with shape [tokens] for one

nemo_automodel.components.models.mimo_v2_flash.cp.make_mimo_te_cp_sharder(
    model: torch.nn.Module,
    num_chunks: int,
    image_token_id: int | None,
    video_token_id: int | None
) -> nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder

Create MiMo’s thin adapter around the framework TE THD sharder.

Parameters:

model
torch.nn.Module

MiMo model or pipeline-local part to configure from the runtime CP mesh.

num_chunks
int

Number of pipeline microbatch streams.

image_token_id
int | None

Optional image placeholder token ID.

video_token_id
int | None

Optional video placeholder token ID.

Returns: ContextParallelSharder

An unresolved :class:ContextParallelSharder configured by the caller’s

nemo_automodel.components.models.mimo_v2_flash.cp.shard_batch_for_mimo_te(
    cp_mesh,
    tp_mesh,
    batch: dict[str, typing.Any],
    model: torch.nn.Module | None = None,
    loss_mask: torch.Tensor | None = None,
    padding_token_id: int = 0,
    num_chunks: int = 1,
    image_token_id: int | None = None,
    video_token_id: int | None = None
)

Delegate MiMo packed CP to the framework TE THD sharder.

The wrapper records global VLM placeholder masks and TE’s local-token index map before the framework mutates the packed batch. Those tensors let MiMo select exactly the image/video features owned by each DualChunkSwap shard. It also preserves the PP media side channel, which is intentionally not a token-aligned tensor and therefore must not be split by the THD helper.

Parameters:

cp_mesh

Optional one-dimensional context-parallel mesh.

tp_mesh

Unused tensor-parallel mesh required by the sharder protocol.

batch
dict[str, Any]

Packed batch whose token tensors have shape [batch, sequence].

model
torch.nn.Module | NoneDefaults to None

MiMo model or pipeline-local part whose TE attention is configured from the runtime CP mesh before the first forward.

loss_mask
torch.Tensor | NoneDefaults to None

Optional loss mask passed by the sharder protocol. Labels already carry the loss ignore value, so this is unsupported here.

padding_token_id
intDefaults to 0

Token ID used for physical THD padding.

num_chunks
intDefaults to 1

Number of pipeline microbatch streams.

image_token_id
int | NoneDefaults to None

Optional image placeholder token ID.

video_token_id
int | NoneDefaults to None

Optional video placeholder token ID.

Returns:

A null transport context, the TE-prepared batch, and its

nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_GLOBAL_IMAGE_MASK = '_mimo_global_image_mask'
 
nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_GLOBAL_VIDEO_MASK = '_mimo_global_video_mask'
 
nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_THD_LOCAL_INDICES = '_mimo_thd_local_indices'
 
nemo_automodel.components.models.mimo_v2_flash.cp._SEQ_LENS_PADDING_VALUE = -1000
 
nemo_automodel.components.models.mimo_v2_flash.cp._VLM_PP_MEDIA_KEY = '_vlm_pp_media_chunks'
 
nemo_automodel.components.models.mimo_v2_flash.cp.__all__ = ['_MIMO_GLOBAL_IMAGE_MASK', '_MIMO_GLOBAL_VIDEO_MASK', '_MIMO_THD_LOCAL_INDICES'...