ARLoop-baseline-0.6B

This is the fixed-loop PreNorm baseline for ARLoop. This release contains the final checkpoint of the approximately 180B-token run prenorm640m-l3-hybrid500b23-180b-lr3e3-gbs256-mbs4-s4096. It is a base pretrained model for text completion, not a chat/SFT model.

Training code: ykyaol7/ARLoop. The loading implementation and file layout follow ARLoop-1.3B.

Load and generate

The included implementation uses PyTorch SDPA and standard pre-norm residual additions. No ARLoop source checkout, FLA, Triton or flash-attn installation is required.

pip install 'transformers==5.22.0' safetensors
# Install a PyTorch build suitable for your CPU or CUDA environment separately.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "osunlp/ARLoop-baseline-0.6B"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.float32,
).to(device).eval()
inputs = tokenizer("The capital of France is", return_tensors="pt").to(device)
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Keep the weights in FP32: the implementation preserves FP32 residuals and value embeddings with BF16 autocast internally, matching the training model's mixed-precision policy. Total prompt plus generation length must not exceed 4096. For batches, the tokenizer uses left padding; pass its attention_mask to the model.

Architecture

Property	Value
Physical transformer layers	18, shared across 3 recurrent passes
Effective transformer depth	54
Hidden / MLP intermediate width	1536 / 4608
Attention / KV heads	16 / 4
Context length	4096
RoPE base	50000
Residual connections	Standard PreNorm additions; no attention-residual routing
Value embeddings	Shared across alternating physical layers
Tokenizer vocabulary	49152, SmolLM2 tokenizer used in training
Output head	Tied to token embeddings

The inference checkpoint contains 582,804,096 unique parameters. The training model also allocated an unused 75,497,472-parameter output head despite using tied embeddings; that unused weight is omitted here. Duplicate shared value embedding tensors are stored once. The exported weights retain their original FP32 values without lossy conversion.

Training checkpoint

main contains the final checkpoint:

Property	Value
Original filename	step-01373280-prenorm640m-l3-hybrid500b23-180b-lr3e3-gbs256-mbs4-s4096
Saved microbatch step	1,373,280
Saved optimizer step	171,660
Global batch size	256 sequences
Sequence length	4096 tokens
Training tokens	179,998,556,160

The token count is calculated from the saved optimizer step × 256 × 4096. Step labels in original filenames are microbatch steps, not optimizer steps. For reproducible experiments, pin the model repository's commit SHA.

Files and validation

• configuration_baseline.py, modeling_baseline.py: standalone Transformers loading code matching the native fixed-loop PreNorm architecture.
• config.json, generation_config.json: architecture and generation defaults.
• Two model-*.safetensors shards and model.safetensors.index.json: FP32 weights.
• Tokenizer JSON, vocabulary, merges and special-token metadata.

Export checks cover every tensor's shape, dtype and finiteness, equality of shared embeddings, and tokenizer vocabulary/merge consistency with the training run. Standalone loading, forward/loss and batched cached generation were tested with Python 3.12, PyTorch 2.6.0 (CPU) and Transformers 5.22.0, with imports of the training packages and FLA blocked. On the tested CPU input, logits exactly matched the native PreNorm model, and cached chunk inference exactly matched full forward inference. These are functional checks, not benchmark evaluations. GPU inference has not been tested for this release; different attention kernels may differ numerically in BF16.

License and attribution

Model release: Apache License 2.0. Attribution and MIT copyright/permission notices for adapted Parcae/ARLoop code are retained in NOTICE. The tokenizer files are copied from the training run's SmolLM2 tokenizer assets; these model weights were trained independently of SmolLM2.