Delete mythos-fineweb.py

mythos-fineweb.py

@@ -1,530 +0,0 @@
-#!/usr/bin/env python3
-"""
-OpenMythos pretraining on FineWeb-Edu with FSDP + AdamW.
-
-Single GPU:
-    python training/3b_fine_web_edu.py
-
-Multi-GPU:
-    torchrun --nproc_per_node=$(python -c "import torch; print(torch.cuda.device_count())") training/3b_fine_web_edu.py
-"""
-
-import os
-import math
-import time
-import torch
-import torch.nn as nn
-import torch.distributed as dist
-from loguru import logger
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-    ShardingStrategy,
-    MixedPrecision,
-    FullStateDictConfig,
-    StateDictType,
-)
-from torch.distributed.fsdp.wrap import ModuleWrapPolicy
-from torch.utils.data import IterableDataset, DataLoader, get_worker_info
-from contextlib import nullcontext
-
-from datasets import load_dataset
-
-from open_mythos import OpenMythos
-from open_mythos.main import TransformerBlock, RecurrentBlock
-from open_mythos.variants import mythos_3b
-from open_mythos.tokenizer import MythosTokenizer
-
-
-# ---------------------------------------------------------------------------
-# Dataset
-# ---------------------------------------------------------------------------
-
-
-class FineWebEduDataset(IterableDataset):
-    """
-    Streaming FineWeb-Edu loader yielding fixed-length (input, target) pairs.
-
-    FineWeb-Edu is trillions of tokens, so `streaming=True` pulls shards on
-    demand instead of materializing to disk. Sharding is two-dimensional —
-    `world_size` ranks × `num_workers` DataLoader workers per rank — and each
-    `(rank, worker_id)` deterministically owns one shard of the global stream.
-    That gives disjoint coverage without any cross-process coordination.
-
-    Streaming datasets are not seekable, so a resumed run re-enters its shard
-    from the beginning. Acceptable at pretraining scale: the chance of
-    re-playing the same tokens before the run ends is negligible versus the
-    cost of a true resumable loader.
-    """
-
-    def __init__(self, encoding, seq_len: int, subset: str, rank: int, world_size: int):
-        """
-        Args:
-            encoding   -- tokenizer exposing `.encode(str) -> list[int]`
-            seq_len    -- context length; every yielded pair has this many tokens
-            subset     -- FineWeb-Edu config name (e.g. "sample-10BT", "default")
-            rank       -- global rank of this process within the distributed job
-            world_size -- total number of distributed processes
-        """
-        self.encoding = encoding
-        self.seq_len = seq_len
-        self.subset = subset
-        self.rank = rank
-        self.world_size = world_size
-
-    def __iter__(self):
-        """
-        Yield `(input_ids, target_ids)` tensors of length `seq_len` forever.
-
-        Inputs and targets are shifted by one for next-token prediction —
-        `target[i] == input[i + 1]`. Documents are concatenated into a rolling
-        buffer and sliced into fixed-length chunks, packing short docs together
-        and splitting long ones. This keeps every step at the same shape,
-        which under FSDP avoids recompute from variable-length inputs and
-        removes the need for a pad-aware attention mask.
-        """
-        worker = get_worker_info()
-        num_workers = worker.num_workers if worker else 1
-        worker_id = worker.id if worker else 0
-
-        total_shards = self.world_size * num_workers
-        shard_index = self.rank * num_workers + worker_id
-
-        ds = load_dataset(
-            "HuggingFaceFW/fineweb-edu",
-            name=self.subset,
-            split="train",
-            streaming=True,
-        ).shard(num_shards=total_shards, index=shard_index)
-
-        buf = []
-        for sample in ds:
-            buf.extend(self.encoding.encode(sample["text"]))
-            while len(buf) >= self.seq_len + 1:
-                chunk = buf[: self.seq_len + 1]
-                buf = buf[self.seq_len + 1 :]
-                yield (
-                    torch.tensor(chunk[:-1], dtype=torch.long),
-                    torch.tensor(chunk[1:], dtype=torch.long),
-                )
-
-
-# ---------------------------------------------------------------------------
-# LR schedule: linear warmup → cosine decay
-# ---------------------------------------------------------------------------
-
-
-def get_lr(step: int, warmup: int, total: int, max_lr: float, min_lr: float) -> float:
-    """
-    Linear warmup → half-cosine decay to `min_lr`.
-
-    Standard language-model pretraining schedule. The warmup phase prevents
-    Adam's second-moment estimate from collapsing to a huge LR in the first
-    few steps when gradients are noisy. The cosine tail lets the model make
-    small, increasingly conservative updates near the end of training rather
-    than crashing to `min_lr` at a fixed step.
-
-    Behavior by region:
-        step < warmup                 → linear ramp 0 → max_lr
-        warmup ≤ step < total         → cosine decay max_lr → min_lr
-        step ≥ total                  → clamped at min_lr (safety for
-                                        off-by-one step counters at the end
-                                        of training)
-
-    Args:
-        step    -- current global optimizer step (0-indexed)
-        warmup  -- number of warmup steps before cosine decay begins
-        total   -- step at which the cosine reaches `min_lr`
-        max_lr  -- peak learning rate reached at the end of warmup
-        min_lr  -- floor learning rate at and after `total` steps
-
-    Returns:
-        Scalar learning rate for this step.
-    """
-    if step < warmup:
-        return max_lr * step / warmup
-    if step >= total:
-        return min_lr
-    decay = (step - warmup) / (total - warmup)
-    return min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * decay))
-
-
-# ---------------------------------------------------------------------------
-# Checkpointing — weights-only every 500 steps, full at epoch end + best
-# ---------------------------------------------------------------------------
-
-
-def save_weights_only(model, step, epoch, ckpt_dir, ddp):
-    """Save model weights only (~1.3GB for 3B bf16). For testing/transfer."""
-    if ddp:
-        with FSDP.state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-        ):
-            model_state = model.state_dict()
-    else:
-        model_state = model.state_dict()
-
-    ckpt_path = os.path.join(ckpt_dir, f"spiderportal-v5-ep{epoch}-step{step}.pt")
-    tmp_path = ckpt_path + ".tmp"
-    torch.save(model_state, tmp_path)
-    os.replace(tmp_path, ckpt_path)
-    size_mb = os.path.getsize(ckpt_path) / (1024 * 1024)
-    return ckpt_path, size_mb
-
-
-def save_full_checkpoint(model, optimizer, step, epoch, cfg, vocab_size, ckpt_dir, ddp, master, ckpt_name="full"):
-    """Save model + optimizer state (~18GB for 3B bf16). For resume training."""
-    if ddp:
-        with FSDP.state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-        ):
-            model_state = model.state_dict()
-            optim_state = FSDP.optim_state_dict(model, optimizer)
-    else:
-        model_state = model.state_dict()
-        optim_state = optimizer.state_dict()
-
-    if not master:
-        return None, 0
-
-    os.makedirs(ckpt_dir, exist_ok=True)
-    final_path = os.path.join(ckpt_dir, f"spiderportal-v5-{ckpt_name}.pt")
-    tmp_path = final_path + ".tmp"
-    torch.save(
-        {
-            "step": step,
-            "epoch": epoch,
-            "model_state_dict": model_state,
-            "optimizer_state_dict": optim_state,
-            "cfg": cfg,
-            "vocab_size": vocab_size,
-        },
-        tmp_path,
-    )
-    os.replace(tmp_path, final_path)
-    size_mb = os.path.getsize(final_path) / (1024 * 1024)
-    return final_path, size_mb
-
-
-def delete_step_checkpoints(ckpt_dir):
-    """Delete all weights-only step checkpoints to free disk space."""
-    deleted = 0
-    for f in os.listdir(ckpt_dir):
-        if f.startswith("spiderportal-v5-ep") and "-step" in f and f.endswith(".pt"):
-            path = os.path.join(ckpt_dir, f)
-            try:
-                os.remove(path)
-                deleted += 1
-            except OSError:
-                pass
-    return deleted
-
-
-def load_checkpoint(model, optimizer, path, ddp):
-    """Restore model + optimizer from full checkpoint."""
-    ckpt = torch.load(path, map_location="cpu", weights_only=False)
-
-    if ddp:
-        with FSDP.state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(offload_to_cpu=True, rank0_only=False),
-        ):
-            model.load_state_dict(ckpt["model_state_dict"])
-            optim_state = FSDP.optim_state_dict_to_load(
-                model=model,
-                optim=optimizer,
-                optim_state_dict=ckpt["optimizer_state_dict"],
-            )
-            optimizer.load_state_dict(optim_state)
-    else:
-        model.load_state_dict(ckpt["model_state_dict"])
-        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
-
-    return int(ckpt["step"]), int(ckpt.get("epoch", 0))
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-
-def main():
-    """
-    End-to-end pretraining entry point.
-
-    Order matters: distributed init must run before any CUDA allocation, the
-    tokenizer must exist before the model is built (vocab_size flows into
-    cfg), and FSDP must wrap the model before the optimizer is constructed
-    (FSDP re-flattens parameters, so an optimizer built on the unwrapped
-    model would track stale param objects). Resume then loads state into the
-    already-constructed optimizer in-place.
-
-    Lifecycle:
-        1. Initialize torch.distributed (NCCL) if launched under torchrun.
-        2. Build tokenizer → derive vocab_size.
-        3. Construct OpenMythos with the 3B variant config.
-        4. Wrap in FSDP with FULL_SHARD + bf16/fp16 mixed precision (multi-GPU)
-           or move to device + autocast (single-GPU).
-        5. Build fused AdamW on (possibly sharded) parameters.
-        6. Resume from the latest checkpoint in `ckpt_dir` if one exists.
-        7. Stream FineWeb-Edu through grad-accumulation microbatches with
-           cosine LR schedule, per-step logging, and periodic checkpoints.
-        8. Write a final checkpoint if the last save wasn't aligned to
-           `ckpt_every`, then barrier + tear down the process group.
-
-    All hyperparameters are literal constants in this function by design —
-    pretraining runs are long-lived and each run pins exact settings; a
-    CLI/config layer is deliberately avoided to keep the file self-auditable.
-    """
-    # ------------------------------------------------------------------
-    # Distributed init
-    # ------------------------------------------------------------------
-    ddp = int(os.environ.get("RANK", -1)) != -1
-    if ddp:
-        dist.init_process_group("nccl")
-        rank = int(os.environ["RANK"])
-        local_rank = int(os.environ["LOCAL_RANK"])
-        world_size = int(os.environ["WORLD_SIZE"])
-        device = f"cuda:{local_rank}"
-        torch.cuda.set_device(device)
-    else:
-        rank = local_rank = 0
-        world_size = 1
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    master = rank == 0
-
-    if master:
-        logger.info(
-            f"GPUs: {torch.cuda.device_count()}  |  World size: {world_size}  |  Device: {device}"
-        )
-
-    # ------------------------------------------------------------------
-    # Tokenizer
-    # ------------------------------------------------------------------
-    encoding = MythosTokenizer()
-    vocab_size = encoding.vocab_size
-
-    if master:
-        logger.info(f"Tokenizer: gpt-oss-20b  |  Vocab size: {vocab_size:,}")
-
-    # ------------------------------------------------------------------
-    # Hyperparameters
-    # ------------------------------------------------------------------
-    seq_len = 2048
-    micro_batch = 32
-    target_tokens = 1_000_000_000
-    grad_accum = max(1, 256 // (world_size * micro_batch))
-    global_batch_tok = world_size * micro_batch * grad_accum * seq_len
-    total_steps = target_tokens // global_batch_tok
-    warmup_steps = 200
-    lr = 3e-4
-    wd = 0.1
-    log_every = 10
-    ckpt_every = 500
-    ckpt_dir = "checkpoints"
-    dataset_subset = "sample-1BT"
-
-    if master:
-        logger.info(
-            f"seq_len={seq_len} | micro_batch={micro_batch} | grad_accum={grad_accum} | "
-            f"global_batch_tokens={global_batch_tok:,} | total_steps={total_steps:,}"
-        )
-
-    # ------------------------------------------------------------------
-    # Model
-    # ------------------------------------------------------------------
-    cfg = mythos_3b()
-    cfg.vocab_size = vocab_size
-    cfg.max_seq_len = seq_len
-
-    bf16_ok = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
-    amp_dtype = torch.bfloat16 if bf16_ok else torch.float16
-
-    model = OpenMythos(cfg)
-
-    if ddp:
-        mp_policy = MixedPrecision(
-            param_dtype=amp_dtype,
-            reduce_dtype=amp_dtype,
-            buffer_dtype=amp_dtype,
-        )
-        wrap_policy = ModuleWrapPolicy({TransformerBlock, RecurrentBlock})
-        model = FSDP(
-            model,
-            sharding_strategy=ShardingStrategy.FULL_SHARD,
-            mixed_precision=mp_policy,
-            auto_wrap_policy=wrap_policy,
-            device_id=local_rank,
-        )
-    else:
-        model = model.to(device)
-        amp_ctx = (
-            torch.amp.autocast(device_type="cuda", dtype=amp_dtype)
-            if "cuda" in device
-            else nullcontext()
-        )
-
-    # FSDP handles its own mixed precision; only need autocast for single-GPU
-    amp_ctx = nullcontext() if ddp else amp_ctx  # type: ignore[possibly-undefined]
-
-    if master:
-        n_params = sum(p.numel() for p in model.parameters())
-        logger.info(f"Parameters: {n_params:,}  |  AMP dtype: {amp_dtype}")
-
-    # Compile for 20-30% speedup (requires PyTorch 2.0+)
-    try:
-        model = torch.compile(model, mode="reduce-overhead")
-        if master:
-            logger.info("torch.compile: enabled")
-    except Exception:
-        if master:
-            logger.info("torch.compile: not available, using eager mode")
-
-    # ------------------------------------------------------------------
-    # Optimizer
-    # ------------------------------------------------------------------
-    optimizer = torch.optim.AdamW(
-        model.parameters(), lr=lr, weight_decay=wd, betas=(0.9, 0.95), fused=True
-    )
-
-    # ------------------------------------------------------------------
-    # Resume from latest checkpoint (if any)
-    # ------------------------------------------------------------------
-    start_step = 0
-    start_epoch = 1
-    best_loss = float("inf")
-    existing_ckpts = [f for f in os.listdir(ckpt_dir) if f.startswith("spiderportal-v5-ep") and f.endswith(".pt") and "-step" not in f] if os.path.isdir(ckpt_dir) else []
-    if existing_ckpts:
-        latest = os.path.join(ckpt_dir, sorted(existing_ckpts)[-1])
-        if master:
-            logger.info(f"Resuming from checkpoint: {latest}")
-        start_step, start_epoch = load_checkpoint(model, optimizer, latest, ddp)
-        if master:
-            logger.success(f"Resumed at step {start_step}, epoch {start_epoch}")
-
-    # ------------------------------------------------------------------
-    # Dataset + DataLoader
-    # ------------------------------------------------------------------
-    dataset = FineWebEduDataset(encoding, seq_len, dataset_subset, rank, world_size)
-    loader = DataLoader(dataset, batch_size=micro_batch, num_workers=8, pin_memory=True, prefetch_factor=2)
-
-    # ------------------------------------------------------------------
-    # Training loop
-    # ------------------------------------------------------------------
-    if master:
-        os.makedirs(ckpt_dir, exist_ok=True)
-
-    model.train()
-    data_iter = iter(loader)
-    t0 = time.perf_counter()
-    step = start_step
-    epoch = start_epoch
-    step_ckpt_files = []
-    tokens_in_epoch = 0
-    tokens_per_epoch = target_tokens
-
-    while step < total_steps:
-        cur_lr = get_lr(step, warmup_steps, total_steps, lr, lr * 0.1)
-        for g in optimizer.param_groups:
-            g["lr"] = cur_lr
-
-        optimizer.zero_grad()
-        loss_accum = 0.0
-
-        for micro_step in range(grad_accum):
-            try:
-                x, y = next(data_iter)
-            except StopIteration:
-                data_iter = iter(loader)
-                x, y = next(data_iter)
-
-            x = x.to(device if not ddp else f"cuda:{local_rank}", non_blocking=True)
-            y = y.to(device if not ddp else f"cuda:{local_rank}", non_blocking=True)
-
-            sync = (
-                nullcontext()
-                if (not ddp or micro_step == grad_accum - 1)
-                else model.no_sync()
-            )
-            with sync, amp_ctx:
-                logits = model(x)
-                loss = nn.functional.cross_entropy(
-                    logits.view(-1, vocab_size), y.view(-1)
-                )
-                loss = loss / grad_accum
-
-            loss.backward()
-            loss_accum += loss.item()
-
-        if ddp:
-            grad_norm = model.clip_grad_norm_(1.0)
-        else:
-            grad_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-        optimizer.step()
-        step += 1
-        tokens_in_epoch += global_batch_tok
-
-        if master and step % log_every == 0:
-            dt = time.perf_counter() - t0
-            tok_per_sec = global_batch_tok * log_every / dt
-            tokens_seen = step * global_batch_tok
-            logger.info(
-                f"Epoch {epoch} | step {step:6d}/{total_steps} | loss {loss_accum:.4f} "
-                f"| gnorm {float(grad_norm):.2f} | lr {cur_lr:.2e} "
-                f"| {tok_per_sec / 1e6:.2f}M tok/s "
-                f"| {tokens_seen / 1e9:.2f}B tokens seen"
-            )
-            t0 = time.perf_counter()
-
-        if step % ckpt_every == 0 and master:
-            ckpt_path, size_mb = save_weights_only(model, step, epoch, ckpt_dir, ddp)
-            step_ckpt_files.append(ckpt_path)
-            logger.info(f"Saved weights-only: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")
-
-        if tokens_in_epoch >= tokens_per_epoch:
-            epoch_loss = loss_accum
-            if master:
-                epoch_time = (time.perf_counter() - t0) / 60
-                logger.info(f"Epoch {epoch} complete | loss={epoch_loss:.4f} | Time: {epoch_time:.1f}min")
-
-                for f in step_ckpt_files:
-                    if os.path.exists(f):
-                        os.remove(f)
-                        logger.info(f"  Deleted step checkpoint: {os.path.basename(f)}")
-                step_ckpt_files.clear()
-
-                ckpt_path, size_mb = save_full_checkpoint(model, optimizer, step, epoch, cfg, vocab_size, ckpt_dir, ddp, master, f"ep{epoch}")
-                if ckpt_path:
-                    logger.info(f"Saved epoch checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")
-
-                if epoch_loss < best_loss:
-                    best_loss = epoch_loss
-                    ckpt_path, size_mb = save_full_checkpoint(model, optimizer, step, epoch, cfg, vocab_size, ckpt_dir, ddp, master, "best")
-                    if ckpt_path:
-                        logger.info(f"Saved best checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")
-
-            epoch += 1
-            tokens_in_epoch = 0
-
-    if step > start_step and master:
-        ckpt_path, size_mb = save_full_checkpoint(model, optimizer, step, epoch, cfg, vocab_size, ckpt_dir, ddp, master, f"final-ep{epoch}")
-        if ckpt_path:
-            logger.info(f"Saved final checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")
-
-    if ddp:
-        # Barrier so no rank exits while another is still finishing its
-        # checkpoint gather — avoids NCCL "process group destroyed" noise.
-        dist.barrier()
-        dist.destroy_process_group()
-
-    if master:
-        logger.success("Training complete.")
-
-
-if __name__ == "__main__":
-    main()