levanter-backpack-1b

Model Card for Levanter-Backpack-1.4B

This is 1.4B parameter version of Backpack architecture, intended to combine strong modeling performance with an interface for interpretability and control.

Training Details

Training Data

This model was trained on the OpenWebText corpus.

Training Procedure

This model was trained for 450k gradient steps and cosine decaying learning rate from 1e-4 to zero, with a linear warmup of 5k steps.

Environmental Impact

• Hardware Type: v3-128 TPU (128 cores, 2TB Memory)
• Hours used: Roughly 8.6 days.
• Cloud Provider: Google Cloud Patform
• Compute Region: North America.

Model Architecture and Objective

This model was trained to minimize the cross-entropy loss, and is a Backpack language model.

Software

This model was trained with Levanter and Jax.

Loss Curve

How to Get Started with the Model

Please install transformers, safetensors and torch to use this model.

pip install transformers safetensors torch

Run the following Python code:

import torch
import transformers
from transformers import AutoModelForCausalLM


model_id = "stanford-crfm/levanter-backpack-1b"
config = transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=True)
torch_model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    config=config, 
    trust_remote_code=True
)
torch_model.eval()

input = torch.randint(0, 50264, (1, 512), dtype=torch.long)
torch_out = torch_model(input, position_ids=None,)
torch_out = torch.nn.functional.softmax(torch_out.logits, dim=-1)
print(torch_out.shape)