trillionlabs/Tri-21B-Think-Preview

<p align="center"> <picture> <img src="https://raw.githubusercontent.com/trillion-labs/.github/main/Tri-21B-Think.png" alt="Tri-21B-Think-Preview" style="width: 80%;"> </picture> </p>

Introduction

Tri-21B-Think-Preview is an intermediate checkpoint of Tri-21B-Think, featuring mid-training context length expansion to 32K tokens and instruction tuning for chain-of-thought reasoning and tool use.

Model Specifications

• —Type: Causal Language Model (Reasoning-Enhanced)
• —Base Model: Tri-21B
• —Architecture: Transformer Decoder with RoPE, SwiGLU, RMSNorm, and GQA
• —Number of Parameters: 20.73B
• —Number of Layers: 40
• —Number of Attention Heads: 32 (Query) / 8 (Key, Value)
• —Head Dimension: 160
• —Hidden Size: 5,120
• —Intermediate Size: 27,392
• —Context Length: 32,768 (up to 262,144 with YaRN)
• —Vocab Size: 124,416

Quickstart

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "trillionlabs/Tri-21B-Think-Preview"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "Solve the following step by step: What is the sum of the first 100 prime numbers?"
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=4096,
    temperature=0.6,
    top_p=0.9
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)

vLLM & SGLang Deployment

vLLM and SGLang support for Trillion Model is on the way. Stay tuned!

Fine-tuning Notes

Note on `<think>` tags: This model was trained without <think> and </think> as special tokens. They were added post-training for compatibility with reasoning parsers. If you plan to fine-tune this model, you'll need to modify tokenizer_config.json to avoid indexing errors.

Replace tokens 123975 and 123976 in tokenizer_config.json:

"123975": {
  "content": "<|reserved_special_token_9|>",
  "lstrip": false,
  "normalized": false,
  "rstrip": false,
  "single_word": false,
  "special": true
},
"123976": {
  "content": "<|reserved_special_token_10|>",
  "lstrip": false,
  "normalized": false,
  "rstrip": false,
  "single_word": false,
  "special": true
}

Evaluation

Limitations

• —Language Support: Optimized for English, Korean, and Japanese. Other languages may show degraded performance.
• —Knowledge Cutoff: February 2025.
• —Intermediate Checkpoint: See Tri-21B-Think for the final model.

License

This model is licensed under the Apache 2.0 License.

Contact

For inquiries: info@trillionlabs.co