Network-level configuration for Gemma 4 transformer networks. More...

Inheritance diagram for Mila::Dnn::GemmaConfig:

Public Member Functions
	GemmaConfig (dim_t embedding_dim, dim_t num_layers)
	Construct a Gemma network configuration.
void	fromMetadata (const SerializationMetadata &meta) override
	Populate configuration from provided metadata.
float	getEmbeddingScale () const noexcept
	Embedding scale applied after token lookup: sqrt(embedding_dim).
float	getFinalLogitSoftcapping () const noexcept
dim_t	getGlobalHeadDim () const noexcept
	Global-layer per-head dimension (falls back to head_dim when unset).
dim_t	getGlobalKVProjectionWidth () const noexcept
	Global-layer KV-projection width = num_global_kv_heads * global_head_dim.
dim_t	getGlobalPackedQKVWidth () const noexcept
	Global-layer packed QKV trailing dimension.
dim_t	getGlobalQProjectionWidth () const noexcept
	Global-layer Q-projection width = num_heads * global_head_dim.
dim_t	getGlobalRotaryDim () const noexcept
dim_t	getHeadDim () const noexcept
	Per-head dimension, decoupled from the residual stream.
dim_t	getHeadDimForLayer (dim_t layer_index) const noexcept
dim_t	getHiddenDimension () const noexcept
dim_t	getKVProjectionWidth () const noexcept
	KV-projection output width = num_kv_heads * head_dim.
dim_t	getMaxSequenceLength () const noexcept
dim_t	getModelDim () const noexcept
	Residual-stream dimension (HuggingFace hidden_size).
dim_t	getNumGlobalKVHeads () const noexcept
	Global-layer KV-head count (falls back to num_kv_heads when unset).
dim_t	getNumHeads () const noexcept
dim_t	getNumKVHeads () const noexcept
dim_t	getNumKVHeadsForLayer (dim_t layer_index) const noexcept
dim_t	getNumLayers () const noexcept
dim_t	getPackedQKVWidth () const noexcept
	Packed QKV trailing dimension = (num_heads + 2 * num_kv_heads) * head_dim.
dim_t	getPackedQKVWidthForLayer (dim_t layer_index) const noexcept
	Packed fused-QKV width for this layer (K=V global layers drop the V section).
dim_t	getQProjectionWidth () const noexcept
	Q-projection output width = num_heads * head_dim.
dim_t	getQProjectionWidthForLayer (dim_t layer_index) const noexcept
	Q-projection width for this layer = num_heads * head_dim(layer).
float	getRMSNormEpsilon () const noexcept
float	getRoPEThetaForLayer (dim_t layer_index) const noexcept
float	getRoPEThetaGlobal () const noexcept
float	getRoPEThetaLocal () const noexcept
dim_t	getRotaryDimForLayer (dim_t layer_index) const noexcept
	Global layers use proportional partial-rotary; sliding layers rotate fully (0).
dim_t	getSlidingWindowPattern () const noexcept
bool	getTieWordEmbeddings () const noexcept
dim_t	getVocabSize () const noexcept
dim_t	getWindow () const noexcept
dim_t	getWindowForLayer (dim_t layer_index) const noexcept
	Sliding layers carry the window; global layers are unbounded (window 0).
bool	isGlobalLayer (dim_t layer_index) const noexcept
	True when layer_index is a global (full-attention) layer.
bool	keyEqualsValue () const noexcept
bool	keyEqualsValueForLayer (dim_t layer_index) 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)	withFinalLogitSoftcapping (this Self &&self, float cap)
	Final logit soft-cap value (Gemma 4: 30.0; 0 disables).
template<typename Self>
decltype(auto)	withGlobalHeadDim (this Self &&self, dim_t global_head_dim)
	Set the global-layer per-head dimension (Gemma 4 12B: 512).
template<typename Self>
decltype(auto)	withGlobalRoPETheta (this Self &&self, float theta)
	RoPE base for global (full) layers (Gemma 4: 1e6).
template<typename Self>
decltype(auto)	withGlobalRotaryDim (this Self &&self, dim_t global_rotary_dim)
	Rotated dimension count for global layers (proportional partial-rotary).
template<typename Self>
decltype(auto)	withHeadDim (this Self &&self, dim_t head_dim)
	Set the per-head dimension, decoupled from the residual stream.
template<typename Self>
decltype(auto)	withHiddenDimension (this Self &&self, dim_t hidden_dim)
	Set the FFN intermediate dimension (GeGLU hidden width).
template<typename Self>
decltype(auto)	withKeyEqualsValue (this Self &&self, bool key_equals_value)
	Set whether global layers share K=V (no separate v_proj).
template<typename Self>
decltype(auto)	withMaxSequenceLength (this Self &&self, dim_t max_seq_len)
	Set the trained maximum sequence length (HuggingFace max_position_embeddings).
template<typename Self>
decltype(auto)	withNumGlobalKVHeads (this Self &&self, dim_t num_global_kv_heads)
	Set the global-layer KV-head count (Gemma 4 12B: 1, MQA).
template<typename Self>
decltype(auto)	withNumHeads (this Self &&self, dim_t num_heads)
template<typename Self>
decltype(auto)	withNumKVHeads (this Self &&self, dim_t num_kv_heads)
template<typename Self>
decltype(auto)	withRMSNormEpsilon (this Self &&self, float eps)
template<typename Self>
decltype(auto)	withRoPETheta (this Self &&self, float theta)
	RoPE base for local (sliding) layers (Gemma 4: 10000).
template<typename Self>
decltype(auto)	withSlidingWindowPattern (this Self &&self, dim_t pattern)
	Sliding/global interleave period.
template<typename Self>
decltype(auto)	withTieWordEmbeddings (this Self &&self, bool tie)
	Whether lm_head shares the token embedding table (weight tying).
template<typename Self>
decltype(auto)	withVocabularyLength (this Self &&self, dim_t vocab_size)
template<typename Self>
decltype(auto)	withWindow (this Self &&self, dim_t window)
	Sliding-window size for local (sliding) layers (Gemma 4 12B: 1024).
Public Member Functions inherited from Mila::Dnn::ComponentConfig
virtual	~ComponentConfig ()=default
	Virtual destructor for polymorphic base.

Detailed Description

Network-level configuration for Gemma 4 transformer networks.

Carries the network geometry with head_dim decoupled from the residual stream. The Q-projection width (num_heads * head_dim), the KV-projection width (num_kv_heads * head_dim), and the packed QKV width are derived on demand so they always stay consistent with the primary fields and remain correct when num_heads * head_dim != embedding_dim (the Gemma case).

Fluent setters follow the C++23 explicit-object-parameter pattern used throughout the codebase.

Constructor & Destructor Documentation

◆ GemmaConfig()

Mila::Dnn::GemmaConfig::GemmaConfig	(	dim_t	embedding_dim,
		dim_t	num_layers )

inline

Construct a Gemma network configuration.

Parameters

embedding_dim	Residual-stream dimension (HuggingFace hidden_size). Must be > 0.
num_layers	Number of transformer layers. Must be > 0.

Member Function Documentation

◆ fromMetadata()

void Mila::Dnn::GemmaConfig::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

meta	Metadata to read configuration values from.

Implements Mila::Dnn::ComponentConfig.

◆ getEmbeddingScale()

float Mila::Dnn::GemmaConfig::getEmbeddingScale ( ) const

inlinenoexcept

Embedding scale applied after token lookup: sqrt(embedding_dim).

Gemma multiplies the embedded hidden states by this factor; Llama does not.

◆ getGlobalPackedQKVWidth()

dim_t Mila::Dnn::GemmaConfig::getGlobalPackedQKVWidth ( ) const

inlinenoexcept

Global-layer packed QKV trailing dimension.

With K=V (the Gemma global case) the value section is absent, so the width is (num_heads + num_global_kv_heads) * global_head_dim; otherwise it is (num_heads + 2 * num_global_kv_heads) * global_head_dim.

◆ getHeadDim()

dim_t Mila::Dnn::GemmaConfig::getHeadDim ( ) const

inlinenoexcept

Per-head dimension, decoupled from the residual stream.

Returns the explicitly-set head_dim, or falls back to embedding_dim / num_heads when unset (head_dim_ == 0).

◆ getModelDim()

dim_t Mila::Dnn::GemmaConfig::getModelDim ( ) const

inlinenoexcept

Residual-stream dimension (HuggingFace hidden_size).

For Gemma this is NOT num_heads * head_dim. The attention output (num_heads * head_dim) projects back down to this width via a non-square output projection.

◆ getPackedQKVWidth()

dim_t Mila::Dnn::GemmaConfig::getPackedQKVWidth ( ) const

inlinenoexcept

Packed QKV trailing dimension = (num_heads + 2 * num_kv_heads) * head_dim.

The fused QKV projection output width for a standard (non-K=V) layer. The global-layer K=V variant (Step 1) uses (num_heads + num_kv_heads)

head_dim instead.

◆ getQProjectionWidth()

dim_t Mila::Dnn::GemmaConfig::getQProjectionWidth ( ) const

inlinenoexcept

Q-projection output width = num_heads * head_dim.

This is the value to feed as GqaConfig/RopeConfig's first constructor argument (both derive head_dim from it). For Gemma it differs from the residual stream (4096 vs 3840 for the sliding geometry).

◆ isGlobalLayer()

bool Mila::Dnn::GemmaConfig::isGlobalLayer ( dim_t layer_index ) const

inlinenoexcept

True when layer_index is a global (full-attention) layer.

Pattern N (sliding_window_pattern) makes every Nth layer global; Gemma 4's period of 6 yields the 5 sliding : 1 global interleave with a global final layer.

◆ toMetadata()

SerializationMetadata Mila::Dnn::GemmaConfig::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::GemmaConfig::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::GemmaConfig::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_argument If the configuration is invalid.

Implements Mila::Dnn::ComponentConfig.

◆ withGlobalHeadDim()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withGlobalHeadDim	(	this Self &&	self,
		dim_t	global_head_dim )

inline

Set the global-layer per-head dimension (Gemma 4 12B: 512).

Applies to the 1-in-6 full-attention layers. Default 0 means "use the sliding head_dim" (no distinct global width).

◆ withGlobalRotaryDim()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withGlobalRotaryDim	(	this Self &&	self,
		dim_t	global_rotary_dim )

inline

Rotated dimension count for global layers (proportional partial-rotary).

Gemma 4: 128 of global_head_dim 512. 0 = full rotation.

◆ withHeadDim()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withHeadDim	(	this Self &&	self,
		dim_t	head_dim )

inline

Set the per-head dimension, decoupled from the residual stream.

This is the Gemma 4 departure from Llama: head_dim is an independent field, not embedding_dim / num_heads. Default 0 means "derive from embedding_dim / num_heads" (the Llama-compatible fallback); Gemma sets it explicitly (256 for the sliding/base geometry).

◆ withKeyEqualsValue()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withKeyEqualsValue	(	this Self &&	self,
		bool	key_equals_value )

inline

Set whether global layers share K=V (no separate v_proj).

When true, the global-layer packed QKV width drops to (num_heads + num_global_kv_heads) * global_head_dim; the block aliases the value projection to the key projection.

◆ withNumGlobalKVHeads()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withNumGlobalKVHeads	(	this Self &&	self,
		dim_t	num_global_kv_heads )

inline

Set the global-layer KV-head count (Gemma 4 12B: 1, MQA).

Default 0 means "use num_kv_heads".

◆ withSlidingWindowPattern()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withSlidingWindowPattern	(	this Self &&	self,
		dim_t	pattern )

inline

Sliding/global interleave period.

Pattern N means every Nth layer is global (full attention) and the rest are sliding; Gemma 4 uses 6 (5 sliding : 1 global), which also makes the final layer global for layer counts that are multiples of the period.

◆ withTieWordEmbeddings()

template<typename Self>

decltype(auto) Mila::Dnn::GemmaConfig::withTieWordEmbeddings	(	this Self &&	self,
		bool	tie )

inline

Whether lm_head shares the token embedding table (weight tying).

Set from checkpoint metadata; when true the transformer installs the shared table into lm_head BEFORE build so the head never allocates its own weight (WeightTying.md).

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

Mila/Src/Dnn/Components/Transformers/Gemma/Gemma.Config.ixx

Public Member Functions

Detailed Description

Constructor & Destructor Documentation

◆ GemmaConfig()

Member Function Documentation

◆ fromMetadata()

◆ getEmbeddingScale()

◆ getGlobalPackedQKVWidth()

◆ getHeadDim()

◆ getModelDim()

◆ getPackedQKVWidth()

◆ getQProjectionWidth()

◆ isGlobalLayer()

◆ toMetadata()

◆ toString()

◆ validate()

◆ withGlobalHeadDim()

◆ withGlobalRotaryDim()

◆ withHeadDim()

◆ withKeyEqualsValue()

◆ withNumGlobalKVHeads()

◆ withSlidingWindowPattern()

◆ withTieWordEmbeddings()