Upload folder using huggingface_hub

README.md

@@ -27,6 +27,37 @@ tags:
 - **Iterative Latent Reasoning**: Performs reasoning through recurrent computation in latent space
 - **Adaptive Computation**: Supports early exit mechanisms for dynamic compute allocation
 
+## Configuration
+
+### Recurrent Steps and Adaptive Exit
+
+The model's computational behavior can be configured through the `config.json` file:
+
+```json
+{
+  "total_ut_steps": 4,
+  "early_exit_threshold": 1.0
+}
+```
+
+- **`total_ut_steps`**: Controls the number of recurrent steps (default: 4). You can adjust this value to trade off between performance and computation time.
+- **`early_exit_threshold`**: Controls the adaptive exit mechanism (default: 1.0). Lower values encourage earlier exit, while 1.0 means always use all steps.
+
+**Example: Modify recurrent steps**
+```python
+from transformers import AutoConfig, AutoModelForCausalLM
+
+config = AutoConfig.from_pretrained("ByteDance/Ouro-1.4B")
+config.total_ut_steps = 3  # Use 3 recurrent steps instead of 4
+model = AutoModelForCausalLM.from_pretrained(
+    "ByteDance/Ouro-1.4B",
+    config=config,
+    device_map="auto"
+)
+```
+
+> **Note**: vLLM does not currently support the adaptive exit feature due to its inference optimization characteristics. When using vLLM, the model will always execute the full number of `total_ut_steps`.
+
 ## Model Architecture
 
 Ouro-1.4B is based on the decoder-only Transformer architecture with parameter sharing across recurrent steps:

config.json

@@ -54,6 +54,7 @@
   "tie_word_embeddings": false,
   "torch_dtype": "bfloat16",
   "total_ut_steps": 4,
+  "early_exit_threshold": 1.0,
   "transformers_version": "4.55.0",
   "use_cache": true,
   "use_sliding_window": false,

configuration_ouro.py

@@ -169,6 +169,8 @@ class OuroConfig(PretrainedConfig):
         max_window_layers=28,
         layer_types=None,
         attention_dropout=0.0,
+        total_ut_steps=4,
+        early_exit_threshold=1.0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -193,6 +195,8 @@ class OuroConfig(PretrainedConfig):
         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
         self.attention_dropout = attention_dropout
+        self.total_ut_steps = total_ut_steps
+        self.early_exit_threshold = early_exit_threshold
         # Validate the correctness of rotary position embeddings parameters
         # BC: if there is a 'type' field, move it to 'rope_type'.
         if self.rope_scaling is not None and "type" in self.rope_scaling: