AXIOM-TECH/Ouro-2.6B-Thinking-IBNN

Copy

import torch, time, gc, sys
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

assert torch.cuda.is_available()
print(f"GPU  : {torch.cuda.get_device_name(0)}")
print(f"VRAM : {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

tok = AutoTokenizer.from_pretrained("AXIOM-TECH/Ouro-2.6B-Thinking-IBNN", trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(
    "AXIOM-TECH/Ouro-2.6B-Thinking-IBNN",
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True,
    attn_implementation="sdpa",
)

# Configure custom architecture parameters
model.config.total_ut_steps = 4
model.config.early_exit_threshold = 1.0
model.config.use_cache = True
model.generation_config = GenerationConfig.from_model_config(model.config)
model.generation_config.use_cache = True
model.eval()
model = torch.compile(model, mode="max-autotune")

print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f} GB")

print("Model ready")

Citation

@article{zhu2025scaling,
  title={Scaling Latent Reasoning via Looped Language Models},
  author={Zhu, Rui-Jie and Wang, Zixuan and Hua, Kai and Zhang, Tianyu and Li, Ziniu and Que, Haoran and Wei, Boyi and Wen, Zixin and Yin, Fan and Xing, He and others},
  journal={arXiv preprint arXiv:2510.25741},
  year={2025}
}


## License

This model is licensed under Apache-2.0. See the LICENSE file for details.

## Project Links

- **Paper**: [Updating the standard neuron model in artificial neural networks](https://arxiv.org/pdf/2605.30370)
- **Paper**: [Scaling Latent Reasoning via Looped Language Models](https://huggingface.co/papers/2510.25741)
- **Project Page**: [https://ouro-llm.github.io](https://ouro-llm.github.io)
- **Code**: [https://github.com/ByteDance/Ouro](https://github.com/ByteDance/Ouro)

---