> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.models.mimo_v2_flash.cp MiMo adapters for the framework Transformer Engine THD sharder. ## Module Contents ### Functions | Name | Description | | --------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ | | [`_chunk_partition_indices`](#nemo_automodel-components-models-mimo_v2_flash-cp-_chunk_partition_indices) | Reproduce TE's data-dependent THD token partition for every chunk. | | [`_flatten_chunks`](#nemo_automodel-components-models-mimo_v2_flash-cp-_flatten_chunks) | Flatten batch rows into one token stream per pipeline chunk. | | [`_media_mask`](#nemo_automodel-components-models-mimo_v2_flash-cp-_media_mask) | Build a global placeholder mask before Transformer Engine shards tokens. | | [`make_mimo_te_cp_sharder`](#nemo_automodel-components-models-mimo_v2_flash-cp-make_mimo_te_cp_sharder) | Create MiMo's thin adapter around the framework TE THD sharder. | | [`shard_batch_for_mimo_te`](#nemo_automodel-components-models-mimo_v2_flash-cp-shard_batch_for_mimo_te) | Delegate MiMo packed CP to the framework TE THD sharder. | ### Data [`_MIMO_GLOBAL_IMAGE_MASK`](#nemo_automodel-components-models-mimo_v2_flash-cp-_MIMO_GLOBAL_IMAGE_MASK) [`_MIMO_GLOBAL_VIDEO_MASK`](#nemo_automodel-components-models-mimo_v2_flash-cp-_MIMO_GLOBAL_VIDEO_MASK) [`_MIMO_THD_LOCAL_INDICES`](#nemo_automodel-components-models-mimo_v2_flash-cp-_MIMO_THD_LOCAL_INDICES) [`_SEQ_LENS_PADDING_VALUE`](#nemo_automodel-components-models-mimo_v2_flash-cp-_SEQ_LENS_PADDING_VALUE) [`_VLM_PP_MEDIA_KEY`](#nemo_automodel-components-models-mimo_v2_flash-cp-_VLM_PP_MEDIA_KEY) [`__all__`](#nemo_automodel-components-models-mimo_v2_flash-cp-__all__) ### API ```python 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 ```python 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 ```python 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 ```python 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 ```python 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 | None` — default: 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 | None` — default: None Optional loss mask passed by the sharder protocol. Labels already carry the loss ignore value, so this is unsupported here. --- **`padding_token_id`** `int` — default: 0 Token ID used for physical THD padding. --- **`num_chunks`** `int` — default: 1 Number of pipeline microbatch streams. --- **`image_token_id`** `int | None` — default: None Optional image placeholder token ID. --- **`video_token_id`** `int | None` — default: None Optional video placeholder token ID. --- **Returns:** A null transport context, the TE-prepared batch, and its ```python nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_GLOBAL_IMAGE_MASK = '_mimo_global_image_mask' ``` ```python nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_GLOBAL_VIDEO_MASK = '_mimo_global_video_mask' ``` ```python nemo_automodel.components.models.mimo_v2_flash.cp._MIMO_THD_LOCAL_INDICES = '_mimo_thd_local_indices' ``` ```python nemo_automodel.components.models.mimo_v2_flash.cp._SEQ_LENS_PADDING_VALUE = -1000 ``` ```python nemo_automodel.components.models.mimo_v2_flash.cp._VLM_PP_MEDIA_KEY = '_vlm_pp_media_chunks' ``` ```python nemo_automodel.components.models.mimo_v2_flash.cp.__all__ = ['_MIMO_GLOBAL_IMAGE_MASK', '_MIMO_GLOBAL_VIDEO_MASK', '_MIMO_THD_LOCAL_INDICES'... ```