train_multi.py

from torch.utils.checkpoint import checkpoint
import os, time, torch, torch.nn as nn, torch.optim as optim, torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from dataset import get_or_create_tokenizer, PackedChineseDataset
from srcn_model import SRCNv3_1B

def train():
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    device = torch.device(f"cuda:{local_rank}")

    C_default, M_default, K_default = 160, 384, 8
    C = int(os.environ.get("SRCN_C", str(C_default)))
    M = int(os.environ.get("SRCN_M", str(M_default)))
    K = int(os.environ.get("SRCN_K", str(K_default)))
    B = int(os.environ.get("SRCN_B", "64"))
    seq_len = 512
    bptt_steps = 32
    pool_steps = 4
    lr = 3e-4
    lr_enc = 3e-4
    lr_w = 5e-5
    lr_w_wd = 5e-4
    grad_clip = 0.3
    save_interval = 1800  # 30 min

    script_dir = os.path.dirname(os.path.abspath(__file__))
    corpus = os.path.join(script_dir, "annotated_corpus.jsonl")
    tokenizer = get_or_create_tokenizer(corpus, os.path.join(script_dir, "vocab_tokenizer_v3.pkl"))
    V_size = tokenizer.vocab_size

    dataset = PackedChineseDataset(corpus, tokenizer, chunk_len=seq_len, cache_path=os.path.join(script_dir, "packed_dataset_340m.pkl"))
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=B, sampler=sampler, drop_last=True, num_workers=2, pin_memory=True)
    if rank == 0:
        print(f"World: {world_size} GPUs | B/GPU: {B} | Eff B: {B*world_size} | bptt: {bptt_steps}")
        print(f"Dataset: {len(dataset)} chunks | Batches/epoch: {len(loader)}")

    model = SRCNv3_1B(vocab_size=V_size, num_columns=C, neurons_per_column=M, num_partners=K, num_motor_pool_steps=pool_steps).to(device)
    model = DDP(model, device_ids=[local_rank], find_unused_parameters=False)
    total = sum(p.numel() for p in model.module.parameters())
    if rank == 0:
        print(f"Params: {total:,} ({total/1e9:.3f}B) | C={C} M={M} K={K}")

    motor_start = model.module.motor_start_col
    num_motor = model.module.num_motor_neurons

    # 3 param groups: encoder (high LR), W_raw (slow), MLP (no wd)
    enc_params = []
    w_raw_params = []
    mlp_params = []
    for name, param in model.named_parameters():
        if 'W_raw' in name:
            w_raw_params.append(param)
        elif 'vocab_head' in name:
            mlp_params.append(param)
        else:
            enc_params.append(param)
    opt = optim.AdamW([
        {'params': enc_params, 'lr': lr_enc, 'weight_decay': 1e-4},
        {'params': w_raw_params, 'lr': lr_w, 'weight_decay': lr_w_wd},
        {'params': mlp_params, 'lr': lr, 'weight_decay': 0.0},
    ])
    if rank == 0:
        print(f"LR: enc={lr_enc}, mlp={lr}, W_raw={lr_w}(wd={lr_w_wd})")

    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_id)

    V_th_persist = None

    ckpt_path = os.path.join(script_dir, "checkpoint.pt")
    start_epoch = 0
    start_batch = 0
    if os.path.exists(ckpt_path):
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
        model.module.load_state_dict(ckpt["model"])
        opt.load_state_dict(ckpt["optimizer"])
        start_epoch = ckpt["epoch"]
        start_batch = ckpt.get("batch_idx", 0)
        if "V_th_persist" in ckpt:
            V_th_persist = ckpt["V_th_persist"].to(device)
        if rank == 0:
            print(f"Resumed from epoch {start_epoch+1} batch {start_batch}")

    if rank == 0:
        print("Starting training...\n")

    torch.cuda.synchronize()
    t0 = time.time()
    last_save_time = time.time()
    total_tokens = 0
    num_epochs = 20

    for epoch in range(start_epoch, num_epochs):
        sampler.set_epoch(epoch)
        epoch_loss = 0.0
        n_batches = 0

        for batch_idx, batch_x in enumerate(loader):
            # Skip batches already processed before checkpoint
            if batch_idx < start_batch:
                continue
            start_batch = 0  # reset after first epoch
            batch_x = batch_x.to(device)
            # Precompute W once per batch (saves 377MB MP16 recompute per timestep)
            W_fp16 = model.module.precompute_W()
            S = torch.zeros(B, C, M, device=device)
            V = torch.zeros(B, C, M, device=device)
            if V_th_persist is None or V_th_persist.shape[0] != B:
                V_th_persist = torch.full((B, C, M), 2.0, device=device)
            V_th = V_th_persist.detach()
            I_ampa = torch.zeros(B, C, M, device=device)
            I_nmda = torch.zeros(B, C, M, device=device)
            I_psc = torch.zeros(B, num_motor, device=device)
            window_loss = 0.0
            n_windows = 0
            total_spikes = 0.0
            total_steps = 0

            t_ranges = list(range(0, seq_len - 1, bptt_steps))
            for wi, t_start in enumerate(t_ranges):
                t_end = min(t_start + bptt_steps, seq_len - 1)
                win_tokens = batch_x[:, t_start:t_end]
                target = batch_x[:, t_start + 1:t_end + 1]
                opt.zero_grad(set_to_none=True)

                W_fp16 = model.module.precompute_W()

                # Flat iteration: all tokens in window processed sequentially
                pooled_list = []
                num_win_tokens = win_tokens.shape[1]
                for t in range(num_win_tokens):
                    token = win_tokens[:, t]
                    ts_start = t_start + t * pool_steps
                    ts_start_tensor = torch.tensor(ts_start, device=device)
                    
                    S, V, V_th, I_ampa, I_nmda, I_psc, pooled_token, spikes_sum = checkpoint(
                        model.module.forward_token_with_psc,
                        S, V, V_th, I_ampa, I_nmda, I_psc, token, ts_start_tensor, W_fp16,
                        use_reentrant=True
                    )
                    total_spikes += spikes_sum.item()
                    total_steps += pool_steps
                    pooled_list.append(pooled_token)

                window_pooled = torch.stack(pooled_list, dim=1)
                logits = model.module.vocab_head(window_pooled)
                loss = criterion(logits.view(-1, V_size), target.reshape(-1))

                if not torch.isfinite(loss):
                    # Reset all state to zero — detach cuts old graph
                    S = torch.zeros(B, C, M, device=device)
                    V = torch.zeros(B, C, M, device=device)
                    I_ampa = torch.zeros(B, C, M, device=device)
                    I_nmda = torch.zeros(B, C, M, device=device)
                    I_psc = torch.zeros(B, num_motor, device=device)
                    # If V_th_persist is NaN, reinit from scratch
                    if not torch.isfinite(V_th_persist).all():
                        V_th_persist = torch.full((B, C, M), 2.0, device=device)
                    V_th = V_th_persist.detach()
                    # DO NOT backward — NaN grads poison DDP bucket
                    # Old graph freed when loss is overwritten next iteration
                    opt.zero_grad(set_to_none=True)
                    window_loss += float('nan')
                    n_windows += 1
                    if rank == 0:
                        t = time.time() - t0
                        print(f"  [!] NaN/Inf at E{epoch+1}B{batch_idx+1:04d} W{wi} | VRAM: {torch.cuda.memory_allocated()/1e9:.2f}GB | time={t:.0f}s")
                    continue

                if wi < len(t_ranges) - 1:
                    with model.no_sync():
                        loss.backward()
                else:
                    loss.backward()

                # Check for NaN/Inf in gradients to prevent parameter poisoning
                grad_ok = True
                for name, param in model.named_parameters():
                    if param.grad is not None:
                        if not torch.isfinite(param.grad).all():
                            grad_ok = False
                            if rank == 0:
                                print(f"  [!] Infinite/NaN gradient in {name} at E{epoch+1}B{batch_idx+1:04d} W{wi}")
                            break

                if not grad_ok:
                    opt.zero_grad(set_to_none=True)
                    S = torch.zeros(B, C, M, device=device)
                    V = torch.zeros(B, C, M, device=device)
                    I_ampa = torch.zeros(B, C, M, device=device)
                    I_nmda = torch.zeros(B, C, M, device=device)
                    I_psc = torch.zeros(B, num_motor, device=device)
                    V_th = V_th_persist.detach()
                    continue

                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
                opt.step()
                S, V, V_th, I_ampa, I_nmda, I_psc = [x.detach() for x in [S, V, V_th, I_ampa, I_nmda, I_psc]]
                window_loss += loss.item()
                n_windows += 1

            avg = window_loss / max(n_windows, 1)
            epoch_loss += avg
            n_batches += 1
            total_tokens += B * seq_len * world_size
            elapsed = time.time() - t0
            V_th_persist = V_th.detach()
            cur_mem = torch.cuda.memory_allocated() / 1e9
            max_mem = torch.cuda.max_memory_allocated() / 1e9
            if rank == 0:
                spike_rate = total_spikes / max(total_steps * B * C * M, 1) if total_steps > 0 else 0.0
                print(f"E{epoch+1}B{batch_idx+1:04d} | Loss: {avg:.6f} | Tok/s: {total_tokens/elapsed:.0f} | VRAM: {cur_mem:.2f}/{max_mem:.2f}GB | SR: {spike_rate:.4f}")

            # Periodic checkpoint
            if rank == 0 and time.time() - last_save_time > save_interval:
                torch.save({"model": model.module.state_dict(), "optimizer": opt.state_dict(),
                            "epoch": epoch, "batch_idx": batch_idx + 1,
                            "V_th_persist": V_th_persist.detach().cpu()}, ckpt_path)
                last_save_time = time.time()
                print(f"  [Checkpoint saved at {time.time()-t0:.0f}s] E{epoch+1}B{batch_idx+1}")

            # Defragment CUDA allocator every 100 batches
            if batch_idx > 0 and batch_idx % 100 == 0:
                torch.cuda.empty_cache()

        avg_epoch = epoch_loss / max(n_batches, 1)
        if rank == 0:
            print(f"\n=== Epoch {epoch+1} done | Avg loss: {avg_epoch:.4f} | Elapsed: {time.time()-t0:.0f}s ===\n")
            torch.save({"model": model.module.state_dict(), "optimizer": opt.state_dict(),
                        "epoch": epoch, "batch_idx": 0,
                        "V_th_persist": V_th_persist.detach().cpu()}, ckpt_path)
            last_save_time = time.time()
        start_batch = 0  # Ensure start_batch is reset if the previous epoch was completed early or skipped

    dist.destroy_process_group()
    if rank == 0:
        print(f"Done in {time.time()-t0:.0f}s | Total tokens: {total_tokens:,}")

if __name__ == "__main__":
    train()