Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::RmsNormConfig Class Referenceexport
Inheritance diagram for Mila::Dnn::RmsNormConfig:
Mila::Dnn::ComponentConfig

Public Member Functions

 RmsNormConfig (int64_t axis)
 Construct in axis mode.
 RmsNormConfig (shape_t normalized_shape)
 Construct in shape mode.
void fromMetadata (const SerializationMetadata &meta) override
 Populate configuration from provided metadata.
std::optional< int64_t > getAxis () const noexcept
float getEpsilon () const noexcept
const shape_tgetNormalizedShape () const noexcept
float getUnitOffset () const noexcept
bool hasBias () const noexcept
bool hasNormalizedShape () const noexcept
SerializationMetadata toMetadata () const override
 Convert configuration into a SerializationMetadata object.
std::string toString () const override
 Produce a short, human-readable summary of the configuration.
void validate () const override
 Validate configuration parameters.
template<typename Self>
decltype(auto) withBias (this Self &&self, bool has_bias)
 Enable or disable learnable bias.
template<typename Self>
decltype(auto) withEpsilon (this Self &&self, float epsilon)
 Set epsilon for numerical stability.
template<typename Self>
decltype(auto) withUnitOffset (this Self &&self, float unit_offset)
 Set the unit offset added to the loaded weight before scaling.
Public Member Functions inherited from Mila::Dnn::ComponentConfig
virtual ~ComponentConfig ()=default
 Virtual destructor for polymorphic base.

Constructor & Destructor Documentation

◆ RmsNormConfig() [1/2]

Mila::Dnn::RmsNormConfig::RmsNormConfig ( shape_t normalized_shape)
inlineexplicit

Construct in shape mode.

Normalizes over the trailing dimensions described by normalized_shape.

Parameters
normalized_shapeTrailing dimensions to normalize over (e.g. shape_t{ model_dim }).

◆ RmsNormConfig() [2/2]

Mila::Dnn::RmsNormConfig::RmsNormConfig ( int64_t axis)
inlineexplicit

Construct in axis mode.

Normalizes over a single axis.

Parameters
axisAxis along which to normalize (negative indexing supported).

Member Function Documentation

◆ fromMetadata()

void Mila::Dnn::RmsNormConfig::fromMetadata ( const SerializationMetadata & meta)
inlineoverridevirtual

Populate configuration from provided metadata.

Implementations should read available keys and leave missing keys at their current/default values to preserve forward/backward compatibility.

Parameters
metaMetadata to read configuration values from.

Implements Mila::Dnn::ComponentConfig.

◆ toMetadata()

SerializationMetadata Mila::Dnn::RmsNormConfig::toMetadata ( ) const
inlineoverridevirtual

Convert configuration into a SerializationMetadata object.

Implementations should include any fields required to fully reconstruct the configuration via fromMetadata.

Returns
SerializationMetadata Metadata representation of the config.

Implements Mila::Dnn::ComponentConfig.

◆ toString()

std::string Mila::Dnn::RmsNormConfig::toString ( ) const
inlineoverridevirtual

Produce a short, human-readable summary of the configuration.

Implementations should return a compact, single-line description suitable for logging and debugging.

Returns
std::string Human-readable summary of the configuration.

Implements Mila::Dnn::ComponentConfig.

◆ validate()

void Mila::Dnn::RmsNormConfig::validate ( ) const
inlineoverridevirtual

Validate configuration parameters.

Called by callers to ensure the configuration represents a valid, constructible component. Implementations must throw std::invalid_argument (or a derived exception) when validation fails.

Exceptions
std::invalid_argumentIf the configuration is invalid.

Implements Mila::Dnn::ComponentConfig.

◆ withBias()

template<typename Self>
decltype(auto) Mila::Dnn::RmsNormConfig::withBias ( this Self && self,
bool has_bias )
inline

Enable or disable learnable bias.

Default: true. Llama 3 uses false.

◆ withEpsilon()

template<typename Self>
decltype(auto) Mila::Dnn::RmsNormConfig::withEpsilon ( this Self && self,
float epsilon )
inline

Set epsilon for numerical stability.

Default: 1e-5f. Llama 3 uses 1e-5f; some models use 1e-6f.

◆ withUnitOffset()

template<typename Self>
decltype(auto) Mila::Dnn::RmsNormConfig::withUnitOffset ( this Self && self,
float unit_offset )
inline

Set the unit offset added to the loaded weight before scaling.

The normalized activation is scaled by (weight + unit_offset). Default 0.0 reproduces standard RMSNorm (x_norm * weight) – used by Llama 3 / GPT-2. Gemma sets 1.0: its RMSNorm is x_norm * (1 + weight), with weights stored raw (zero-centered, weight-decay-friendly). The offset is applied at the kernel so the stored/loaded weights remain identical to the source checkpoint.


The documentation for this class was generated from the following file: