nemo_automodel.shared.matrix_layout | NVIDIA NeMo AutoModel

Storage-orientation contract between grouped-GEMM modules and layout-aware optimizers.

Module Contents

Classes

Name	Description
`TransposedMatrixModule`	Marker base class for modules whose matrix parameters store `[..., input, output]`.

API

class nemo_automodel.shared.matrix_layout.TransposedMatrixModule()

Marker base class for modules whose matrix parameters store [..., input, output].

nn.Linear stores [output, input]. Grouped-GEMM expert modules (GroupedExperts, GroupedExpertsDeepEP and their LoRA and MXFP4 subclasses) store every trainable matrix parameter, including LoRA factors, as [experts, input, output] so each expert computes x @ W. Optimizers whose update depends on the output-neuron axis (Muown) check isinstance(module, TransposedMatrixModule) to normalize along that axis. Biases are not matrices and are routed by the optimizer’s own rules. The class carries no state or methods.