Views
No views yet
1import torch, time, gc, sys
2from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
3
4assert torch.cuda.is_available()
5print(f"GPU : {torch.cuda.get_device_name(0)}")
6print(f"VRAM : {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
7
8torch.backends.cuda.matmul.allow_tf32 = True
9torch.backends.cudnn.allow_tf32 = True
10
11tok = AutoTokenizer.from_pretrained("AXIOM-TECH/Ouro-2.6B-Thinking-IBNN", trust_remote_code=True)
12
13model = AutoModelForCausalLM.from_pretrained(
14 "AXIOM-TECH/Ouro-2.6B-Thinking-IBNN",
15 device_map="auto",
16 torch_dtype=torch.float16,
17 trust_remote_code=True,
18 attn_implementation="sdpa",
19)
20
21# Configure custom architecture parameters
22model.config.total_ut_steps = 4
23model.config.early_exit_threshold = 1.0
24model.config.use_cache = True
25model.generation_config = GenerationConfig.from_model_config(model.config)
26model.generation_config.use_cache = True
27model.eval()
28model = torch.compile(model, mode="max-autotune")
29
30print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f} GB")
31
32print("Model ready")
1@article{zhu2025scaling,
2 title={Scaling Latent Reasoning via Looped Language Models},
3 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},
4 journal={arXiv preprint arXiv:2510.25741},
5 year={2025}
6}
7
8
9## License
10
11This model is licensed under Apache-2.0. See the LICENSE file for details.
12
13## Project Links
14
15- **Paper**: [Updating the standard neuron model in artificial neural networks](https://arxiv.org/pdf/2605.30370)
16- **Paper**: [Scaling Latent Reasoning via Looped Language Models](https://huggingface.co/papers/2510.25741)
17- **Project Page**: [https://ouro-llm.github.io](https://ouro-llm.github.io)
18- **Code**: [https://github.com/ByteDance/Ouro](https://github.com/ByteDance/Ouro)
19
20---