Make BF16 the canonical GPU weight format

README.md

@@ -71,9 +71,13 @@ These scores show clear progress over the project’s earlier checkpoints, but t
 
 This repository is a native safetensors export rather than a Transformers `AutoModel` package. It includes the minimal `tiny_english` inference module.
 
+## Weights
+
+`model.safetensors` contains the single BF16 release (624 MB; SHA-256 `fce44c1f7fc64a51a775d8802cfa75277dc68a34ef5d2908f541caa21fad66c2`). It requires a BF16-capable CUDA GPU, such as the RTX 3060 Ti used for training. The bundled inference command defaults to CUDA.
+
 ```bash
 pip install -r requirements.txt
-python -m tiny_english.generate --model-dir . --prompt "The invention of the printing press changed" --max-new-tokens 128 --temperature 0.7 --top-k 40 --seed 42
+python -m tiny_english.generate --model-dir . --device cuda --prompt "The invention of the printing press changed" --max-new-tokens 128 --temperature 0.7 --top-k 40 --seed 42
 ```
 
 Generation can repeat phrases, make factual errors, or produce incoherent continuations. Do not use it for high-stakes decisions or as an assistant without additional evaluation and alignment work.

config.json

@@ -11,14 +11,14 @@
     "heads": 16,
     "intermediate_size": 2752
   },
-  "dtype": "float32",
+  "dtype": "bfloat16",
   "parameter_count": 311866368,
   "checkpoint_step": 15259,
   "stage_tokens": 2000027648,
   "previous_tokens": 3000238080,
   "training_tokens": 5000265728,
   "tokenizer_sha256": "40d36acc35e5364d90fdcdae773143b7571b5ec3b7ace2030d12d4bd5a3a1025",
-  "weights_sha256": "1255d10b05add1834ebe0ee05cd65f0173b57448b26b2b0ff203c29a761e9a50",
+  "weights_sha256": "fce44c1f7fc64a51a775d8802cfa75277dc68a34ef5d2908f541caa21fad66c2",
   "data_manifest_sha256": "cf6ca78377e03aef753b073d25f710d665de31f0a91d285032955feecc565934",
   "eos_token_id": 1,
   "unk_token_id": 0

model.bf16.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:fce44c1f7fc64a51a775d8802cfa75277dc68a34ef5d2908f541caa21fad66c2
-size 623748472

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1255d10b05add1834ebe0ee05cd65f0173b57448b26b2b0ff203c29a761e9a50
-size 1247481064
+oid sha256:fce44c1f7fc64a51a775d8802cfa75277dc68a34ef5d2908f541caa21fad66c2
+size 623748472

tiny_english/generate.py

@@ -25,12 +25,16 @@ def load_model(checkpoint_path, tokenizer_path, device="cpu"):
     return model.to(device).eval(), tokenizer, step
 
 
-def load_export(directory, device="cpu"):
+def load_export(directory, device="cpu"):
     from safetensors.torch import load_file
     directory = Path(directory)
     metadata = json.loads((directory / "config.json").read_text(encoding="utf-8"))
-    if metadata.get("format") != "tiny-english-safetensors-v1":
-        raise ValueError("Unsupported exported model format")
+    if metadata.get("format") != "tiny-english-safetensors-v1":
+        raise ValueError("Unsupported exported model format")
+    if metadata.get("dtype") != "bfloat16":
+        raise ValueError("This export requires BF16 weights")
+    if device != "cuda":
+        raise ValueError("This BF16 export requires a CUDA GPU; use --device cuda")
     tokenizer_path = directory / "tokenizer.json"
     weights_path = directory / "model.safetensors"
     if digest_file(tokenizer_path) != metadata["tokenizer_sha256"]:
@@ -42,7 +46,7 @@ def load_export(directory, device="cpu"):
     if tokenizer.get_vocab_size() != model.config.vocab_size:
         raise ValueError("Export tokenizer vocabulary does not match model")
     model.load_state_dict(load_file(str(weights_path), device="cpu"), strict=True)
-    return model.to(device).eval(), tokenizer, metadata["checkpoint_step"]
+    return model.to(device=device, dtype=torch.bfloat16).eval(), tokenizer, metadata["checkpoint_step"]
 
 
 def complete(model, tokenizer, prompt, max_new_tokens=150, temperature=0.8, top_k=40, seed=42, deadline=None, strict_context=False):
@@ -90,7 +94,7 @@ def main():
     parser.add_argument("--max-new-tokens", type=int, default=150)
     parser.add_argument("--temperature", type=float, default=0.8)
     parser.add_argument("--top-k", type=int, default=40)
-    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda")
     parser.add_argument("--threads", type=int, default=8)
     parser.add_argument("--seed", type=int, default=42)
     args = parser.parse_args()

training_summary.json

@@ -30,6 +30,6 @@
   "previous_tokens": 3000238080,
   "training_tokens": 5000265728,
   "sampled_validation_loss": 2.609314309898764,
-  "export_dtype": "float32",
+  "export_dtype": "bfloat16",
   "optimizer_state_included": false
 }