nemo_automodel.components.models.gemma4_moe.cp_batch

Gemma4’s aux-only contiguous-shard context-parallel batch prep.

Gemma4 runs its own p2p ring FlexAttention over contiguous per-rank sequence slices (no collective — the transport lives in Gemma4’s attention, see cp_attention.py). Under the sunk (Megatron-style per-microbatch) CP path the model embeds, splices vision, builds per_layer_inputs and the flex-ring mask metadata inside its forward and contiguously slices them there; the dispatch-time sharder therefore only touches the no-grad auxiliary streams.

The generic slicing lives in components/distributed/context_parallel/sharder.py; this module owns the one Gemma4-specific piece the aux-only shard still needs: the _packed_seq_ids document-boundary synthesis its manual CP attention mask builder requires (its pad-region zeros depend on the global pad tail, which the forward — holding only this rank’s slice — cannot reconstruct). Gemma4’s prepare_model_inputs_for_cp exposes it through the ContextParallelSharder it returns under the "cp_sharder" batch key, which the CP dispatch invokes in place of the default load-balanced context_parallel path.

Module Contents

Functions

Name	Description
`_synthesize_single_document_seq_ids`	Materialize the trivial single-document `_packed_seq_ids` map.
`make_contiguous_aux_only_shard_cp_batch_and_ctx`	Aux-only contiguous CP shard for Gemma4’s sunk (in-forward) pre-embed.

API

nemo_automodel.components.models.gemma4_moe.cp_batch._synthesize_single_document_seq_ids(
    batch: dict,
    seq_len: int
) -> None

Materialize the trivial single-document _packed_seq_ids map.

Collates emit _packed_seq_ids only when 2+ documents are packed, but Gemma4’s manual CP attention mask builder needs document boundaries even for one document (1 = real token, 0 = pad). Derived from padding_mask when present, else all-ones. A no-op when _packed_seq_ids already exists (genuinely packed input).

Parameters:

batch

dict

The CP batch dict; mutated in place to add _packed_seq_ids.

seq_len

int

The pre-pad sequence length.

nemo_automodel.components.models.gemma4_moe.cp_batch.make_contiguous_aux_only_shard_cp_batch_and_ctx(
    cp_mesh,
    tp_mesh,
    batch,
    loss_mask = None,
    padding_token_id: int = 0,
    extra_seq_keys: dict[str, int] | None = None,
    extra_pad_values: dict[str, typing.Any] | None = None
)

Aux-only contiguous CP shard for Gemma4’s sunk (in-forward) pre-embed.

Exposed as ContextParallelSharder.shard_batch by Gemma4’s sharder-only prepare_model_inputs_for_cp. It shards only the no-grad auxiliary streams (labels / position_ids / loss_mask / padding_mask plus the synthesized _packed_seq_ids document map) and leaves input_ids / pixel_values / mm_token_type_ids FULL-length in the batch. The model forward then embeds, splices vision, builds per_layer_inputs and the _gemma4_vision_group_ids / mm_token_type_ids ring metadata on the full sequence and contiguously slices them per microbatch (see shard_sequence_for_cp_contiguous), so the embeddings and vision tower are trainable under CP and the PP×CP shared pre-embed graph no longer exists.

_packed_seq_ids is synthesized full (from the full padding_mask) and sharded here rather than in the forward: its pad-region zeros depend on the global pad tail, which the forward — holding only this rank’s slice — cannot reconstruct. Every other flex-ring mask input the forward owns is a pure per-token or cumsum-over-full-then-slice quantity, so it slices to the same contiguous layout this sharder applies.