nemo_automodel.components.models.gemma4_moe.loss

Loss adapters for Gemma4 tensor parallelism.

Module Contents

Classes

Name	Description
`Gemma4TensorParallelFusedLinearCrossEntropy`	Fused linear CE for Gemma4’s tied, vocabulary-sharded LM head.

API

class nemo_automodel.components.models.gemma4_moe.loss.Gemma4TensorParallelFusedLinearCrossEntropy()

Bases: FusedLinearCrossEntropy

Fused linear CE for Gemma4’s tied, vocabulary-sharded LM head.

Gemma4 E-series TP keeps the tied embedding/LM-head weight sharded on the TP vocabulary axis. Cut cross entropy needs a rank-local full vocabulary weight, but the generic loss intentionally rejects a DTensor mesh that does not match the DP/CP loss-reduction group. This adapter gathers only the Gemma4 TP layout and marks its backward gradient replicated: TP peers see the same tokens and therefore must not sum duplicate weight gradients. FSDP continues to reduce independent DP/CP contributions normally.

nemo_automodel.components.models.gemma4_moe.loss.Gemma4TensorParallelFusedLinearCrossEntropy.materialize_lm_weight(
    lm_weight: torch.Tensor,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor

staticmethod

Materialize a Gemma4 TP LM head with correct TP gradient semantics.

Parameters:

lm_weight

torch.Tensor

Regular or DTensor LM-head weight shaped [vocab, hidden].

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

DP/CP group whose ranks own independent tokens. It remains owned by FSDP when the visible DTensor mesh contains only TP, as it does while an FSDP unit is unsharded.

Returns: torch.Tensor

A rank-local full-vocabulary tensor. Backward slices, rather than