Upload train_cascade_1b.py with huggingface_hub

train_cascade_1b.py

@@ -63,6 +63,7 @@ def arguments() -> argparse.Namespace:
     p.add_argument("--validation-batches", type=int, default=4)
     p.add_argument("--seed", type=int, default=3407)
     p.add_argument("--compile", action="store_true")
+    p.add_argument("--compile-blocks", action="store_true")
     p.add_argument("--resume", action="store_true")
     p.add_argument("--probe-only", action="store_true")
     p.add_argument("--probe-batches", default="64,96,128,160,192,224,256")
@@ -420,6 +421,8 @@ def train(args: argparse.Namespace) -> None:
     optimizers = Optimizers(model, args)
     initial_validation = evaluate(model, validation, config.route_paths)
     print(json.dumps({"parameters": counts, "threshold": threshold, "initial_validation": initial_validation, "sources": packer.documents}), flush=True)
+    if args.compile_blocks:
+        model.compile_blocks()
     compiled = torch.compile(model, mode="max-autotune-no-cudagraphs") if args.compile else model
     history = args.output / "training.jsonl"
     tokens = start_step * args.batch_size * args.sequence_length