bridge.models.common.te_layers#
Shared Transformer Engine layer extensions for Post-LayerNorm architectures.
Module Contents#
Classes#
Row-parallel linear with an additional Post-LayerNorm on the output. |
API#
- class bridge.models.common.te_layers.TERowParallelLinearLayerNorm(
- input_size: int,
- output_size: int,
- *,
- config: megatron.core.transformer.TransformerConfig,
- **kwargs,
Bases:
bridge.models.common.te_layers.TERowParallelLinearRow-parallel linear with an additional Post-LayerNorm on the output.
Used by models that attach a Post-LN module to row-parallel projection outputs, such as Gemma2, Gemma3, and EXAONE 4.0.
Initialization
- forward(x)#
Forward with additional Post-LN on output.