ARLoop-0.6B

ARLoop is a recurrent causal language model with attention residuals. This release contains the final checkpoint of the three-stage, approximately 180B-token run lsopt640m-l3-b18-three-stage-blend180b-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 FP32 attention-residual routing. 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-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 uses BF16 residuals and 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
Attention-residual block size 18 sublayers
Recent residual blocks 3, with input embedding retained
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 583,138,944 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-lsopt640m-l3-b18-three-stage-blend180b-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_arloop.py, modeling_arloop.py: Transformers custom model code, reused from ARLoop-1.3B revision c04485151cb408a2fbc35e05a240fde3d6263fe0.
  • 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. GPU inference and benchmark scores have not been revalidated for this release. The portable PyTorch implementation follows the native fused model's computation but need not be bit-identical in BF16.

License and attribution

Model release: Apache License 2.0. Attribution and MIT copyright/permission notices for adapted Parcae/ARLoop code and the FLA routing reference 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.

Downloads last month
10
Safetensors
Model size
0.6B params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support