train_adam_finewebedu.py

import importlib
import math
import os
import pickle
import random
import time
from contextlib import nullcontext
from datetime import datetime

import numpy as np
import torch
import torch.distributed as dist
from torch.distributed import destroy_process_group, init_process_group
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import AdamW

from optim_utils import get_optimizer_param_groups
from reproducibility import fold_in_seed
from seed_utils import seed_everything

# -----------------------------------------------------------------------------
# default config values designed to train a gpt2 (124M) on Fineweb-edu 100B
# I/O
data_path = "data"
out_dir = 'output/out'
resume_dir = '.'
eval_interval = 1000
log_interval = 1
eval_iters = 2000
eval_only = False  # if True, script exits right after the first eval
always_save_checkpoint = True  # if True, always save a checkpoint after each eval
init_from = 'scratch'  # 'scratch' or 'resume' or 'gpt2*'
# init_from = 'resume'
# wandb logging
wandb_log = False  # disabled by default
wandb_project = 'nanogpt-next'
wandb_run_name = 'gpt2'  # 'run' + str(time.time())
# data
dataset = 'fineweb-edu100B'
gradient_accumulation_steps = 3  # used to simulate larger batch sizes
batch_size = 20  # if gradient_accumulation_steps > 1, this is the micro-batch size
block_size = 4096
# reproducibility
seed = 42
deterministic = False  # if True, enable deterministic (may reduce throughput)
bitwise_deterministic = False  # if True, best-effort bitwise-identical training + exact resume (forces deterministic, disables compile/TF32)
data_seed = -1  # if <0, defaults to seed
eval_seed = -1  # if <0, defaults to seed
data_rng_mode = 'stateless'  # 'stateful' (Generator+Random) or 'stateless' (seed from (rank, step))
# model
n_layer = 24
n_head = 8
n_embd = 1024
head_dim = 128
tpa_kvrank = 2
tpa_qrank = 16
dropout = 0.0  # for pretraining 0 is good, for finetuning try 0.1+
bias = False  # do we use bias inside LayerNorm and Linear layers?
using_groupnorm = False
# KV shifting (optional)
use_k_shift = False
use_v_shift = False
# optimizer
optimizer_name = 'adamw'
learning_rate = 1e-3  # max learning rate
max_iters = 25000  # total number of training iterations
weight_decay = 0.01
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0  # clip gradients at this value, or disable if == 0.0
rho = 0.1
interval = 10
variant = 4
# learning rate decay settings
decay_lr = True  # whether to decay the learning rate
warmup_iters = 2000  # how many steps to warm up for
lr_decay_iters = 25000  # should be ~= max_iters per Chinchilla
min_lr = 3e-5  # minimum learning rate, should be ~= learning_rate/10 per Chinchilla
# DDP settings
backend = 'nccl'  # 'nccl', 'gloo', etc.
schedule = 'cosine'
model_type = '__base_model_placeholder__'
group_size = 11
# system
device = 'cuda'  # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
dtype = 'bfloat16'  # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
compile = True  # use PyTorch 2.0 to compile the model to be faster
scale_attn_by_inverse_layer_idx = True
# -----------------------------------------------------------------------------
config_keys = [
    k
    for k, v in globals().items()
    if not k.startswith('_') and isinstance(v, (int, float, bool, str))
]
exec(open('configurator.py').read())  # overrides from command line or config file
if data_seed < 0:
    data_seed = seed
if eval_seed < 0:
    eval_seed = seed
if bitwise_deterministic:
    deterministic = True
    compile = False
# Some CUDA determinism also requires setting CUBLAS_WORKSPACE_CONFIG before CUDA context init.
if deterministic:
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":16:8")
config = {k: globals()[k] for k in config_keys}  # will be useful for logging
# -----------------------------------------------------------------------------
model_file = importlib.import_module(f'model.{model_type}')
GPTConfig = model_file.GPTConfig
GPT = model_file.GPT


def get_num_params(self, non_embedding=False):
    """
    Return the number of parameters in the model.
    For non-embedding count (default), the position embeddings get subtracted.
    The token embeddings would too, except due to the parameter sharing these
    params are actually used as weights in the final layer, so we include them.
    """
    n_params = sum(p.numel() for p in self.parameters())
    if non_embedding:
        n_params -= self.transformer.wpe.weight.numel()
    return n_params


# Get current date and job ID
current_date = datetime.now().strftime("%Y%m%d_%H%M%S")
job_id = os.environ.get('SLURM_JOB_ID', '0')

# various inits, derived attributes, I/O setup
ddp = int(os.environ.get('RANK', -1)) != -1  # is this a ddp run?
if ddp:
    print(
        f"WORLD_SIZE: {os.environ.get('WORLD_SIZE')}, "
        f"RANK: {os.environ.get('RANK')}, "
        f"LOCAL_RANK: {os.environ.get('LOCAL_RANK')}"
    )
    init_process_group(backend=backend)
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    master_process = ddp_rank == 0  # this process will do logging, checkpointing etc.
    seed_offset = ddp_rank  # each process gets a different seed
else:
    # if not ddp, we are running on a single gpu, and one process
    master_process = True
    seed_offset = 0
    ddp_rank = 0
    world_size = 1
    gradient_accumulation_steps *= 8  # simulate 8 gpus

# Calculate total tokens in billions
tokens_per_iter = batch_size * block_size * gradient_accumulation_steps * world_size
total_tokens_B = tokens_per_iter * max_iters / (1000 ** 3)

# Add after the initial variable declarations
tokens_trained = 0  # track total tokens trained

# Initialize random seed and torch settings
seed_everything(seed + seed_offset, deterministic=deterministic)
torch.backends.cuda.matmul.allow_tf32 = not bitwise_deterministic  # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = not bitwise_deterministic  # allow tf32 on cudnn
if bitwise_deterministic:
    torch.set_float32_matmul_precision('highest')
    # Force deterministic SDPA kernel selection (slower, but stable).
    if torch.cuda.is_available():
        torch.backends.cuda.enable_flash_sdp(False)
        torch.backends.cuda.enable_mem_efficient_sdp(False)
        torch.backends.cuda.enable_cudnn_sdp(False)
        torch.backends.cuda.enable_math_sdp(True)
device_type = 'cuda' if 'cuda' in device else 'cpu'  # for later use in torch.autocast
# Note: float16 data type will automatically use a GradScaler
ptdtype = {
    'float32': torch.float32,
    'bfloat16': torch.bfloat16,
    'float16': torch.float16,
}[dtype]
if device_type == 'cpu':
    ctx = nullcontext()
else:
    ctx = torch.autocast(device_type=device_type, dtype=ptdtype)

# Poor man's data loader
data_dir = os.path.join(data_path, dataset)
# train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
train_file_list = sorted(list([
    file_name
    for file_name in os.listdir(data_dir)
    if file_name.endswith('.bin') and file_name.startswith('fineweb_train')
]))
train_data_list = [
    np.memmap(os.path.join(data_dir, file_name), dtype=np.uint16, mode='r')
    for file_name in train_file_list
]
val_data = np.memmap(os.path.join(data_dir, 'fineweb_val_000000.bin'), dtype=np.uint16, mode='r')
train_data_rng = torch.Generator(device='cpu')
train_data_rng.manual_seed(data_seed + ddp_rank)
train_py_rng = random.Random(data_seed + ddp_rank)
eval_data_rng = torch.Generator(device='cpu')
eval_data_rng.manual_seed(eval_seed)
eval_py_rng = random.Random(eval_seed)
stateless_train_rng = torch.Generator(device='cpu')
stateless_eval_rng = torch.Generator(device='cpu')


def get_batch(
    split,
    *,
    rng: torch.Generator | None = None,
    py_rng: random.Random | None = None,
    batch_id: int | None = None,
    base_seed: int | None = None,
    rank: int | None = None,
):
    if data_rng_mode == 'stateful':
        if rng is None:
            raise ValueError("stateful data_rng_mode requires rng=")
        batch_rng = rng
        if split == 'train':
            if py_rng is None:
                raise ValueError("stateful data_rng_mode requires py_rng= for split='train'")
            data = py_rng.choice(train_data_list)
        else:
            data = val_data
    elif data_rng_mode == 'stateless':
        if batch_id is None:
            raise ValueError("stateless data_rng_mode requires batch_id=")
        base = base_seed if base_seed is not None else (data_seed if split == 'train' else eval_seed)
        eff_rank = rank if rank is not None else (ddp_rank if split == 'train' else 0)
        batch_rng = stateless_train_rng if (base_seed is None and rank is None and split == 'train') else stateless_eval_rng
        batch_rng.manual_seed(fold_in_seed(base, eff_rank, batch_id))
        if split == 'train':
            file_idx = torch.randint(len(train_data_list), (1,), generator=batch_rng).item()
            data = train_data_list[file_idx]
        else:
            data = val_data
    else:
        raise ValueError(f"Unknown data_rng_mode={data_rng_mode!r}")

    offset = 512
    ix = torch.randint(
        len(data) - block_size - offset, (batch_size,), generator=batch_rng
    )
    x = torch.stack(
        [
            torch.from_numpy(
                data[offset + i : offset + i + block_size].astype(np.int64)
            )
            for i in ix
        ]
    )
    y = torch.stack(
        [
            torch.from_numpy(
                data[offset + i + 1 : offset + i + 1 + block_size].astype(np.int64)
            )
            for i in ix
        ]
    )
    if device_type == 'cuda':
        # pin arrays x, y, which allows us to move them to GPU asynchronously (non_blocking=True)
        x, y = (
            x.pin_memory().to(device, non_blocking=True),
            y.pin_memory().to(device, non_blocking=True),
        )
    else:
        x, y = x.to(device), y.to(device)
    return x, y


train_data_rng_state_for_batch: torch.ByteTensor | None = None
train_py_rng_state_for_batch: object | None = None


def _data_rng_state_filename() -> str:
    return f"data_rng_state_rank{ddp_rank}.pt" if ddp else "data_rng_state.pt"


def _save_data_rng_state(ckpt_dir: str) -> None:
    if data_rng_mode != 'stateful':
        return
    if train_data_rng_state_for_batch is None or train_py_rng_state_for_batch is None:
        raise RuntimeError("Missing train RNG state for stateful RNG mode")
    payload = {
        'iter_num': iter_num,
        'world_size': world_size,
        'ddp_rank': ddp_rank,
        'train_data_rng_state_for_batch': train_data_rng_state_for_batch,
        'train_py_rng_state_for_batch': train_py_rng_state_for_batch,
    }
    torch.save(payload, os.path.join(ckpt_dir, _data_rng_state_filename()))


def _get_train_batch_stateful():
    global train_data_rng_state_for_batch, train_py_rng_state_for_batch
    train_data_rng_state_for_batch = train_data_rng.get_state()
    train_py_rng_state_for_batch = train_py_rng.getstate()
    return get_batch('train', rng=train_data_rng, py_rng=train_py_rng)

def _global_rng_state_filename() -> str:
    return f"global_rng_state_rank{ddp_rank}.pt" if ddp else "global_rng_state.pt"


def _save_global_rng_state(ckpt_dir: str) -> None:
    if not bitwise_deterministic:
        return
    payload = {
        'iter_num': iter_num,
        'world_size': world_size,
        'ddp_rank': ddp_rank,
        'python_random_state': random.getstate(),
        'numpy_random_state': np.random.get_state(),
        'torch_rng_state': torch.get_rng_state(),
    }
    if device_type == 'cuda':
        payload['cuda_rng_state'] = torch.cuda.get_rng_state(device)
    torch.save(payload, os.path.join(ckpt_dir, _global_rng_state_filename()))


# Init these up here, can override if init_from='resume' (i.e. from a checkpoint)
iter_num = 0
best_val_loss = 1e9

# Attempt to derive vocab_size from the dataset
meta_path = os.path.join(data_dir, 'meta.pkl')
meta_vocab_size = None
if os.path.exists(meta_path):
    with open(meta_path, 'rb') as f:
        meta = pickle.load(f)
    meta_vocab_size = meta['vocab_size']
    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")

# Model initialization
# Build model_args with defaults that can be overridden by config globals
model_args = {
    'n_layer': n_layer,
    'n_head': n_head,
    'n_embd': n_embd,
    'block_size': block_size,
    'bias': bias,
    'head_dim': head_dim,
    'tpa_kvrank': tpa_kvrank,
    'tpa_qrank': tpa_qrank,
    'using_groupnorm': using_groupnorm,
    'vocab_size': None,
    'dropout': dropout,
    'scale_attn_by_inverse_layer_idx': scale_attn_by_inverse_layer_idx,
    # Init/normalization knobs (with sensible defaults)
    'embedding_init_std': globals().get('embedding_init_std', 0.02),
    'hidden_init_std_factor': globals().get('hidden_init_std_factor', 0.5),
    'use_qk_rmsnorm': globals().get('use_qk_rmsnorm', True),
    'use_k_shift': globals().get('use_k_shift', False),
    'use_v_shift': globals().get('use_v_shift', False),
    'p_tie_mode': globals().get('p_tie_mode', 'none'),
    'p_head_dim': globals().get('p_head_dim', None),
}

# Fox-specific toggles (only relevant when using model_type == 'fox')
if model_type == 'fox':
    model_args.update({
        'fgate_type': globals().get('fgate_type', 'full'),
    })

if "gqa" in model_type:
    model_args['group_size'] = group_size

# Pass through any GRAPE-specific hyperparameters provided via config files.
if 'grape' in model_type:
    for key, value in list(globals().items()):
        if key.startswith('grape_'):
            model_args[key] = value

# Pass through any key-gated additive bias hyperparameters provided via config files.
if 'keygated' in model_type:
    for key, value in list(globals().items()):
        if key.startswith('keygated_'):
            model_args[key] = value

# Pass through any query-gated additive bias hyperparameters provided via config files.
if 'querygated' in model_type or 'querykeygated' in model_type:
    for key, value in list(globals().items()):
        if key.startswith('querygated_'):
            model_args[key] = value

# Pass through any query+key-gated additive bias hyperparameters provided via config files.
if 'querykeygated' in model_type:
    for key, value in list(globals().items()):
        if key.startswith('querykeygated_'):
            model_args[key] = value

if init_from == 'scratch':
    # Init a new model from scratch
    print("Initializing a new model from scratch")
    # Determine the vocab size we'll use for from-scratch training
    if meta_vocab_size is None:
        print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)")
    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
    gptconf = GPTConfig(**model_args)
    model = GPT(gptconf)
elif init_from == 'resume':
    print(f"Resuming training from {resume_dir}")
    # Resume training from a checkpoint.
    # ckpt_path = os.path.join(out_dir, 'ckpt.pt')
    # checkpoint = torch.load(ckpt_path, map_location=device)
    # checkpoint_model_args = checkpoint['model_args']
    config = GPTConfig.from_json_file(os.path.join(resume_dir, 'config.json'))
    model = GPT.from_pretrained(resume_dir, config=config)
    
    # Force these config attributes to be equal otherwise we can't even resume training
    # The rest of the attributes (e.g. dropout) can stay as desired from command line
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = getattr(config, k)
    model.transformer.wte.weight = model.lm_head.weight
elif init_from.startswith('gpt2'):
    print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
    # Initialize from OpenAI GPT-2 weights
    override_args = {'dropout': dropout}
    model = GPT.from_pretrained(init_from, override_args)
    # Read off the created config params, so we can store them into checkpoint correctly
    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
        model_args[k] = getattr(model.config, k)
# Crop down the model block size if desired, using model surgery
if block_size < model.config.block_size:
    model.crop_block_size(block_size)
    model_args['block_size'] = block_size  # so that the checkpoint will have the right value
model.to(device)

# Now calculate non-embedding parameters
param_count = get_num_params(model, non_embedding=False)
param_count_m = param_count / 1_000_000  # convert to millions

# Update wandb run name and out_dir if not resuming
if init_from != 'resume':
    # Update wandb run name
    wandb_run_name = f"M_{model_type}_{int(param_count_m)}m_Opt_{optimizer_name}_LR_{learning_rate}_D_{dataset}_T_{total_tokens_B:.2f}B_time_{current_date}_jobid_{job_id}"
    # Update output directory
    out_dir = f"output/out_{model_type}_{int(param_count_m)}m_Opt_{optimizer_name}_LR_{learning_rate}_D_{dataset}_T_{total_tokens_B:.2f}B_time_{current_date}_jobid_{job_id}"
else:
    try:
        resume_job_id = resume_dir.split('time_')[1].split('/')[0]
        wandb_run_name = f"M_{model_type}_{int(param_count_m)}m_Opt_{optimizer_name}_LR_{learning_rate}_D_{dataset}_T_{total_tokens_B:.2f}B_time_{resume_job_id}"
        out_dir = f"output/out_{model_type}_{int(param_count_m)}m_Opt_{optimizer_name}_LR_{learning_rate}_D_{dataset}_T_{total_tokens_B:.2f}B_time_{resume_job_id}"
    except Exception:
        pass
# Now create the output directory
if master_process:
    os.makedirs(out_dir, exist_ok=True)
if ddp:
    dist.barrier()

# Initialize a GradScaler. If enabled=False, scaler is a no-op
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))

# Optimizer
param_groups = get_optimizer_param_groups(model, weight_decay)
optimizer = AdamW(
    param_groups,
    lr=learning_rate,
    betas=(beta1, beta2),
    eps=1e-8,
)
if init_from == 'resume':
    optimizer_state_path = os.path.join(resume_dir, 'optimizer.pt')
    optimizer_state = torch.load(optimizer_state_path, map_location=device)
    # param_num = optimizer_state['optimizer']['param_groups'][0]['params'][-1]+1
    # optimizer_state['optimizer']['param_groups'][0]['params'].append(param_num)
    # optimizer_state['optimizer']['state'][param_num] = optimizer_state['optimizer']['state'][0]
    optimizer.load_state_dict(optimizer_state['optimizer'])
    iter_num = optimizer_state['iter_num']
    best_val_loss = optimizer_state['best_val_loss']
    if bitwise_deterministic:
        scaler_state = optimizer_state.get('scaler')
        if scaler_state is None:
            raise RuntimeError(
                f"bitwise_deterministic=True requires scaler state in {optimizer_state_path}"
            )
        scaler.load_state_dict(scaler_state)
    # print(param_num)
    # print(optimizer.state_dict()['param_groups'])
    print(best_val_loss)
    del optimizer_state
    if data_rng_mode == 'stateful':
        rng_state_path = os.path.join(resume_dir, _data_rng_state_filename())
        if os.path.exists(rng_state_path):
            rng_state = torch.load(rng_state_path, map_location='cpu')
            # Sanity-check that we are resuming with the intended DDP topology and step.
            ckpt_iter = rng_state.get('iter_num')
            ckpt_world_size = rng_state.get('world_size')
            ckpt_rank = rng_state.get('ddp_rank')
            if ckpt_iter is not None and int(ckpt_iter) != int(iter_num):
                raise RuntimeError(
                    f"RNG state iter mismatch: rng_state iter_num={ckpt_iter} "
                    f"but optimizer iter_num={iter_num} ({rng_state_path})"
                )
            if ckpt_world_size is not None and int(ckpt_world_size) != int(world_size):
                raise RuntimeError(
                    f"RNG state WORLD_SIZE mismatch: rng_state world_size={ckpt_world_size} "
                    f"but current WORLD_SIZE={world_size} ({rng_state_path})"
                )
            if ckpt_rank is not None and int(ckpt_rank) != int(ddp_rank):
                raise RuntimeError(
                    f"RNG state RANK mismatch: rng_state ddp_rank={ckpt_rank} "
                    f"but current RANK={ddp_rank} ({rng_state_path})"
                )
            torch_state = rng_state.get('train_data_rng_state_for_batch')
            py_state = rng_state.get('train_py_rng_state_for_batch')
            if torch_state is None or py_state is None:
                raise RuntimeError(f"Missing train RNG state in {rng_state_path}")
            train_data_rng_state_for_batch = torch_state
            train_py_rng_state_for_batch = py_state
            train_data_rng.set_state(torch_state)
            train_py_rng.setstate(py_state)
        else:
            raise FileNotFoundError(
                f"stateful data_rng_mode requires {rng_state_path} to resume deterministically"
            )
    if bitwise_deterministic:
        rng_state_path = os.path.join(resume_dir, _global_rng_state_filename())
        if os.path.exists(rng_state_path):
            rng_state = torch.load(rng_state_path, map_location='cpu')
            ckpt_iter = rng_state.get('iter_num')
            ckpt_world_size = rng_state.get('world_size')
            ckpt_rank = rng_state.get('ddp_rank')
            if ckpt_iter is not None and int(ckpt_iter) != int(iter_num):
                raise RuntimeError(
                    f"Global RNG state iter mismatch: rng_state iter_num={ckpt_iter} "
                    f"but optimizer iter_num={iter_num} ({rng_state_path})"
                )
            if ckpt_world_size is not None and int(ckpt_world_size) != int(world_size):
                raise RuntimeError(
                    f"Global RNG state WORLD_SIZE mismatch: rng_state world_size={ckpt_world_size} "
                    f"but current WORLD_SIZE={world_size} ({rng_state_path})"
                )
            if ckpt_rank is not None and int(ckpt_rank) != int(ddp_rank):
                raise RuntimeError(
                    f"Global RNG state RANK mismatch: rng_state ddp_rank={ckpt_rank} "
                    f"but current RANK={ddp_rank} ({rng_state_path})"
                )
            random.setstate(rng_state['python_random_state'])
            np.random.set_state(rng_state['numpy_random_state'])
            torch.set_rng_state(rng_state['torch_rng_state'])
            if device_type == 'cuda':
                cuda_state = rng_state.get('cuda_rng_state')
                if cuda_state is None:
                    raise RuntimeError(f"Missing cuda_rng_state in {rng_state_path}")
                torch.cuda.set_rng_state(cuda_state, device)
        else:
            raise FileNotFoundError(
                f"bitwise_deterministic=True requires {rng_state_path} to resume deterministically"
            )
# Compile the model
if compile:
    print("compiling the model... (takes a ~minute)")
    unoptimized_model = model
    model = torch.compile(model)  # requires PyTorch 2.0

# Wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])


# Helps estimate an arbitrarily accurate loss over either split using many batches
@torch.no_grad()
def estimate_loss():
    out = {}
    eval_model = model.module if ddp else model
    eval_model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            if data_rng_mode == 'stateless':
                X, Y = get_batch(split, batch_id=k, base_seed=eval_seed, rank=0)
            else:
                X, Y = get_batch(split, rng=eval_data_rng, py_rng=eval_py_rng)
            with ctx:
                logits, loss = eval_model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    eval_model.train()
    return out


# Learning rate decay scheduler (cosine with warmup)
def get_lr(it, schedule='cosine'):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1

    return min_lr + coeff * (learning_rate - min_lr)

# Logging
if wandb_log and master_process:
    import wandb
    wandb_config = {
        'model_args': model_args,
        'training_args': {
            'batch_size': batch_size,
            'block_size': block_size,
            'gradient_accumulation_steps': gradient_accumulation_steps,
            'max_iters': max_iters,
            'lr_decay_iters': lr_decay_iters,
            'eval_interval': eval_interval,
            'eval_iters': eval_iters,
            'log_interval': log_interval,
        },
        'optimizer_args': {
            'optimizer_name': optimizer_name,
            'learning_rate': learning_rate,
            'weight_decay': weight_decay,
            'beta1': beta1,
            'beta2': beta2,
            'grad_clip': grad_clip,
            'decay_lr': decay_lr,
            'warmup_iters': warmup_iters,
            'min_lr': min_lr,
            'schedule': schedule,
        },
    }
    wandb.init(project=wandb_project, name=wandb_run_name, config=wandb_config)

# Training loop
if data_rng_mode == 'stateless':
    train_batch_id = iter_num * gradient_accumulation_steps
    X, Y = get_batch('train', batch_id=train_batch_id)  # fetch the very first batch
else:
    train_batch_id = None
    X, Y = _get_train_batch_stateful()  # fetch the very first batch
t0 = time.time()
local_iter_num = 0  # number of iterations in the lifetime of this process
raw_model = model.module if ddp else model  # unwrap DDP container if needed
running_mfu = -1.0
clip_time = 0
while True:

    # Determine and set the learning rate for this iteration
    lr = get_lr(iter_num, schedule=schedule) if decay_lr else learning_rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Evaluate the loss on train/val sets and write checkpoints (keep DDP ranks in sync)
    if iter_num % eval_interval == 0:
        save_latest = False
        if master_process:
            losses = estimate_loss()
            print(
                f"step {iter_num}: train loss {losses['train']:.4f}, "
                f"val loss {losses['val']:.4f}"
            )
            if wandb_log:
                wandb.log(
                    {
                        "iter": iter_num,
                        "train/loss": losses['train'],
                        "val/loss": losses['val'],
                        "lr": lr,
                        "mfu": running_mfu * 100,  # convert to percentage
                    },
                    step=iter_num,
                )
            update_best = losses['val'] < best_val_loss or always_save_checkpoint
            if update_best:
                best_val_loss = losses['val']
                if iter_num > 0:
                    print(f"saving checkpoint to {out_dir}")
                    raw_model.save_pretrained(out_dir)
                    optimizer_state = {
                        'optimizer': optimizer.state_dict(),
                        'iter_num': iter_num,
                        'best_val_loss': best_val_loss,
                    }
                    if bitwise_deterministic:
                        optimizer_state['scaler'] = scaler.state_dict()
                    torch.save(optimizer_state, os.path.join(out_dir, 'optimizer.pt'))
                    save_latest = True

        if ddp:
            comm_device = device if backend == 'nccl' else 'cpu'
            save_latest_flag = torch.tensor([1 if save_latest else 0], device=comm_device)
            dist.broadcast(save_latest_flag, src=0)
            save_latest = bool(save_latest_flag.item())

        if save_latest:
            _save_data_rng_state(out_dir)
            _save_global_rng_state(out_dir)
            if ddp:
                dist.barrier()

        save_snapshot = iter_num % (eval_interval * 5) == 0
        if save_snapshot:
            checkpoint_dir = os.path.join(out_dir, f'checkpoint-{iter_num}')
            if master_process:
                os.makedirs(checkpoint_dir, exist_ok=True)
                raw_model.save_pretrained(checkpoint_dir)
                optimizer_state = {
                    'optimizer': optimizer.state_dict(),
                    'iter_num': iter_num,
                    'best_val_loss': best_val_loss,
                }
                if bitwise_deterministic:
                    optimizer_state['scaler'] = scaler.state_dict()
                torch.save(optimizer_state, os.path.join(checkpoint_dir, 'optimizer.pt'))
            if ddp:
                dist.barrier()
            _save_data_rng_state(checkpoint_dir)
            _save_global_rng_state(checkpoint_dir)
            if ddp:
                dist.barrier()
        if ddp:
            dist.barrier()

    if iter_num == 0 and eval_only:
        break

    # Forward backward update, with optional gradient accumulation to simulate larger batch size
    # and using the GradScaler if data type is float16
    for micro_step in range(gradient_accumulation_steps):
        if ddp:
            # In DDP training we only need to sync gradients at the last micro step.
            # The official way to do this is with model.no_sync() context manager, but
            # I really dislike that this bloats the code and forces us to repeat code
            # Looking at the source of that context manager, it just toggles this variable
            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
        with ctx:
            logits, loss = model(X, Y)
        # Immediately async prefetch next batch while model is doing the forward pass on the GPU
        if data_rng_mode == 'stateless':
            train_batch_id += 1
            X, Y = get_batch('train', batch_id=train_batch_id)
        else:
            X, Y = _get_train_batch_stateful()
        # Backward pass, with gradient scaling if training in fp16
        scaler.scale(loss).backward()
    # Clip the gradient
    if grad_clip != 0.0:
        scaler.unscale_(optimizer)
        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        if total_norm.item() > grad_clip:
            clip_time += 1
    # Step the optimizer and scaler if training in fp16
    scaler.step(optimizer)
    scaler.update()
    # Flush the gradients as soon as we can, no need for this memory anymore
    optimizer.zero_grad(set_to_none=True)

    # Timing and logging
    t1 = time.time()
    dt = t1 - t0
    t0 = t1
    # Update total tokens trained
    tokens_trained += tokens_per_iter
    tokens_trained_B = tokens_trained / 1e9  # Convert to billions

    if iter_num % log_interval == 0 and master_process:
        lossf = loss.item()  # loss as float. note: this is a CPU-GPU sync point
        if local_iter_num >= 5:  # let the training loop settle a bit
            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
            running_mfu = mfu if running_mfu == -1.0 else 0.9 * running_mfu + 0.1 * mfu
        tokens_per_sec = tokens_per_iter / dt
        tokens_per_sec_M = tokens_per_sec / 1_000_000
        print(
            f"iter {iter_num}: loss {lossf:.4f}, time {dt * 1000:.2f}ms, "
            f"mfu {running_mfu * 100:.2f}%, tps (M) {tokens_per_sec_M:.2f}, "
            f"tokens trained {tokens_trained_B:.2f}B"
        )

        params = [param for _, param in model.named_parameters()]
        total_param_norm = 0.0
        for param in params:
            param_norm = param.data.norm(2)
            total_param_norm += param_norm.item() ** 2
        total_param_norm = total_param_norm ** 0.5

        optimizer_state = optimizer.state_dict()['state']
        momentum_norm = 0
        momentum_norm_sq = 0
        for state in optimizer_state.values():
            momentum_norm += state['exp_avg'].detach().norm(2) ** 2
            momentum_norm_sq += state['exp_avg_sq'].detach().norm(2) ** 2
        momentum_norm = torch.sqrt(momentum_norm).item()
        momentum_norm_sq = torch.sqrt(momentum_norm_sq).item()
        momentum_div = momentum_norm / (np.sqrt(momentum_norm_sq) + 1e-8)
        if wandb_log:
            wandb.log(
                {
                    "iter": iter_num,
                    "train/loss": lossf,
                    "lr": lr,
                    "param_norm": total_param_norm,
                    "momentum_norm": momentum_norm,
                    "momentum_norm_sq": momentum_norm_sq,
                    "momentum_div": momentum_div,
                    "train/clip_rate": clip_time / (iter_num + 1),
                    "train/grad_norm": total_norm.item() if grad_clip != 0.0 else 0.0,
                    "train/iter_time_ms": dt * 1000,
                    "train/mfu": running_mfu * 100,
                    "train/tokens_per_sec_M": tokens_per_sec_M,
                    "train/tokens_trained_B": tokens_trained_B,
                    "gpu/memory_allocated_MB": torch.cuda.memory_allocated() / (1024 * 1024),
                    "gpu/max_memory_allocated_MB": torch.cuda.max_memory_allocated() / (1024 * 1024),
                },
                step=iter_num,
            )
    iter_num += 1
    local_iter_num += 1

    # Termination conditions
    if iter_num > max_iters:
        break

if ddp:
    destroy_process_group()