SpiceeChat/Train-From-Scratch-Via-Flash3-Kernels

train.py =>

parser.add_argument("--attention-backend", default="sage",
choices=["sage", "torch", "flash2", "flash3"], )

parser.add_argument("--src", required=True)
parser.add_argument("--tokenizer", required=True)
parser.add_argument("--out-dir", default="LM_SAGE9")
parser.add_argument("--cache-dir", default="lm_token_cache_sage9_marker_special")

parser.add_argument("--activity-column", default="activites")
parser.add_argument("--row-batch-size", type=int, default=100_000)
parser.add_argument("--rebuild-cache", action="store_true")

parser.add_argument("--shuffle-before-tokenize", action="store_true")
parser.add_argument("--no-shuffle-before-tokenize", action="store_true")
parser.add_argument("--shuffle-buffer-size", type=int, default=500_000)

parser.add_argument("--ctx-len", type=int, default=512)
parser.add_argument("--batch-size", type=int, default=4)
parser.add_argument("--num-workers", type=int, default=0)

parser.add_argument("--shuffle-blocks", action="store_true")
parser.add_argument("--max-train-tokens", type=int, default=0)
parser.add_argument("--max-val-tokens", type=int, default=0)

parser.add_argument("--val-ratio", type=float, default=0.01)
parser.add_argument("--val-every", type=int, default=2000)
parser.add_argument("--val-batches", type=int, default=10)

parser.add_argument("--n-layer", type=int, default=4)
parser.add_argument("--n-head", type=int, default=4)
parser.add_argument("--n-embd", type=int, default=384)
parser.add_argument("--dropout", type=float, default=0.0)

parser.add_argument(
    "--attention-backend",
    default="sage",
    choices=["sage", "torch", "flash2", "flash3"],
)

parser.add_argument("--lr", type=float, default=3e-4)
parser.add_argument("--weight-decay", type=float, default=0.1)
parser.add_argument("--max-steps", type=int, default=50_000)
parser.add_argument("--save-every", type=int, default=10_000)
parser.add_argument("--log-every", type=int, default=20)
parser.add_argument("--grad-clip", type=float, default=1.0)

parser.add_argument("--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"])
parser.add_argument("--device", default="cuda")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--lowercase", action="store_true")

parser.add_argument(
    "--special-marker",
    action="append",
    default=[],
    help='Map a dataset marker to a tokenizer special token. Example: --special-marker "<|BOC|>=[BOC]". Can be repeated.',
)

parser.add_argument(
    "--no-special-boundaries",
    action="store_true",
    help="Disable BOS/EOS insertion and marker replacement during pretokenization.",
)

parser.add_argument(
    "--no-append-sep",
    action="store_true",
    help="Legacy alias: disables BOS/EOS insertion too.",
)

parser.add_argument("--compile", action="store_true")

return parser.parse_args()