import argparse
import gc
import logging
import math
import os
import time
from contextlib import nullcontext
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import sys

_script_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(_script_dir)
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as grad_ckpt

from model_cpu_gpt2 import (
    CPUGPT,
    CPUGPTConfig,
    get_config,
    gpt2_small_config,
    smoke_config,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

BYTE_PAD = 0
BYTE_BOS = 1
BYTE_EOS = 2
BYTE_MASK = 3
BYTE_SEP = 4
BYTE_OFFSET = 5
BYTE_VOCAB = 261

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
torch.backends.cuda.enable_flash_sdp(False)

polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


def adamw_step(p, grad, m, v, step, lr, b1, b2, eps, wd):
    p.mul_(1 - lr * wd)
    m.lerp_(grad, 1 - b1)
    v.lerp_(grad.square(), 1 - b2)
    bc1 = 1 - b1**step
    bc2 = 1 - b2**step
    p.addcdiv_(m / bc1, (v / bc2).sqrt_().add_(eps), value=-lr)


def muon_step(grads_stack, params, mom_buf, lr, momentum=0.95, ns_steps=3):
    mom_buf.lerp_(grads_stack, 1 - momentum)
    X = mom_buf.float()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
    for a, b, c in polar_express_coeffs[:ns_steps]:
        if X.size(-2) >= X.size(-1):
            A = X.mT @ X
            X = a * X + X @ (b * A + c * (A @ A))
        else:
            A = X @ X.mT
            X = a * X + (b * A + c * (A @ A)) @ X
    torch._foreach_sub_(params, list((X * lr).to(params[0].dtype).unbind(0)))


class MuonAdamW(torch.optim.Optimizer):
    def __init__(self, param_groups):
        super().__init__(param_groups, defaults={})

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            if g["kind"] == "adamw":
                for p in g["params"]:
                    if p.grad is None:
                        continue
                    st = self.state[p]
                    if not st:
                        st["step"] = 0
                        st["m"] = torch.zeros_like(p)
                        st["v"] = torch.zeros_like(p)
                    st["step"] += 1
                    adamw_step(
                        p,
                        p.grad,
                        st["m"],
                        st["v"],
                        st["step"],
                        g["lr"],
                        *g["betas"],
                        g["eps"],
                        g.get("wd", 0),
                    )
            elif g["kind"] == "muon":
                params = g["params"]
                if not params:
                    continue
                p0 = params[0]
                st = self.state[p0]
                stacked = torch.stack(
                    [p.grad for p in params if p.grad is not None]
                ).float()
                if not st:
                    st["mom"] = torch.zeros_like(stacked)
                lr = g["lr"] * max(1.0, p0.shape[-2] / max(p0.shape[-1], 1)) ** 0.5
                muon_step(stacked, params, st["mom"], lr, g.get("momentum", 0.95))


def build_optimizer(
    model: CPUGPT,
    cfg: CPUGPTConfig,
    lr_matrix=0.02,
    lr_emb=0.2,
    lr_lm=0.004,
    wd=0.0,
    betas=(0.8, 0.95),
) -> MuonAdamW:
    scale = (cfg.n_embd / 768) ** -0.5
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    matrix_params, scalar_params = [], []
    for block in raw.blocks:
        for p in block.parameters():
            (matrix_params if p.ndim == 2 else scalar_params).append(p)

    groups = [
        dict(
            kind="adamw",
            params=list(raw.wte.parameters()),
            lr=lr_emb * scale,
            betas=betas,
            eps=1e-8,
            wd=0,
        ),
        dict(
            kind="adamw",
            params=list(raw.lm_head.parameters())
            if raw.lm_head.weight is not raw.wte.weight
            else [],
            lr=lr_lm * scale,
            betas=betas,
            eps=1e-8,
            wd=0,
        ),
        dict(
            kind="adamw",
            params=scalar_params,
            lr=lr_matrix * scale,
            betas=betas,
            eps=1e-8,
            wd=0,
        ),
    ]
    for shape in sorted({p.shape for p in matrix_params}):
        ps = [p for p in matrix_params if p.shape == shape]
        groups.append(dict(kind="muon", params=ps, lr=lr_matrix, momentum=0.95))

    opt = MuonAdamW(groups)
    for g in opt.param_groups:
        g["initial_lr"] = g["lr"]
    return opt


def lr_multiplier(
    progress: float, warmup: float = 0.01, min_ratio: float = 0.1
) -> float:
    if progress < warmup:
        return progress / warmup
    t = (progress - warmup) / (1.0 - warmup)
    return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * t))


def _load_shard(path: str, seq_len: int):
    import pyarrow.parquet as pq

    tbl = pq.read_table(path, columns=["tokens"])
    tokens = torch.tensor(tbl["tokens"].to_pylist()[0], dtype=torch.long)
    n = (len(tokens) // seq_len) * seq_len
    return tokens[:n].view(-1, seq_len)


def _load_shard_bin(path: str, seq_len: int):
    import numpy as np

    data = np.memmap(path, dtype=np.int32, mode="r")
    n = (len(data) // seq_len) * seq_len
    return torch.from_numpy(data[:n].copy()).long().view(-1, seq_len)


def _load_shard_text(path: str, seq_len: int):
    import tiktoken

    enc = tiktoken.get_encoding("r50k_base")
    with open(path, "r", encoding="utf-8", errors="replace") as f:
        text = f.read()
    tokens = torch.tensor(enc.encode(text), dtype=torch.long)
    n = (len(tokens) // seq_len) * seq_len
    if n == 0:
        return torch.zeros(0, seq_len, dtype=torch.long)
    return tokens[:n].view(-1, seq_len)


def _load_wikitext_bytes(max_bytes: int = 2_000_000) -> torch.Tensor:
    from datasets import load_dataset

    ds = load_dataset(
        "Salesforce/wikitext",
        "wikitext-103-raw-v1",
        split="test",
        trust_remote_code=True,
    )
    buf: list[int] = []
    for row in ds:
        text = row["text"]
        if not text.strip():
            continue
        raw = text.encode("utf-8", errors="replace")
        buf.append(BYTE_BOS)
        buf.extend(b + BYTE_OFFSET for b in raw)
        buf.append(BYTE_EOS)
        if len(buf) >= max_bytes:
            break
    return torch.tensor(buf[:max_bytes], dtype=torch.long)


def _owt_bytes_producer(
    seq_len: int,
    q: "queue.Queue[tuple[torch.Tensor, torch.Tensor, int]]",
    batch_size: int,
    seed: int,
) -> None:
    from datasets import load_dataset

    ep = 1
    buf: list[int] = []
    batch_x: list[torch.Tensor] = []
    batch_y: list[torch.Tensor] = []
    while True:
        ds = load_dataset(
            "Skylion007/openwebtext",
            split="train",
            streaming=True,
            trust_remote_code=True,
        )
        ds = ds.shuffle(seed=seed + ep, buffer_size=10_000)
        for doc in ds:
            raw = doc["text"].encode("utf-8", errors="replace")
            buf.append(BYTE_BOS)
            buf.extend(b + BYTE_OFFSET for b in raw)
            buf.append(BYTE_EOS)
            while len(buf) >= seq_len + 1:
                chunk = buf[: seq_len + 1]
                buf = buf[seq_len + 1 :]
                batch_x.append(torch.tensor(chunk[:-1], dtype=torch.long))
                batch_y.append(torch.tensor(chunk[1:], dtype=torch.long))
                if len(batch_x) == batch_size:
                    q.put((torch.stack(batch_x), torch.stack(batch_y), ep))
                    batch_x, batch_y = [], []
        ep += 1


def make_loader(
    data_dir: str,
    seq_len: int,
    device: str,
    batch_size: int = 1,
    data_format: str = "parquet",
):
    import glob
    import queue
    import threading

    buf: queue.Queue = queue.Queue(maxsize=4)

    if data_format == "bytes":
        t = threading.Thread(
            target=_owt_bytes_producer,
            args=(seq_len, buf, batch_size, 42),
            daemon=True,
        )
        t.start()
        while True:
            x, y, ep = buf.get()
            yield x.to(device), y.to(device), ep
        return

    if data_format == "bin":
        shards = sorted(glob.glob(os.path.join(data_dir, "*.bin")))
        load_fn = _load_shard_bin
    elif data_format == "text":
        shards = sorted(glob.glob(os.path.join(data_dir, "*.txt")))
        load_fn = _load_shard_text
    else:
        shards = sorted(glob.glob(os.path.join(data_dir, "*.parquet")))
        load_fn = _load_shard

    if not shards:
        raise FileNotFoundError(f"No *.{data_format} files found in {data_dir}")

    def _producer():
        ep = 1
        while True:
            for shard in shards:
                seqs = load_fn(shard, seq_len + 1)
                if len(seqs) == 0:
                    continue
                idx = torch.randperm(len(seqs))
                batch_x, batch_y = [], []
                for i in range(len(idx)):
                    row = seqs[idx[i]]
                    batch_x.append(row[:-1])
                    batch_y.append(row[1:])
                    if len(batch_x) == batch_size:
                        buf.put((torch.stack(batch_x), torch.stack(batch_y), ep))
                        batch_x, batch_y = [], []
            ep += 1

    t = threading.Thread(target=_producer, daemon=True)
    t.start()

    while True:
        x, y, epoch = buf.get()
        yield x.to(device), y.to(device), epoch


def _make_ckpt_forward(original_forward):
    def _ckpt_fwd(x):
        return grad_ckpt(original_forward, x, use_reentrant=False)

    return _ckpt_fwd


def save_ckpt(model, opt, step, path, keep=2):
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    torch.save({"step": step, "model": raw.state_dict(), "opt": opt.state_dict()}, path)
    log.info(f"checkpoint saved → {path}")
    old = sorted(Path(path).parent.glob("step_[0-9]*.pt"))[:-keep]
    for p in old:
        p.unlink(missing_ok=True)


def load_ckpt(model, opt, path):
    ck = torch.load(path, map_location="cpu")
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.load_state_dict(ck["model"])
    opt.load_state_dict(ck["opt"])
    return ck["step"]


def _eval_val(model, val_tokens: torch.Tensor, seq_len: int, device: str) -> float:
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.eval()
    T = seq_len
    total_loss = 0.0
    total_cnt = 0
    num_win = (len(val_tokens) - 1) // T
    with torch.no_grad():
        for w in range(num_win):
            x = val_tokens[w * T : (w + 1) * T].unsqueeze(0).to(device)
            y = val_tokens[w * T + 1 : (w + 1) * T + 1].unsqueeze(0).to(device)
            if y.shape[1] < T:
                break
            loss = raw(x, y)
            total_loss += loss.item() * T
            total_cnt += T
    raw.train()
    return total_loss / max(total_cnt, 1)


def gpu_mem_mb(device) -> float:
    try:
        return torch.cuda.max_memory_allocated(device) / (1024 * 1024)
    except Exception:
        return -1.0


def train(args):
    use_ddp = args.num_gpus > 1
    local_rank = 0
    global_rank = 0

    if use_ddp:
        dist.init_process_group(backend="nccl")
        local_rank = dist.get_rank() % args.num_gpus
        global_rank = dist.get_rank()
        torch.cuda.set_device(local_rank)

    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
    is_master = global_rank == 0

    torch.manual_seed(args.seed)

    try:
        cfg = get_config(args.config)
    except ValueError:
        cfg = gpt2_small_config()

    if args.seq_len:
        cfg.seq_len = args.seq_len
    if args.n_embd:
        cfg.n_embd = args.n_embd
    if args.n_layer:
        cfg.n_layer = args.n_layer
    if args.gla_chunk:
        cfg.gla_chunk = args.gla_chunk

    model = CPUGPT(cfg).to(device)
    nparams = model.param_count()
    if is_master:
        log.info(f"model: {nparams / 1e6:.1f}M params  config={cfg}")

    if args.grad_checkpoint:
        for block in model.blocks:
            block.forward = _make_ckpt_forward(block.forward)
        if is_master:
            log.info("gradient checkpointing enabled")

    if args.compile:
        import torch._dynamo as _dynamo

        _dynamo.config.suppress_errors = True
        model = torch.compile(model, backend="inductor", fullgraph=False)
        if is_master:
            log.info("torch.compile active (inductor)")

    if use_ddp:
        model = nn.parallel.DistributedDataParallel(
            model, device_ids=[local_rank], output_device=local_rank
        )

    opt = build_optimizer(
        model, cfg, lr_matrix=args.matrix_lr, lr_emb=args.emb_lr, wd=args.weight_decay
    )

    total_tokens = int(float(args.tokens))
    tokens_per_step = args.total_batch
    seq_len = cfg.seq_len
    dev_batch = args.device_batch
    world_size = dist.get_world_size() if use_ddp else 1
    tokens_per_micro = dev_batch * seq_len * world_size
    grad_accum = max(1, tokens_per_step // max(tokens_per_micro, 1))
    total_steps = max(1, total_tokens // tokens_per_step)

    if is_master:
        log.info(
            f"total_tokens={total_tokens / 1e9:.2f}B  steps={total_steps}  "
            f"grad_accum={grad_accum}  device_batch={dev_batch}  world_size={world_size}"
        )

    if args.precision == "bf16" and torch.cuda.is_available():
        amp_ctx = torch.autocast("cuda", dtype=torch.bfloat16)
    else:
        amp_ctx = nullcontext()

    loader = make_loader(
        args.data_dir,
        seq_len,
        device,
        batch_size=dev_batch,
        data_format=args.data_format,
    )

    ckpt_dir = Path(args.ckpt_dir)
    if is_master:
        ckpt_dir.mkdir(parents=True, exist_ok=True)
    if use_ddp:
        dist.barrier()

    step = 0
    resume = sorted(ckpt_dir.glob("step_*.pt"))
    if resume:
        saved_initial_lrs = [g["initial_lr"] for g in opt.param_groups]
        step = load_ckpt(model, opt, resume[-1])
        for g, ilr in zip(opt.param_groups, saved_initial_lrs):
            g["initial_lr"] = ilr
        if is_master:
            log.info(f"resumed from step {step}")

    if step == 0 and is_master:
        init_pt = ckpt_dir / "init.pt"
        raw = (
            model.module
            if isinstance(model, nn.parallel.DistributedDataParallel)
            else model
        )
        torch.save(raw.state_dict(), init_pt)
        log.info(f"init.pt saved → {init_pt}")
        if args.btm and args.s3_bucket:
            import io as _io

            import boto3 as _boto3

            _s3i = _boto3.client("s3")
            _buf = _io.BytesIO()
            torch.save(raw.state_dict(), _buf)
            _buf.seek(0)
            _init_key = f"{args.run_name}/init.pt"
            _s3i.upload_fileobj(_buf, args.s3_bucket, _init_key)
            log.info(f"init.pt uploaded → s3://{args.s3_bucket}/{_init_key}")

    val_tokens = None
    if args.val_shard and os.path.exists(args.val_shard):
        import pyarrow.parquet as _pq

        _tbl = _pq.read_table(args.val_shard, columns=["tokens"])
        val_tokens = torch.tensor(_tbl["tokens"].to_pylist()[0], dtype=torch.long)
        if is_master:
            log.info(f"val shard: {len(val_tokens):,} tokens from {args.val_shard}")
    elif args.data_format == "bytes":
        if is_master:
            log.info("loading WikiText-103 test split as byte val ...")
        val_tokens = _load_wikitext_bytes(max_bytes=2_000_000)
        if is_master:
            log.info(
                f"byte val: {len(val_tokens):,} byte tokens from WikiText-103 test"
            )

    diloco = args.num_nodes > 1 and args.master_ip and not args.btm
    s3 = None
    velocity = None
    ref_state = None
    inner_since_sync = 0
    outer_step = 0
    if diloco:
        from diloco_sync import (
            diloco_outer_step,
            init_gloo,
            load_latest_checkpoint,
            save_outer_checkpoint,
        )

        init_gloo(args.master_ip, args.gloo_port, args.node_rank, args.num_nodes)
        if args.s3_bucket:
            import boto3

            s3 = boto3.client("s3")
            ckpt = load_latest_checkpoint(
                s3, args.s3_bucket, args.run_name, args.node_rank
            )
            if ckpt is not None:
                raw = (
                    model.module
                    if isinstance(model, nn.parallel.DistributedDataParallel)
                    else model
                )
                raw.load_state_dict(ckpt["model"])
                opt.load_state_dict(ckpt["optimizer"])
                velocity = ckpt.get("velocity")
                outer_step = ckpt["outer_step"]
                if is_master:
                    log.info(f"resumed outer_step={outer_step}")
        raw = (
            model.module
            if isinstance(model, nn.parallel.DistributedDataParallel)
            else model
        )
        ref_state = {k: v.clone().cpu() for k, v in raw.state_dict().items()}
        if is_master:
            log.info(
                f"DiLoCo/GLOO: {args.num_nodes} nodes, inner_steps={args.inner_steps}, "
                f"master={args.master_ip}:{args.gloo_port}"
            )

    if is_master:
        log.info(f"training for {total_steps} steps on {device}")

    if step == 0 and is_master:
        x, y, _ = next(loader)
        with torch.no_grad(), amp_ctx:
            loss = model(x, y)
        log.info(
            f"step=0 initial loss={loss.item():.4f}  gpu_mem={gpu_mem_mb(device):.0f}MB"
        )

    t0 = time.time()
    tokens_trained = step * tokens_per_step

    while step < total_steps:
        progress = step / total_steps
        lrm = lr_multiplier(progress, min_ratio=args.lr_min_ratio)
        for g in opt.param_groups:
            g["lr"] = g["initial_lr"] * lrm

        opt.zero_grad(set_to_none=True)
        last_loss = 0.0
        for _ in range(grad_accum):
            x, y, epoch = next(loader)
            with amp_ctx:
                loss = model(x, y) / grad_accum
            loss.backward()
            last_loss += loss.item()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

        dt = time.time() - t0
        t0 = time.time()
        tokens_trained += tokens_per_step
        step += 1

        if diloco:
            inner_since_sync += 1

        if step % 10 == 0 and is_master:
            tok_per_sec = int(tokens_per_step / dt)
            cost_1m = (
                (args.instance_price / 3600.0) / (tok_per_sec / 1e6)
                if tok_per_sec > 0
                else 0.0
            )
            log.info(
                f"step={step}/{total_steps} ({100 * progress:.1f}%)  "
                f"loss={last_loss:.4f}  tok/s={tok_per_sec:,}  "
                f"gpu_mem={gpu_mem_mb(device):.0f}MB  epoch={epoch}  cost_1m={cost_1m:.3f}$"
            )

        if step % args.checkpoint_every == 0 and is_master:
            save_ckpt(model, opt, step, ckpt_dir / f"step_{step:06d}.pt")

        if (
            val_tokens is not None
            and args.val_every > 0
            and step % args.val_every == 0
            and is_master
        ):
            val_nats = _eval_val(model, val_tokens, seq_len, device)
            if args.data_format == "bytes":
                log.info(
                    f"val step={step}  val_nats={val_nats:.4f}  "
                    f"val_bpb={val_nats / 0.6931:.4f}"
                )
            else:
                log.info(
                    f"val step={step}  val_nats={val_nats:.4f}  "
                    f"val_bpb_approx={val_nats / 0.6931 / 4.0:.4f}"
                )

        if step == 1:
            gc.collect()
            gc.freeze()
            gc.disable()

        if (
            diloco
            and inner_since_sync >= args.inner_steps
            and (total_steps - step) >= args.inner_steps
        ):
            raw = (
                model.module
                if isinstance(model, nn.parallel.DistributedDataParallel)
                else model
            )
            velocity = diloco_outer_step(
                raw,
                ref_state,
                velocity,
                args.outer_lr,
                args.outer_momentum,
            )
            outer_step += 1
            inner_since_sync = 0
            ref_state = {k: v.clone().cpu() for k, v in raw.state_dict().items()}
            if is_master:
                log.info(f"outer_step={outer_step} complete")
            save_outer_checkpoint(
                s3,
                args.s3_bucket,
                args.run_name,
                outer_step,
                step,
                args.node_rank,
                raw,
                opt.state_dict(),
                velocity,
            )

    if is_master:
        log.info(f"training complete  final_loss={last_loss:.4f}")
        save_ckpt(model, opt, step, ckpt_dir / f"step_{step:06d}_final.pt")

        if args.btm and args.s3_bucket:
            import io

            import boto3

            s3_btm = boto3.client("s3")
            buf = io.BytesIO()
            raw = (
                model.module
                if isinstance(model, nn.parallel.DistributedDataParallel)
                else model
            )
            torch.save(raw.state_dict(), buf)
            buf.seek(0)
            btm_key = f"{args.run_name}/node_{args.node_rank:04d}/latest.pt"
            s3_btm.upload_fileobj(buf, args.s3_bucket, btm_key)
            log.info(f"BTM upload complete → s3://{args.s3_bucket}/{btm_key}")

    if use_ddp:
        dist.destroy_process_group()


def parse_args():
    p = argparse.ArgumentParser(
        description="GPU training script for FNO+GLA language model."
    )
    p.add_argument(
        "--config",
        default="gpt2-small",
        choices=["smoke", "gpt2-small", "gpt2-1b", "gpt2-8b", "byte-125m"],
        help="Model config name",
    )
    p.add_argument("--tokens", type=float, default=2e9, help="Total training tokens")
    p.add_argument(
        "--total-batch",
        type=int,
        default=131072,
        help="Global batch size in tokens per optimizer step",
    )
    p.add_argument(
        "--device-batch", type=int, default=4, help="Sequences per GPU per micro-step"
    )
    p.add_argument("--matrix-lr", type=float, default=0.01)
    p.add_argument("--emb-lr", type=float, default=0.02)
    p.add_argument("--weight-decay", type=float, default=0.1)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--seq-len",
        type=int,
        default=None,
        help="Override config seq_len (e.g. 32768 for long-context)",
    )
    p.add_argument("--n-embd", type=int, default=None)
    p.add_argument("--n-layer", type=int, default=None)
    p.add_argument(
        "--gla-chunk",
        type=int,
        default=None,
        help="GLA intra-chunk size (default: from config)",
    )
    p.add_argument(
        "--lr-min-ratio",
        type=float,
        default=0.1,
        help="Cosine decay floor as fraction of peak LR (default 0.1 = 10%%)",
    )
    p.add_argument("--checkpoint-every", type=int, default=50)
    p.add_argument("--data-dir", default=os.path.expanduser("~/data"))
    p.add_argument("--ckpt-dir", default="checkpoints/gpu_gpt2")
    p.add_argument(
        "--num-gpus",
        type=int,
        default=1,
        help="Number of GPUs on this node (enables DDP when > 1)",
    )
    p.add_argument(
        "--grad-checkpoint",
        action="store_true",
        help="Enable gradient checkpointing per block (needed for seq_len=32K)",
    )
    p.add_argument(
        "--precision",
        default="bf16",
        choices=["bf16", "fp32"],
        help="Training precision (default bf16, A100 supports BF16 natively)",
    )
    p.add_argument(
        "--compile", action="store_true", help="Wrap model with torch.compile(inductor)"
    )
    p.add_argument(
        "--no-compile",
        action="store_true",
        help="Explicitly disable torch.compile (useful for smoke tests)",
    )
    p.add_argument(
        "--data-format",
        default="parquet",
        choices=["parquet", "bin", "text", "bytes"],
        help="Shard format: parquet (default), bin (numpy memmap int32), "
        "text (raw .txt tokenized with tiktoken), or bytes "
        "(streams OWT raw UTF-8 bytes via HuggingFace — no tokenizer)",
    )
    p.add_argument("--node-rank", type=int, default=0)
    p.add_argument("--num-nodes", type=int, default=1)
    p.add_argument("--inner-steps", type=int, default=500)
    p.add_argument("--outer-lr", type=float, default=0.7)
    p.add_argument("--outer-momentum", type=float, default=0.9)
    p.add_argument("--s3-bucket", default=None)
    p.add_argument("--run-name", default="gpu_gpt2")
    p.add_argument(
        "--master-ip",
        default=None,
        help="IP of rank-0 node for GLOO rendezvous (enables DiLoCo when set)",
    )
    p.add_argument("--gloo-port", type=int, default=23456)
    p.add_argument(
        "--btm",
        action="store_true",
        help="Branch-Train-Merge: upload final model to S3 for merge",
    )
    p.add_argument("--val-shard", default=None)
    p.add_argument("--val-every", type=int, default=500)
    p.add_argument(
        "--instance-price",
        type=float,
        default=12.0,
        help="EC2 spot price per hour (default 12.0 for p4d.24xlarge)",
    )
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if args.no_compile:
        args.compile = False
    train(args)