| 292 | # ────────────────────────────────────────────────────────────────────────────── |
| 293 | |
| 294 | def parse_args() -> argparse.Namespace: |
| 295 | p = argparse.ArgumentParser(description="Train a BitNet b1.58 LLaMA model.") |
| 296 | p.add_argument("--steps", type=int, default=CONFIG["max_steps"]) |
| 297 | p.add_argument("--batch-size", type=int, default=CONFIG["batch_size"]) |
| 298 | p.add_argument("--lr", type=float, default=CONFIG["learning_rate"]) |
| 299 | p.add_argument("--checkpoint-dir", type=str, default=CONFIG["checkpoint_dir"]) |
| 300 | p.add_argument("--wandb", action="store_true") |
| 301 | return p.parse_args() |
| 302 | |
| 303 | |
| 304 | if __name__ == "__main__": |