Burn-LM AI Engine

Use the #[derive(Config)] macro to define how model components and variants (e.g., different parameter sizes) are initialized.

Key Features of Config:

• Default values: Use #[config(default = ...)] to specify defaults.
• Automatic constructors: Provides new(...) and with_ methods for fields.
• Serialization: Enables easy saving and loading of configuration files.

Pattern:

1. Define a Config struct for each module.
2. Implement an init method on the config struct that takes a &Device<B> and returns the initialized module.
3. For complex models, create a top-level composite config that delegates initialization to the sub-module configs.

#[derive(Config, Debug)]
pub struct FeedForwardConfig {
    pub d_model: usize,
    pub hidden_size: usize,
}

impl FeedForwardConfig {
    pub fn init<B: Backend>(&self, device: &Device<B>) -> FeedForward<B> {
        let swiglu = SwiGluConfig::new(self.d_model, self.hidden_size)
            .with_bias(false)
            .init(device);
        let w2 = LinearConfig::new(self.hidden_size, self.d_model)
            .with_bias(false)
            .init(device);

        FeedForward { swiglu, w2 }
    }
}

To implement a model architecture, define custom modules using structs or enums that derive the Module trait.

Best Practice: While not mandatory, closely following the architecture and naming structure of the original source model (e.g., PyTorch) makes importing weights significantly easier by reducing the amount of remapping logic required.

Implementation Steps:

1. Define the struct with fields representing the model's components (e.g., Linear, Embedding, RmsNorm).
2. Implement the forward method (or a similarly named method) to define the data flow through the module.

Note: Pay special attention to tokenization and positional encoding, as discrepancies here can lead to subtle bugs.

/// Feed-forward transformation network.
#[derive(Module, Debug)]
pub struct FeedForward<B: Backend> {
    swiglu: SwiGlu<B>,
    w2: Linear<B>,
}

impl<B: Backend> FeedForward<B> {
    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
        self.w2.forward(self.swiglu.forward(input))
    }
}

The TransformerCache manages the Key-Value states for all layers to enable efficient autoregressive generation.

• Initialization: Use TransformerCache::new(config, max_batch_size, device).
• Preparation: Call prepare(seq_len) before a forward pass. This updates the current sequence length and returns an optional attention mask (Option<Tensor<4, Bool>>).
• Error Handling: prepare returns Err(GenerationError::MaxSequenceLengthExceeded) if the sequence length exceeds the configured max_seq_len.
• Resetting: Call reset() to clear all cached states and set curr_seq_len to 0.

let mut cache = TransformerCache::new(&config, batch_size, &device);

// Prepare for a new sequence of length 'n'
match cache.prepare(n) {
    Ok(Some(mask)) => { /* use mask in transformer.forward */ },
    Ok(None) => { /* no mask needed (e.g. seq_len <= 1) */ },
    Err(e) => { /* handle MaxSequenceLengthExceeded */ },
}

// Clear cache for a new session
cache.reset();

The LlamaConfig provides high-level helper functions to initialize a model and load weights from a checkpoint in one step. These functions return an inference::Llama<T> instance.

Llama-3 Series (requires llama3 feature)

These methods use the Tiktoken tokenizer:

• load_llama3_2_3b(checkpoint, tokenizer_path, max_seq_len, device)
• load_llama3_2_1b(checkpoint, tokenizer_path, max_seq_len, device)
• load_llama3_1_8b(checkpoint, tokenizer_path, max_seq_len, device)
• load_llama3_8b(checkpoint, tokenizer_path, max_seq_len, device)

TinyLlama (requires tiny feature)

This method uses the SentencePieceTokenizer:

• load_tiny_llama(checkpoint, tokenizer_path, max_seq_len, device)

// Example: Loading a pre-trained Llama 3.2 1B model
let llama = LlamaConfig::load_llama3_2_1b(
    "path/to/checkpoint",
    "path/to/tokenizer",
    2048,
    &device
).expect("Failed to load model");

An InferenceJob represents a unit of work to be performed by the inference server. To run a job, use InferenceJob::create, which returns both the InferenceJob (used to emit generated items) and a JobHandle (used to wait for completion and retrieve results).

Workflow

1. Create: Call InferenceJob::create(task, listener).
2. Execute: Perform the inference logic in a separate thread/process. As tokens or text are generated, call job.emitter.completed(GeneratedItem::Text(text)) to stream them to the listener.
3. Finish: Call job_handle.join() to block until the job is complete and retrieve the final CompletedItem produced by your listener.

// Example workflow
let task = InferenceTask::Prompt(Prompt::new("Hello!"));
let listener = TextGenerationListener::default();
let (job, handle) = InferenceJob::create(task, listener);

// In your inference loop:
job.emitter.completed(GeneratedItem::Text("Hello".to_string()));
job.emitter.completed(GeneratedItem::Text(" world!".to_string()));

// Wait for result:
let final_text = handle.join();
println!("Result: {}", final_text);