🚀 Refined BitTransformerLM: Organized codebase with best practices

scripts/benchmarks/wikitext_benchmark.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from bit_transformer import text_to_bits, collapse_submodel
+from progressive_scaleup import progressive_scale_up_text
+
+
+def lines_to_bits(lines, max_len=64):
+    data = []
+    for text in lines:
+        bits = text_to_bits(text)[:max_len]
+        if len(bits) < max_len:
+            bits.extend([0] * (max_len - len(bits)))
+        data.append(bits)
+    return data
+
+
+def main():
+    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]")
+    val_ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation[:1%]")
+    train_lines = [item["text"] for item in ds][:256]
+    valid_lines = [item["text"] for item in val_ds][:64]
+
+    train_bits = lines_to_bits(train_lines)
+    valid_bits = lines_to_bits(valid_lines)
+
+    progressive_scale_up_text(
+        eps=0.65,
+        steps=4,
+        width_mult=2.0,
+        max_len=64,
+        dataset_size=min(64, len(train_bits)),
+    )
+
+    target_params = dict(d_model=16, nhead=4, num_layers=1, dim_feedforward=64, max_seq_len=64)
+    model, _ = collapse_submodel(train_bits[:64], target_params, max_rounds=1)
+
+    val_tensor = torch.tensor(valid_bits, dtype=torch.long)
+    logits, _ = model(val_tensor)
+    pred = logits[:, :-1, :].reshape(-1, 2)
+    target = val_tensor[:, 1:].reshape(-1)
+    loss = F.cross_entropy(pred, target)
+    print("Collapsed model validation loss:", loss.item())
+
+
+if __name__ == "__main__":
+    main()