Struct Transformer 

pub struct Transformer {
    pub config: TransformerConfig,
    pub embed_tokens: Embedding,
    pub layers: Vec<TransformerBlock>,
    pub norm: RMSNorm,
    pub lm_head: Option<Tensor>,
}

Complete transformer model

Fields

config: TransformerConfig

Configuration

embed_tokens: Embedding

Token embedding layer

layers: Vec<TransformerBlock>

Transformer layers

norm: RMSNorm

Final layer normalization

lm_head: Option<Tensor>

Language model head (tied to embeddings or separate)

Implementations

impl Transformer

pub fn new(config: &TransformerConfig) -> Self

Create new transformer with initialized weights

pub fn from_params( config: &TransformerConfig, params: &HashMap<String, Tensor>, ) -> Option<Self>

Create transformer from parameter map

Expected parameter names (following HuggingFace LLaMA convention):

• model.embed_tokens.weight
• model.layers.{i}.*
• model.norm.weight
• lm_head.weight (optional, uses tied embeddings if not present)

pub fn from_safetensors( model_path: impl AsRef<Path>, config: &TransformerConfig, ) -> Result<Self>

Load transformer from SafeTensors file(s)

Reads SafeTensors weights from model_path, converts BF16/F16 to F32, validates shapes against config, checks for NaN/Inf, and constructs the complete Transformer.

Arguments

• model_path - Path to model directory or single SafeTensors file
• config - Transformer configuration specifying model dimensions

Errors

Returns Error::ConfigError if:

• No SafeTensors files found
• Required weight tensors are missing
• Weight shapes do not match config dimensions
• Weights contain NaN or Inf values
• Layer count does not match config

pub fn from_apr( apr_path: impl AsRef<Path>, config: &TransformerConfig, ) -> Result<Self>

Load transformer from APR file (.apr format)

Reads tensor data from an APR binary file, dequantizing from any stored dtype (F16, Q4K, etc.) to F32. Uses the same validation pipeline as from_safetensors: structural, shape, and NaN/Inf checks.

Arguments

• apr_path - Path to the .apr model file
• config - Transformer configuration (typically read from APR metadata)

Errors

Returns Error::ConfigError if tensors are missing, shapes mismatch, or weights contain NaN/Inf values.

pub fn forward(&self, token_ids: &[u32]) -> Tensor

Forward pass for language modeling

Arguments

• token_ids - Input token IDs

Returns

Logits tensor (seq_len * vocab_size, flattened)

pub fn forward_hidden(&self, token_ids: &[u32]) -> Tensor

Forward pass returning hidden states (before lm_head)

Arguments

• token_ids - Input token IDs

Returns

Hidden states tensor (seq_len * hidden_size, flattened)

pub fn forward_hidden_with_lora( &self, token_ids: &[u32], lora_layers: &[LoRALayer], ) -> Tensor

Forward pass returning hidden states with LoRA adjusts (KAIZEN-011)

Like forward_hidden but applies LoRA adapters to Q/V projections in each transformer layer’s attention. Enables non-CUDA LoRA training by putting LoRA parameters into the autograd graph.

Arguments

• token_ids - Input token IDs
• lora_layers - LoRA layers in [Q_0, V_0, Q_1, V_1, …] order

Returns

Hidden states tensor (seq_len * hidden_size, flattened)

pub fn forward_with_lora( &self, token_ids: &[u32], lora_layers: &[LoRALayer], ) -> Tensor

Forward pass with LoRA adapters (ENT-LoRA-001)

Like forward but applies LoRA adapters to Q/V projections. Returns full logits (seq_len * vocab_size).

Arguments

• token_ids - Input token IDs
• lora_layers - LoRA layers in [Q_0, V_0, Q_1, V_1, …] order

pub fn forward_last(&self, token_ids: &[u32]) -> Tensor

Get the last token’s logits (for generation)

pub fn parameters(&self) -> Vec<&Tensor>

Get all parameters as a vector

pub fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Get all parameters as mutable references for optimizer

pub fn config(&self) -> &TransformerConfig

Get configuration

pub fn embed_token(&self, token_id: u32) -> Vec<f32>

Embed a single token, returning hidden_size floats.

pub fn output_norm_weight_slice(&self) -> &[f32]

Get the output norm weight as a slice.

pub fn lm_head_weight_slice(&self) -> &[f32]

Get the lm_head weight as a slice (vocab_size × hidden_size, row-major).

pub fn lm_head_weight(&self) -> &Tensor

Get the language model head weight tensor.

Returns the dedicated lm_head weight if present, otherwise falls back to tied embedding weights.

pub fn named_parameters(&self) -> Vec<(String, &Tensor)>

Get named parameters for checkpoint serialization.

Returns (name, tensor) pairs matching HuggingFace weight conventions. This handles variable parameter counts (e.g., models with/without attention biases) correctly, unlike the hardcoded 9-params-per-layer assumption.

pub fn set_named_parameter(&mut self, name: &str, value: Tensor) -> bool

ENT-282: Set a named parameter by name (for delta checkpoint overlay).

Returns true if the parameter was found and set.