nemo_automodel.components.models.gemma4_moe.parallelization

Gemma4-owned tensor-parallel plan and FSDP2 strategy registration.

Module Contents

Classes

Name	Description
`_Gemma4RowwiseParallel`	Shard a Gemma4 linear on input features with a synchronous reduction.
`_Gemma4VocabParallelEmbedding`	Shard a Gemma4 embedding by vocabulary without DTensor `MaskPartial`.
`_ReduceFromTensorParallelRegion`	Sum local TP values in forward and leave replicated gradients local.

Functions

Name	Description
`_gemma4_rowwise_linear_forward`	Apply one Gemma4 row-parallel linear and synchronously reduce its output.
`_gemma4_tp_plan`	Return the TP plan for the concrete Gemma4 checkpoint variant.
`_gemma4_vocab_parallel_forward`	Look up replicated token ids in one Gemma4 vocabulary shard.
`register_gemma4_parallel_strategy`	Register Gemma4’s model-owned FSDP2 strategy once.

API

class nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4RowwiseParallel(
    output_holder: list[torch.Tensor | None]
)

Bases: ParallelStyle

Shard a Gemma4 linear on input features with a synchronous reduction.

nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4RowwiseParallel._apply(
    module: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.nn.Module

Partition [output, input] weight as Shard(1) and install the local linear.

class nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4VocabParallelEmbedding()

Bases: ParallelStyle

Shard a Gemma4 embedding by vocabulary without DTensor MaskPartial.

The input ids are replicated [batch, sequence]. The weight changes from [vocab, embedding] to a DTensor with Shard(0); the forward returns a local, replicated [batch, sequence, embedding] Tensor.

nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4VocabParallelEmbedding._apply(
    module: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.nn.Module

Partition one Gemma4 embedding weight and install its local lookup.

class nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion()

Bases: Function

Sum local TP values in forward and leave replicated gradients local.

nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion.backward(
    ctx,
    grad_output: torch.Tensor
)

staticmethod

Pass replicated gradients unchanged to each owning TP shard.

nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion.forward(
    ctx,
    local_output: torch.Tensor,
    process_group
) -> torch.Tensor

staticmethod

Reduce a local contribution Tensor into one replicated Tensor.

nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_rowwise_linear_forward(
    module: torch.nn.Linear,
    input_: torch.Tensor
) -> torch.Tensor

Apply one Gemma4 row-parallel linear and synchronously reduce its output.

Parameters:

module

nn.Linear

Linear whose weight is a DTensor sharded on input features.

input_

torch.Tensor

Local Tensor shaped [..., input_features / tp_size].

Returns: torch.Tensor

Local Tensor shaped [..., output_features], replicated across the

nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_tp_plan(
    model: torch.nn.Module,
    sequence_parallel: bool = False
) -> dict[str, torch.distributed.tensor.parallel.ParallelStyle]

Return the TP plan for the concrete Gemma4 checkpoint variant.

E2B/E4B add a packed per-layer embedding table that is absent from 31B. Row-sharding that table by vocabulary is important: at E4B dimensions it is roughly 2.8 billion parameters, so leaving it replicated defeats much of the memory benefit of TP.

Parameters:

model

nn.Module

Gemma4 conditional-generation model whose text config selects the dense variant.

sequence_parallel

boolDefaults to False

Whether sequence parallelism was requested. Gemma4 does not currently support it, so the request is ignored.

Returns: dict[str, ParallelStyle]

Mapping from Gemma4 module paths to PyTorch parallel styles.

nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_vocab_parallel_forward(
    module: torch.nn.Embedding,
    input_ids: torch.Tensor
) -> torch.Tensor

Look up replicated token ids in one Gemma4 vocabulary shard.

Parameters:

module

nn.Embedding

Gemma4 embedding whose weight is a DTensor sharded on vocabulary.

input_ids

torch.Tensor

Replicated integer Tensor shaped [batch, sequence].

Returns: torch.Tensor

Local Tensor shaped [batch, sequence, embedding], replicated across

nemo_automodel.components.models.gemma4_moe.parallelization.register_gemma4_parallel_strategy() -> None