Upload train_inverter_v6.py with huggingface_hub

train_inverter_v6.py

@@ -0,0 +1,695 @@
+#!/usr/bin/env python
+import argparse
+import json
+import math
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+
+from v6_model import EncoderOnlyModel
+
+
+@dataclass
+class TrainState:
+    tokens_seen: int = 0
+    example_index: int = 0
+    example_token_offset: int = 0
+    step: int = 0
+
+
+def _write_json_atomic(path, payload):
+    tmp = f"{path}.tmp"
+    with open(tmp, "w") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+    os.replace(tmp, path)
+
+
+def _torch_save_atomic(path, payload):
+    tmp = f"{path}.tmp"
+    torch.save(payload, tmp)
+    os.replace(tmp, path)
+
+
+def _unwrap_model(model):
+    return model._orig_mod if hasattr(model, "_orig_mod") else model
+
+
+class ExpertStream:
+    def __init__(
+        self,
+        idx_path,
+        tokens_path,
+        doc_starts_path,
+        dataset_name,
+        dataset_revision,
+        tokenizer,
+        seq_len,
+        min_seq_len,
+        stride,
+        max_tokens,
+        batch_size,
+        state: TrainState,
+    ):
+        self.idx_path = idx_path
+        self.tokens_path = tokens_path
+        self.doc_starts_path = doc_starts_path
+        self.dataset_name = dataset_name
+        self.dataset_revision = dataset_revision
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.min_seq_len = min(min_seq_len, seq_len)
+        self.stride = stride if stride > 0 else max(1, seq_len // 2)
+        self.max_tokens = max_tokens
+        self.batch_size = batch_size
+        self.state = state
+
+        self.idx_mmap = np.load(self.idx_path, mmap_mode="r")
+        if self.idx_mmap.shape != (self.max_tokens, 24, 4):
+            raise ValueError(f"Unexpected idx shape {self.idx_mmap.shape}.")
+        self.tokens_mmap = None
+        self.doc_starts = None
+        if self.tokens_path:
+            self.tokens_mmap = np.load(self.tokens_path, mmap_mode="r")
+            if self.tokens_mmap.shape != (self.max_tokens,):
+                raise ValueError(f"Unexpected token shape {self.tokens_mmap.shape}.")
+            if self.doc_starts_path and os.path.exists(self.doc_starts_path):
+                doc_starts_mmap = np.load(self.doc_starts_path, mmap_mode="r")
+                if doc_starts_mmap.shape != (self.max_tokens,):
+                    raise ValueError(
+                        f"Unexpected doc-start shape {doc_starts_mmap.shape}."
+                    )
+                starts = np.flatnonzero(doc_starts_mmap[: self.max_tokens])
+                if len(starts) == 0 or starts[0] != 0:
+                    starts = np.concatenate(
+                        [np.array([0], dtype=np.int64), starts[starts > 0]]
+                    )
+                self.doc_starts = np.concatenate(
+                    [
+                        starts.astype(np.int64, copy=False),
+                        np.array([self.max_tokens], dtype=np.int64),
+                    ]
+                )
+
+    def _yield_batch(self, batch_tokens, batch_idx, batch_mask, state):
+        return {
+            "input_ids": torch.tensor(np.asarray(batch_tokens, dtype=np.int64)),
+            "expert_idx": torch.tensor(np.asarray(batch_idx, dtype=np.int64)),
+            "attention_mask": torch.tensor(np.asarray(batch_mask, dtype=np.bool_)),
+            "state": state,
+        }
+
+    def _pad_chunk(self, chunk, idx_chunk, current_len):
+        pad_len = self.seq_len - current_len
+        if pad_len:
+            chunk = np.pad(
+                chunk,
+                (0, pad_len),
+                mode="constant",
+                constant_values=self.tokenizer.pad_token_id,
+            )
+            idx_chunk = np.pad(
+                idx_chunk,
+                ((0, pad_len), (0, 0), (0, 0)),
+                mode="constant",
+                constant_values=0,
+            )
+        mask = [1] * current_len + [0] * pad_len
+        return chunk, idx_chunk, mask
+
+    def _pick_window_len(self, doc_start: int, start_pos: int, max_len: int) -> int:
+        if max_len <= self.min_seq_len:
+            return max_len
+        span = max_len - self.min_seq_len + 1
+        window_idx = (start_pos - doc_start) // self.stride
+        mix = (1103515245 * (window_idx + 1) + 12345 * (doc_start + 1)) & 0x7FFFFFFF
+        return self.min_seq_len + (mix % span)
+
+    def _local_doc_bounds(self, token_pos):
+        if self.doc_starts is None:
+            return 0, self.max_tokens, 0
+        doc_idx = max(0, int(np.searchsorted(self.doc_starts, token_pos, side="right")) - 1)
+        doc_start = int(self.doc_starts[doc_idx])
+        doc_end = int(self.doc_starts[doc_idx + 1])
+        return doc_start, doc_end, doc_idx
+
+    def _iter_local(self):
+        next_start = self.state.tokens_seen
+        batch_tokens = []
+        batch_idx = []
+        batch_mask = []
+
+        while next_start < self.max_tokens:
+            doc_start, doc_end, doc_idx = self._local_doc_bounds(next_start)
+            if next_start < doc_start:
+                next_start = doc_start
+            if next_start >= doc_end:
+                next_start = doc_end
+                continue
+
+            current_max_len = min(
+                self.seq_len,
+                doc_end - next_start,
+                self.max_tokens - next_start,
+            )
+            if current_max_len <= 0:
+                break
+
+            current_len = self._pick_window_len(doc_start, next_start, current_max_len)
+            chunk = self.tokens_mmap[next_start : next_start + current_len].astype(
+                np.int64, copy=False
+            )
+            idx_chunk = self.idx_mmap[next_start : next_start + current_len]
+            chunk, idx_chunk, mask = self._pad_chunk(chunk, idx_chunk, current_len)
+
+            new_offset = min(doc_end - doc_start, next_start + self.stride - doc_start)
+            state = TrainState(
+                tokens_seen=min(doc_end, next_start + self.stride),
+                example_index=doc_idx,
+                example_token_offset=new_offset,
+                step=self.state.step,
+            )
+            next_start = state.tokens_seen
+
+            batch_tokens.append(chunk)
+            batch_idx.append(idx_chunk)
+            batch_mask.append(mask)
+
+            if len(batch_tokens) >= self.batch_size:
+                yield self._yield_batch(batch_tokens, batch_idx, batch_mask, state)
+                batch_tokens, batch_idx, batch_mask = [], [], []
+
+        if batch_tokens:
+            yield self._yield_batch(
+                batch_tokens,
+                batch_idx,
+                batch_mask,
+                TrainState(
+                    tokens_seen=next_start,
+                    example_index=0,
+                    example_token_offset=0,
+                    step=self.state.step,
+                ),
+            )
+
+    def __iter__(self):
+        if self.tokens_mmap is not None:
+            yield from self._iter_local()
+            return
+
+        from datasets import load_dataset
+
+        ds = load_dataset(
+            self.dataset_name,
+            split="train",
+            streaming=True,
+            revision=self.dataset_revision,
+        )
+        tokens_seen = self.state.tokens_seen
+        example_index = self.state.example_index
+        example_token_offset = self.state.example_token_offset
+
+        batch_tokens = []
+        batch_idx = []
+        batch_mask = []
+
+        for idx, example in enumerate(ds):
+            if idx < example_index:
+                continue
+            if tokens_seen >= self.max_tokens:
+                break
+
+            token_ids = self.tokenizer.encode(example["text"], add_special_tokens=False)
+            doc_len = len(token_ids)
+            if not token_ids:
+                example_index = idx + 1
+                example_token_offset = 0
+                continue
+
+            if idx == example_index:
+                pos = example_token_offset
+                doc_abs_start = tokens_seen - example_token_offset
+            else:
+                pos = 0
+                doc_abs_start = tokens_seen
+
+            while pos < doc_len and tokens_seen < self.max_tokens:
+                abs_pos = doc_abs_start + pos
+                current_max_len = min(
+                    self.seq_len,
+                    doc_len - pos,
+                    self.max_tokens - abs_pos,
+                )
+                if current_max_len <= 0:
+                    break
+
+                current_len = self._pick_window_len(doc_abs_start, abs_pos, current_max_len)
+                chunk = np.asarray(token_ids[pos : pos + current_len], dtype=np.int64)
+                idx_chunk = self.idx_mmap[abs_pos : abs_pos + current_len]
+                chunk, idx_chunk, mask = self._pad_chunk(chunk, idx_chunk, current_len)
+
+                next_pos = min(doc_len, pos + self.stride)
+                next_abs = doc_abs_start + next_pos
+                state = TrainState(
+                    tokens_seen=next_abs,
+                    example_index=idx,
+                    example_token_offset=next_pos,
+                    step=self.state.step,
+                )
+
+                batch_tokens.append(chunk)
+                batch_idx.append(idx_chunk)
+                batch_mask.append(mask)
+
+                pos = next_pos
+                tokens_seen = next_abs
+                example_token_offset = next_pos
+
+                if len(batch_tokens) >= self.batch_size:
+                    yield self._yield_batch(batch_tokens, batch_idx, batch_mask, state)
+                    batch_tokens, batch_idx, batch_mask = [], [], []
+
+            example_index = idx + 1
+            example_token_offset = 0
+
+        if batch_tokens:
+            yield self._yield_batch(
+                batch_tokens,
+                batch_idx,
+                batch_mask,
+                TrainState(
+                    tokens_seen=tokens_seen,
+                    example_index=example_index,
+                    example_token_offset=example_token_offset,
+                    step=self.state.step,
+                ),
+            )
+
+
+def _derive_companion_path(idx_path: str, suffix: str) -> str | None:
+    path = Path(idx_path)
+    if path.name.endswith("_idx.npy"):
+        candidate = path.with_name(f"{path.name[:-len('_idx.npy')]}_{suffix}.npy")
+        if candidate.exists():
+            return str(candidate)
+    return None
+
+
+def _load_train_state(args, checkpoint_payload):
+    if checkpoint_payload and "train_state" in checkpoint_payload:
+        return TrainState(**checkpoint_payload["train_state"])
+    if args.resume and os.path.exists(args.state_path):
+        with open(args.state_path, "r") as f:
+            payload = json.load(f)
+        return TrainState(
+            tokens_seen=payload.get("tokens_seen", 0),
+            example_index=payload.get("example_index", 0),
+            example_token_offset=payload.get("example_token_offset", 0),
+            step=payload.get("step", 0),
+        )
+    return TrainState()
+
+
+def _load_checkpoint(path, model, optimizers, schedulers, scaler, device):
+    if not os.path.exists(path):
+        return None
+    payload = torch.load(path, map_location=device)
+    _unwrap_model(model).load_state_dict(payload["model"])
+    optimizer_states = payload.get("optimizers")
+    if optimizer_states is not None:
+        for opt, opt_state in zip(optimizers, optimizer_states):
+            opt.load_state_dict(opt_state)
+    scheduler_states = payload.get("schedulers")
+    if scheduler_states is not None:
+        for sched, sched_state in zip(schedulers, scheduler_states):
+            sched.load_state_dict(sched_state)
+    scaler_state = payload.get("scaler")
+    if scaler_state and scaler.is_enabled():
+        scaler.load_state_dict(scaler_state)
+    return payload
+
+
+def _save_checkpoint(path, model, args, state, step, optimizers, schedulers, scaler):
+    payload = {
+        "model": _unwrap_model(model).state_dict(),
+        "optimizers": [opt.state_dict() for opt in optimizers],
+        "schedulers": [sched.state_dict() for sched in schedulers],
+        "scaler": scaler.state_dict() if scaler.is_enabled() else None,
+        "config": vars(args),
+        "train_state": state.__dict__,
+        "step": step,
+    }
+    _torch_save_atomic(path, payload)
+
+
+def split_muon_params(model):
+    muon_params = []
+    adam_params = []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        is_matrix = p.ndim == 2
+        is_embedding_or_head = name.endswith("head.weight") or name.endswith("pos_emb.weight")
+        if is_matrix and not is_embedding_or_head:
+            muon_params.append(p)
+        else:
+            adam_params.append(p)
+    return muon_params, adam_params
+
+
+def make_trapezoidal_lr(step_idx, max_steps, warmup_ratio, warmdown_ratio):
+    warmup_steps = max(1, int(warmup_ratio * max_steps)) if warmup_ratio > 0 else 0
+    warmdown_steps = max(1, int(warmdown_ratio * max_steps)) if warmdown_ratio > 0 else 0
+    if warmup_steps > 0 and step_idx < warmup_steps:
+        return float(step_idx + 1) / float(warmup_steps)
+    warmdown_start = max_steps - warmdown_steps
+    if step_idx < warmdown_start:
+        return 1.0
+    if warmdown_steps > 0 and step_idx < max_steps:
+        remaining = max_steps - step_idx
+        return max(0.0, float(remaining) / float(warmdown_steps))
+    return 0.0
+
+
+def _set_tf32():
+    if not torch.cuda.is_available():
+        return
+    try:
+        torch.backends.cuda.matmul.fp32_precision = "tf32"
+        torch.backends.cudnn.conv.fp32_precision = "tf32"
+    except AttributeError:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    torch.set_float32_matmul_precision("high")
+
+
+def _estimate_schedule_steps(max_tokens: int, batch_size: int, stride: int, grad_accum: int) -> int:
+    unique_tokens_per_step = max(1, batch_size * stride * grad_accum)
+    return max(1, math.ceil(max_tokens / unique_tokens_per_step))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="V6 inverter trainer (ctx256 + RoPE + QK-Norm + aligned schedule)."
+    )
+    parser.add_argument("--idx", required=True)
+    parser.add_argument("--tokens", default=None)
+    parser.add_argument("--doc-starts", default=None)
+    parser.add_argument("--dataset", default="vietgpt/openwebtext_en")
+    parser.add_argument("--dataset-revision", default=None)
+    parser.add_argument("--model", default="openai/gpt-oss-20b")
+    parser.add_argument("--model-revision", default=None)
+    parser.add_argument("--seq-len", type=int, default=256)
+    parser.add_argument("--min-seq-len", type=int, default=128)
+    parser.add_argument("--stride", type=int, default=128)
+    parser.add_argument("--layers", type=int, default=24)
+    parser.add_argument("--max-tokens", type=int, default=200000000)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--grad-accum", type=int, default=2)
+    parser.add_argument("--steps", type=int, default=50000)
+    parser.add_argument("--schedule-steps", type=int, default=0)
+    parser.add_argument("--save-every", type=int, default=1000)
+    parser.add_argument("--log-every", type=int, default=50)
+    parser.add_argument("--out", default="inverter_v6.pt")
+    parser.add_argument("--state-path", default="train_state_v6.json")
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--compile", action="store_true")
+    parser.add_argument(
+        "--attn-impl",
+        choices=["auto", "flash", "mem_efficient", "math"],
+        default="auto",
+    )
+    parser.add_argument("--position-type", choices=["learned", "rope"], default="rope")
+    parser.add_argument("--rope-theta", type=float, default=10000.0)
+    parser.add_argument("--qk-norm", action="store_true", default=True)
+    parser.add_argument("--no-qk-norm", action="store_false", dest="qk_norm")
+    parser.add_argument("--qk-norm-eps", type=float, default=1e-5)
+    parser.add_argument("--logit-softcap", type=float, default=0.0)
+    parser.add_argument("--layer-gating", action="store_true")
+    parser.add_argument("--d-model", type=int, default=768)
+    parser.add_argument("--n-head", type=int, default=12)
+    parser.add_argument("--d-ff", type=int, default=2048)
+    parser.add_argument("--n-layer", type=int, default=6)
+    parser.add_argument("--layer-hidden", type=int, default=64)
+    parser.add_argument("--layer-proj", type=int, default=64)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--adam-lr", type=float, default=1.75e-4)
+    parser.add_argument("--muon-lr-factor", type=float, default=4.0)
+    parser.add_argument("--weight-decay", type=float, default=0.1)
+    parser.add_argument("--warmup-ratio", type=float, default=0.02)
+    parser.add_argument("--warmdown-ratio", type=float, default=0.40)
+    parser.add_argument("--label-smoothing", type=float, default=0.05)
+    parser.add_argument("--clip-grad-norm", type=float, default=1.0)
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--wandb-project", default="expert-inversion")
+    parser.add_argument("--wandb-entity", default=None)
+    parser.add_argument("--wandb-run-name", default=None)
+    args = parser.parse_args()
+
+    _set_tf32()
+    if torch.cuda.is_available() and args.attn_impl != "auto":
+        try:
+            torch.backends.cuda.enable_flash_sdp(args.attn_impl == "flash")
+            torch.backends.cuda.enable_mem_efficient_sdp(
+                args.attn_impl == "mem_efficient"
+            )
+            torch.backends.cuda.enable_math_sdp(args.attn_impl == "math")
+        except AttributeError:
+            pass
+
+    wandb_run = None
+    if args.wandb:
+        import wandb
+
+        wandb_run = wandb.init(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=args.wandb_run_name,
+            config=vars(args),
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        revision=args.model_revision,
+        local_files_only=True,
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    if args.tokens is None:
+        args.tokens = _derive_companion_path(args.idx, "tok")
+    if args.doc_starts is None:
+        args.doc_starts = _derive_companion_path(args.idx, "doc_start")
+    if args.tokens:
+        print(f"Using local token mmap: {args.tokens}")
+        if args.doc_starts:
+            print(f"Using local doc-start mmap: {args.doc_starts}")
+    else:
+        print("Local token mmap not found; falling back to dataset streaming.")
+
+    schedule_steps = (
+        args.schedule_steps
+        if args.schedule_steps > 0
+        else min(
+            args.steps,
+            _estimate_schedule_steps(
+                max_tokens=args.max_tokens,
+                batch_size=args.batch_size,
+                stride=args.stride,
+                grad_accum=args.grad_accum,
+            ),
+        )
+    )
+    unique_tokens_per_step = args.batch_size * args.stride * args.grad_accum
+    print(
+        f"Schedule horizon: {schedule_steps} steps "
+        f"(estimated unique tokens/step: {unique_tokens_per_step})"
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = EncoderOnlyModel(
+        vocab_size=len(tokenizer),
+        num_experts=32,
+        num_layers=args.layers,
+        topk=4,
+        d_model=args.d_model,
+        n_head=args.n_head,
+        d_ff=args.d_ff,
+        n_layer=args.n_layer,
+        dropout=args.dropout,
+        max_len=args.seq_len,
+        layer_gating=args.layer_gating,
+        logit_softcap=args.logit_softcap,
+        layer_hidden=args.layer_hidden,
+        layer_proj=args.layer_proj,
+        position_type=args.position_type,
+        rope_theta=args.rope_theta,
+        qk_norm=args.qk_norm,
+        qk_norm_eps=args.qk_norm_eps,
+    ).to(device)
+
+    muon_params, adam_params = split_muon_params(model)
+    if not muon_params:
+        raise RuntimeError("No Muon parameters found; check parameter names.")
+    if not hasattr(torch.optim, "Muon"):
+        raise RuntimeError("torch.optim.Muon not available in this environment.")
+
+    optimizer_adam = torch.optim.AdamW(
+        adam_params,
+        lr=args.adam_lr,
+        betas=(0.9, 0.95),
+        weight_decay=args.weight_decay,
+    )
+    optimizer_muon = torch.optim.Muon(
+        muon_params,
+        lr=args.adam_lr * args.muon_lr_factor,
+        weight_decay=args.weight_decay,
+        momentum=0.95,
+        nesterov=True,
+        adjust_lr_fn="match_rms_adamw",
+    )
+    optimizers = [optimizer_adam, optimizer_muon]
+
+    def lr_lambda(step_idx):
+        return make_trapezoidal_lr(
+            step_idx, schedule_steps, args.warmup_ratio, args.warmdown_ratio
+        )
+
+    schedulers = [
+        torch.optim.lr_scheduler.LambdaLR(optimizer_adam, lr_lambda=lr_lambda),
+        torch.optim.lr_scheduler.LambdaLR(optimizer_muon, lr_lambda=lr_lambda),
+    ]
+
+    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
+    checkpoint_payload = None
+    if args.resume:
+        checkpoint_payload = _load_checkpoint(
+            args.out,
+            model,
+            optimizers,
+            schedulers,
+            scaler,
+            device,
+        )
+        if checkpoint_payload is None:
+            print(f"Resume requested but checkpoint not found at {args.out}; starting fresh.")
+    state = _load_train_state(args, checkpoint_payload)
+
+    if args.compile and device.type == "cuda":
+        model = torch.compile(model, dynamic=False)
+
+    stream = ExpertStream(
+        idx_path=args.idx,
+        tokens_path=args.tokens,
+        doc_starts_path=args.doc_starts,
+        dataset_name=args.dataset,
+        dataset_revision=args.dataset_revision,
+        tokenizer=tokenizer,
+        seq_len=args.seq_len,
+        min_seq_len=args.min_seq_len,
+        stride=args.stride,
+        max_tokens=args.max_tokens,
+        batch_size=args.batch_size,
+        state=state,
+    )
+
+    model.train()
+    step = state.step
+    micro_step = 0
+    start_time = time.time()
+    for opt in optimizers:
+        opt.zero_grad(set_to_none=True)
+
+    for batch in stream:
+        if step >= args.steps:
+            break
+
+        micro_step += 1
+        input_ids = batch["input_ids"].to(device, non_blocking=True)
+        expert_idx = batch["expert_idx"][:, :, : args.layers].to(device, non_blocking=True)
+        attention_mask = batch["attention_mask"].to(device, non_blocking=True)
+
+        labels = input_ids.clone()
+        labels[~attention_mask] = -100
+
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+            logits = model(expert_idx, attention_mask)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+                ignore_index=-100,
+                label_smoothing=args.label_smoothing,
+            )
+            scaled_loss = loss / args.grad_accum
+
+        scaler.scale(scaled_loss).backward()
+
+        if micro_step % args.grad_accum != 0:
+            continue
+
+        if args.clip_grad_norm > 0:
+            for opt in optimizers:
+                scaler.unscale_(opt)
+            params = [
+                p
+                for opt in optimizers
+                for group in opt.param_groups
+                for p in group["params"]
+                if p.grad is not None
+            ]
+            torch.nn.utils.clip_grad_norm_(params, args.clip_grad_norm)
+
+        for opt in optimizers:
+            scaler.step(opt)
+        scaler.update()
+        for opt in optimizers:
+            opt.zero_grad(set_to_none=True)
+        for sched in schedulers:
+            sched.step()
+
+        step += 1
+        state = batch["state"]
+        state.step = step
+        micro_step = 0
+
+        if step % args.log_every == 0:
+            elapsed = time.time() - start_time
+            lr_adam = schedulers[0].get_last_lr()[0]
+            lr_muon = schedulers[1].get_last_lr()[0]
+            loss_value = float(loss.item())
+            print(
+                f"step {step} loss {loss_value:.4f} lr_adam {lr_adam:.6e} "
+                f"lr_muon {lr_muon:.6e}"
+            )
+            if wandb_run:
+                wandb_run.log(
+                    {
+                        "train/loss": loss_value,
+                        "train/lr_adam": lr_adam,
+                        "train/lr_muon": lr_muon,
+                        "train/step": step,
+                        "train/time_elapsed_s": elapsed,
+                    },
+                    step=step,
+                )
+
+        if step % args.save_every == 0:
+            _save_checkpoint(
+                args.out, model, args, state, step, optimizers, schedulers, scaler
+            )
+            _write_json_atomic(args.state_path, state.__dict__)
+
+    _save_checkpoint(args.out, model, args, state, step, optimizers, schedulers, scaler)
+    _write_json_atomic(args.state_path, state.__dict__)
+
+    if wandb_run:
+        wandb_run.finish()
+
+
+if __name__ == "__main__":
+    main()