Dhara-250M-AR-Base

A 250M parameter autoregressive language model with Canon layers and 32K context length, designed as the base model for the Dhara masked diffusion LLM family.

Table of Contents

• Model Description
• Architecture
• Training Data
• Training Details
• Benchmark Results
• Context Extension Results
• Usage
• YaRN Context Extension
• Key Insights
• Limitations
• Related Work
• Citation

Model Description

Dhara-250M-AR-Base is an autoregressive language model built on the LLaMA3 architecture enhanced with Canon layers — causal 1D depthwise convolutions placed at all 4 positions (ABCD) within each transformer layer. The model was pretrained on 10.2 billion tokens and extended to 32K context length via progressive RoPE scaling, with built-in YaRN support for inference-time extension to 64K-128K.

This model serves as the AR base for the Dhara diffusion LLM pipeline, where the pretrained weights are transferred to a masked diffusion model (MDM) for parallel text generation.

Key Features

• Canon Layers (ABCD): O(n) depthwise convolutions at 4 positions per layer for enhanced local context mixing with only 0.26% parameter overhead
• 32K Native Context: Extended from 4K via progressive RoPE scaling (3 stages)
• YaRN Support: Built-in inference-time context extension to 64K-128K without retraining
• GQA: Grouped Query Attention (12 heads, 4 KV heads) for efficient inference
• QK Normalization: RMSNorm on Q/K after RoPE for training stability
• Logit Softcapping: Prevents logit explosion during training

Architecture

Canon Layer Positions

Based on "Physics of Language Models: Part 4.1" by Zeyuan Allen-Zhu:

• Position A: After input LayerNorm, before attention
• Position B: Inside attention, applied separately to Q/K/V after linear projections
• Position C: After post-attention LayerNorm, before MLP
• Position D: Inside MLP, after gate*up product

Each Canon layer is a causal depthwise 1D convolution with residual connection, adding local sequential context mixing at O(n) cost.

Layer Flow

x → LayerNorm → [Canon-A] → Attention([Canon-B on Q,K,V]) → + residual
x → LayerNorm → [Canon-C] → MLP([Canon-D]) → + residual

Training Data

Stage 1: Pretraining (10.2B tokens)

• Dataset: codelion/sutra-10B
• Tokens: 10.2 billion
• Context Length: 4,096 tokens
• Content: Diverse English text across science, technology, mathematics, social studies, arts, and language arts
• Domain Distribution: interdisciplinary (35%), technology (21%), science (14%), social studies (8%), mathematics (8%), life skills (5%), arts (4%), language arts (2%)

Stage 2: Long Context Extension (900M tokens)

Progressive RoPE scaling across 3 phases using long-document data:

Training Details

Benchmark Results

Evaluated using lm-evaluation-harness v0.4.11 (0-shot unless noted):

Context extension to 32K preserved short-context benchmark performance with negligible change.

Context Extension Results

Perplexity measured on held-out PG19 (Project Gutenberg) long documents:

Before context extension, perplexity exploded beyond 4K (33 → 703). After extension, perplexity stays nearly flat across all lengths (36 → 42), confirming the model can effectively use the full 32K context.

Usage

Basic Usage

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("codelion/dhara-250m-ar-base", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "codelion/dhara-250m-ar-base",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Generate text
prompt = "The most important discovery in physics was"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(
    inputs.input_ids,
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2,
    do_sample=True,
    pad_token_id=tokenizer.pad_token_id,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Long Context Usage (32K)

# The model natively supports 32K context
long_text = "..." # up to 32K tokens
inputs = tokenizer(long_text, return_tensors="pt", truncation=True, max_length=32768).to(device)
outputs = model(inputs.input_ids)

YaRN Context Extension

The model includes built-in YaRN (Yet another RoPE extensioN) support for inference-time context extension beyond 32K without retraining:

from transformers import AutoModelForCausalLM

# Extend to 64K (2x)
model = AutoModelForCausalLM.from_pretrained(
    "codelion/dhara-250m-ar-base",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    rope_scaling={"type": "yarn", "factor": 2.0, "original_max_position_embeddings": 32768},
)

# Extend to 128K (4x)
model = AutoModelForCausalLM.from_pretrained(
    "codelion/dhara-250m-ar-base",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    rope_scaling={"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768},
)

YaRN splits RoPE frequency dimensions into three groups — high frequencies (local positions) are preserved, low frequencies (global positions) are interpolated — giving better quality than naive scaling.

Key Insights

1. Canon Layers Add Local Context Cheaply: All 4 Canon positions (ABCD) add depthwise causal convolutions throughout the model with only 0.26% parameter overhead, enhancing local sequential pattern recognition.

2. Progressive RoPE Extension Works: Three-stage context extension (4K → 8K → 16K → 32K) with increasing theta preserved short-context quality while enabling flat perplexity out to 32K.

3. AR as Foundation for MDM: This model is designed as a weight donor for masked diffusion model conversion. The embeddings, MLPs, and Canon layers transfer directly; only the attention patterns need retraining for bidirectional denoising.

4. Efficient Training: The full pipeline (10.2B pretraining + 900M context extension) completed in ~55 hours on a single GPU.

Limitations

• This is a base model without instruction tuning — it generates text continuations, not answers
• Performance is limited by the 250M parameter scale and 10.2B token training budget
• The tokenizer is English-focused; non-Latin scripts fall back to byte-level encoding
• Long-context quality depends on having relevant content filling the context window

Related Work

• Scaling Pedagogical Pretraining to 10 Billion Tokens — Prior work on scaling pretraining data
• The Optimal Architecture for Small Language Models — Architecture design principles used in this model
• Dhara-70M — Our 70M diffusion language model
• Physics of Language Models: Part 4.1 — Canon layers (Zeyuan Allen-Zhu)

Contact

For questions or feedback, please open a discussion on the Hugging Face discussions page.