Default to BF16 weights for inference

README.md

@@ -33,7 +33,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_id = "osunlp/ARLoop-1.3B"  # Or the local ARLOOP_hf directory.
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    model_id, trust_remote_code=True, dtype=torch.float32,
+    model_id, trust_remote_code=True, dtype=torch.bfloat16,
 ).to("cuda").eval()
 inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
 with torch.inference_mode():
@@ -41,6 +41,8 @@ with torch.inference_mode():
 print(tokenizer.decode(output[0], skip_special_tokens=True))
 ```
 
+Weights load in BF16 for inference. The stored checkpoint retains the original FP32 weights; routing and normalization statistics use FP32 internally.
+
 ## Architecture
 
 | Property | Value |

config.json

@@ -8,7 +8,7 @@
     "AutoModelForCausalLM": "modeling_arloop.ARLoopForCausalLM"
   },
   "bos_token_id": 0,
-  "dtype": "float32",
+  "dtype": "bfloat16",
   "embedding_scale": 1.0,
   "eos_token_id": 0,
   "head_dim": 128,
@@ -31,5 +31,6 @@
   "tie_word_embeddings": true,
   "transformers_version": "5.12.1",
   "use_cache": true,
-  "vocab_size": 49152
+  "vocab_size": 49152,
+  "torch_dtype": "bfloat16"
 }

modeling_arloop.py

@@ -182,7 +182,7 @@ class ARLoopForCausalLM(PreTrainedModel, GenerationMixin):
             if not bool(attention_mask.all()):
                 mask = (key_indices[None, :] <= query_indices[:, None])[None, None, :, :]
                 mask = mask & attention_mask[:, None, None, :].bool()
-        # Match the training model's BF16 residual stream with FP32 master weights.
+        # Use BF16 activations; parameters follow the requested loading dtype.
         with torch.autocast(device_type=input_ids.device.type, dtype=torch.bfloat16):
             embedded = (self.transformer.wte(input_ids) * self.config.embedding_scale).to(torch.bfloat16)
             ve = self.shared_value_embedding(input_ids)