Upload training/launcher_v2b_multi.py with huggingface_hub

training/launcher_v2b_multi.py

@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+AETHER-2B-7Attn Phase 1.6 v2b — Continued Pretraining
+=======================================================
+Resume:  step_0455000 (6.59B model, 59.8B tokens done)
+Data:    132B pool (phase15_mix 90B + tokenized_v2 42.5B, weighted)
+Goal:    annealing on high-quality data → 132B chinchilla total
+Infra:   multi node 2x8 B200 (CID=5352+5353, 16 GPU)
+LR:      cosine annealing 5e-5 → 5e-6 over 324K steps (~85B new tokens)
+Log:     /NHNHOME/WORKSPACE/0426030036_A/aether_1b_7attn_v3/logs/launcher_v2b.log
+"""
+
+import math, os, sys, time, signal, json, logging
+from collections import deque
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    MixedPrecision, ShardingStrategy, BackwardPrefetch,
+    StateDictType, ShardedStateDictConfig, ShardedOptimStateDictConfig,
+)
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+# ── paths ────────────────────────────────────────────────────────────────────
+BASE      = "/NHNHOME/WORKSPACE/0426030036_A/aether_1b_7attn_v3"
+V2_ROOT   = "/NHNHOME/WORKSPACE/0426030036_A/aether-recovered/aether-v2-current"
+RESUME_CKPT = f"{BASE}/runs/v2b_phase16/ckpt/step_0001905"
+OUTPUT_DIR  = f"{BASE}/runs/v2b_multi_phase16"
+CKPT_DIR    = f"{OUTPUT_DIR}/ckpt"
+LOG_FILE    = f"{BASE}/logs/launcher_v2b_multi.log"
+KILL_SWITCH = "/NHNHOME/STOP_TRAINING_V2B_MULTI"
+
+sys.path.insert(0, V2_ROOT)
+sys.path.insert(0, V2_ROOT + "/scripts")
+
+from aether_pkg.configuration_aether_v2_7way import AETHERV27wayConfig
+from aether_pkg.modeling_aether_v2_7way import (
+    AETHERV27wayForCausalLM as AETHERForCausalLM,
+    AETHERV27wayDecoderLayer as AETHERDecoderLayer,
+)
+from scripts.mini_config import mini_cfg
+
+# ── hyper params ─────────────────────────────────────────────────────────────
+SEQ_LEN    = 4096
+MICRO_BATCH = 1
+GRAD_ACCUM  = 8
+TOTAL_STEPS = 162000   # ~85B new tokens
+LR_PEAK     = 5e-5
+LR_MIN      = 5e-6
+WARMUP      = 0         # already warmed up at step 1905
+GRAD_CLIP   = 1.0
+SAVE_EVERY  = 5000
+CKPT_KEEP   = 4
+LOG_EVERY   = 50
+EMBED_LR_MULT = 0.1
+LWD_DECAY   = 0.97
+
+# ── data pool ────────────────────────────────────────────────────────────────
+# (path, sampling_weight)  — higher weight = more often sampled
+DATA_POOL = [
+    (f"{BASE}/../aether_phase15/tokenized/phase15_mix_90b.bin",              1.0),
+    (f"{BASE}/../aether_phase15/tokenized_v2/eng_fineweb_edu.bin",            2.0),
+    (f"{BASE}/../aether_phase15/tokenized_v2/syn_cosmopedia.bin",             2.0),
+    (f"{BASE}/../aether_phase15/tokenized_v2/math_finemath.bin",              3.5),
+    (f"{BASE}/../aether_phase15/tokenized_v2/math_owm.bin",                   3.5),
+    (f"{BASE}/../aether_phase15/tokenized_v2/code_opc.bin",                   2.5),
+    (f"{BASE}/../aether_phase15/tokenized_v2/kor_webtext.bin",                2.0),
+    (f"{BASE}/../aether_phase15/tokenized_v2/kor_synth.bin",                  2.0),
+]
+
+# ── NCCL / env ───────────────────────────────────────────────────────────────
+os.environ.setdefault("NCCL_TIMEOUT",                  "3600")
+os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF",        "expandable_segments:True")
+
+# ── distributed init ─────────────────────────────────────────────────────────
+dist.init_process_group(backend="nccl")
+local_rank  = int(os.environ["LOCAL_RANK"])
+rank        = dist.get_rank()
+world_size  = dist.get_world_size()
+torch.cuda.set_device(local_rank)
+device = torch.device(f"cuda:{local_rank}")
+is_main = rank == 0
+
+def barrier():
+    dist.barrier()
+
+# ── logging ──────────────────────────────────────────────────────────────────
+Path(LOG_FILE).parent.mkdir(parents=True, exist_ok=True)
+handlers = [logging.StreamHandler(sys.stdout)]
+if is_main:
+    handlers.append(logging.FileHandler(LOG_FILE, mode="a"))
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(asctime)s][r%(process)d] %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=handlers,
+)
+logger = logging.getLogger(__name__)
+def log(msg):
+    if is_main: logger.info(msg)
+
+# ── signal ───────────────────────────────────────────────────────────────────
+shutdown_req = False
+def _handler(sig, frame):
+    global shutdown_req
+    log(f"[signal] {sig} → graceful shutdown")
+    shutdown_req = True
+signal.signal(signal.SIGTERM, _handler)
+signal.signal(signal.SIGINT,  _handler)
+
+# ═══════════════════════════════════════════════════════════════════════════
+log("=" * 68)
+log("AETHER-2B-7Attn Phase 1.6 v2b — Continued Pretraining")
+log(f"  resume ckpt : {RESUME_CKPT}")
+log(f"  output      : {OUTPUT_DIR}")
+log(f"  log         : {LOG_FILE}")
+log(f"  world_size  : {world_size}")
+log(f"  LR          : {LR_PEAK} → {LR_MIN}  warmup={WARMUP}")
+log(f"  steps       : {TOTAL_STEPS}  (~{TOTAL_STEPS*MICRO_BATCH*GRAD_ACCUM*SEQ_LEN*world_size/1e9:.1f}B tok)")
+log("=" * 68)
+barrier()
+
+# ── [1] build model ───────────────────────────────────────────────────────────
+log("[1/5] Building model (mini_cfg, bf16)...")
+cfg = mini_cfg
+cfg.max_position_embeddings = SEQ_LEN
+cfg.use_cache = False
+t0 = time.time()
+model = AETHERForCausalLM(cfg).to(device=device, dtype=torch.bfloat16)
+n_params = sum(p.numel() for p in model.parameters())
+log(f"  {n_params/1e9:.3f}B params  ({time.time()-t0:.1f}s)")
+barrier()
+
+# ── [2] FSDP wrap ─────────────────────────────────────────────────────────────
+log("[2/5] FSDP wrap...")
+strategy = ShardingStrategy.FULL_SHARD if world_size > 1 else ShardingStrategy.NO_SHARD
+model = FSDP(
+    model,
+    sharding_strategy=strategy,
+    mixed_precision=MixedPrecision(
+        param_dtype=torch.bfloat16,
+        reduce_dtype=torch.bfloat16,
+        buffer_dtype=torch.bfloat16,
+    ),
+    auto_wrap_policy=partial(transformer_auto_wrap_policy,
+                              transformer_layer_cls={AETHERDecoderLayer}),
+    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+    forward_prefetch=True,
+    use_orig_params=True,
+    device_id=local_rank,
+)
+log(f"  strategy={strategy.name}")
+barrier()
+
+# ── [3] optimizer ─────────────────────────────────────────────────────────────
+log("[3/5] Building optimizer (AdamW + LWD)...")
+from collections import defaultdict
+used = set()
+groups = []
+
+# embedding (low LR, no WD)
+emb_p = [(n,p) for n,p in model.named_parameters()
+         if "embed" in n.lower() and p.requires_grad]
+if emb_p:
+    groups.append({"params":[p for _,p in emb_p], "lr": LR_PEAK*EMBED_LR_MULT,
+                   "weight_decay":0.0, "name":"embed"})
+    used.update(id(p) for _,p in emb_p)
+
+# scalars / alphas (no WD)
+sc_p = [(n,p) for n,p in model.named_parameters()
+        if id(p) not in used and p.requires_grad and ("_alpha" in n or p.dim() < 2)]
+if sc_p:
+    groups.append({"params":[p for _,p in sc_p], "lr": LR_PEAK,
+                   "weight_decay":0.0, "name":"scalars"})
+    used.update(id(p) for _,p in sc_p)
+
+# body — layer-wise LR decay
+layer_buckets = defaultdict(list)
+other_p = []
+for n, p in model.named_parameters():
+    if id(p) in used or not p.requires_grad: continue
+    parts = n.split(".")
+    lidx = None
+    for i, part in enumerate(parts):
+        if part == "layers" and i+1 < len(parts) and parts[i+1].isdigit():
+            lidx = int(parts[i+1]); break
+    (layer_buckets[lidx] if lidx is not None else other_p).append(p)
+    used.add(id(p))
+for lidx in sorted(layer_buckets):
+    scale = LWD_DECAY ** (cfg.num_hidden_layers - 1 - lidx)
+    groups.append({"params": layer_buckets[lidx], "lr": LR_PEAK*scale,
+                   "weight_decay":0.1, "name":f"layer_{lidx}"})
+if other_p:
+    groups.append({"params":other_p, "lr":LR_PEAK, "weight_decay":0.1, "name":"misc"})
+
+optimizer = torch.optim.AdamW(groups, betas=(0.9, 0.95), eps=1e-8)
+log(f"  param groups: {len(groups)}")
+barrier()
+
+# ── [4] resume from step_0455000 ─────────────────────────────────────────────
+RESUME_STEP = 1905  # world_size 8->16: model only load
+log(f"[4/5] Loading DCP checkpoint: {RESUME_CKPT}  (model-only, optimizer fresh for 16-GPU)")
+t0 = time.time()
+with FSDP.state_dict_type(
+    model,
+    StateDictType.SHARDED_STATE_DICT,
+    ShardedStateDictConfig(offload_to_cpu=True),
+    ShardedOptimStateDictConfig(offload_to_cpu=True),
+):
+    state = {"model": model.state_dict()}
+    dcp.load(state, storage_reader=FileSystemReader(RESUME_CKPT))
+    model.load_state_dict(state["model"])
+log(f"  loaded in {time.time()-t0:.1f}s  (optimizer fresh init)")
+barrier()
+
+# ── [5] data pool ────────────────────────────────────────────────────────────
+log("[5/5] Loading multi-file data pool...")
+arrs, probs_raw = [], []
+for path_rel, wt in DATA_POOL:
+    path = str(Path(path_rel).resolve())
+    if not Path(path).exists():
+        log(f"  SKIP: {Path(path).name}")
+        continue
+    n_tok = Path(path).stat().st_size // 4
+    arr = np.memmap(path, dtype=np.uint32, mode="r")
+    arrs.append(arr)
+    probs_raw.append(wt * n_tok)
+    log(f"  {Path(path).name}: {n_tok/1e9:.2f}B tok  w={wt}")
+total_w = sum(probs_raw)
+probs = [w/total_w for w in probs_raw]
+log(f"  files={len(arrs)}  effective={total_w/1e9:.1f}B weighted-tok")
+
+rng = np.random.default_rng(seed=42 + rank)
+
+def sample_batch():
+    fi = rng.choice(len(arrs), size=MICRO_BATCH, p=probs)
+    out = np.zeros((MICRO_BATCH, SEQ_LEN+1), dtype=np.int64)
+    for i, f in enumerate(fi):
+        arr = arrs[f]
+        mx = len(arr) - SEQ_LEN - 1
+        if mx <= 0: continue
+        s = rng.integers(0, mx)
+        out[i] = arr[s:s+SEQ_LEN+1].astype(np.int64)
+    return torch.from_numpy(out)
+
+# ── LR schedule (cosine annealing) ───────────────────────────────────────────
+def cosine_lr(step):
+    if step < WARMUP:
+        return LR_PEAK * (step + 1) / WARMUP
+    t = (step - WARMUP) / max(TOTAL_STEPS - WARMUP, 1)
+    return LR_MIN + 0.5 * (LR_PEAK - LR_MIN) * (1 + math.cos(math.pi * t))
+
+def apply_lr(step):
+    base = cosine_lr(step)
+    for g in groups:
+        name = g.get("name","")
+        if name == "embed":
+            g["lr"] = base * EMBED_LR_MULT
+        elif name.startswith("layer_"):
+            lidx = int(name.split("_")[1])
+            scale = LWD_DECAY ** (cfg.num_hidden_layers - 1 - lidx)
+            g["lr"] = base * scale
+        else:
+            g["lr"] = base
+    return base
+
+# ── checkpoint helpers ───────────────────────────────────────────────────────
+Path(CKPT_DIR).mkdir(parents=True, exist_ok=True)
+Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+
+def save_ckpt(step):
+    barrier()
+    path = Path(CKPT_DIR) / f"step_{step:07d}"
+    path.mkdir(parents=True, exist_ok=True)
+    with FSDP.state_dict_type(
+        model, StateDictType.SHARDED_STATE_DICT,
+        ShardedStateDictConfig(offload_to_cpu=True),
+        ShardedOptimStateDictConfig(offload_to_cpu=True),
+    ):
+        state = {"model": model.state_dict(),
+                 "optimizer": FSDP.optim_state_dict(model, optimizer),
+                 "step": step}
+        dcp.save(state, storage_writer=FileSystemWriter(str(path)))
+    if is_main:
+        with open(path/"meta.json","w") as f:
+            json.dump({"step":step}, f)
+        existing = sorted([d for d in Path(CKPT_DIR).iterdir()
+                           if d.is_dir() and d.name.startswith("step_")])
+        for old in existing[:-CKPT_KEEP]:
+            import shutil; shutil.rmtree(old, ignore_errors=True)
+        log(f"  [ckpt] saved step_{step:07d}")
+    barrier()
+
+# ═══════════════════════ Training loop ═══════════════════════════════════════
+log("=" * 68)
+log(f"Starting Phase 1.6 v2b training (step 1 → {TOTAL_STEPS})")
+log("=" * 68)
+
+t_log = time.time()
+tok_log = 0
+loss_acc = 0.0
+tok_per_step = MICRO_BATCH * GRAD_ACCUM * SEQ_LEN * world_size
+total_tok_b = 60.3  # tokens done so far in B (Phase 1 + v2b step 1905)
+
+for step in range(RESUME_STEP + 1, TOTAL_STEPS + 1):
+    if Path(KILL_SWITCH).exists():
+        log(f"  [kill] {KILL_SWITCH} → save+exit")
+        save_ckpt(step); dist.destroy_process_group(); sys.exit(0)
+    if shutdown_req:
+        log(f"  [shutdown] step {step} → save+exit")
+        save_ckpt(step); dist.destroy_process_group(); sys.exit(0)
+
+    lr = apply_lr(step)
+    optimizer.zero_grad(set_to_none=True)
+
+    step_loss = 0.0
+    for _ in range(GRAD_ACCUM):
+        batch = sample_batch().to(device, non_blocking=True)
+        inp = batch[:, :-1].contiguous()
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            out = model(input_ids=inp, labels=inp.clone(),
+                        output_router_logits=True)
+        loss = out.loss / GRAD_ACCUM
+        if out.aux_loss is not None:
+            loss = loss + 0.001 * out.aux_loss / GRAD_ACCUM
+        loss.backward()
+        step_loss += loss.item() * GRAD_ACCUM
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+    optimizer.step()
+
+    tok_log  += tok_per_step
+    loss_acc += step_loss
+
+    if math.isnan(step_loss) or math.isinf(step_loss):
+        log(f"  [NaN/Inf] step={step} loss={step_loss:.4f} → ABORT")
+        save_ckpt(step); dist.destroy_process_group(); sys.exit(1)
+
+    if step % LOG_EVERY == 0 and is_main:
+        elapsed = time.time() - t_log
+        tps = tok_log / elapsed
+        total_tok_b += tok_per_step * LOG_EVERY / 1e9
+        log(f"  [step {step:7d}/{TOTAL_STEPS}] "
+            f"loss={loss_acc/LOG_EVERY:.4f}  lr={lr:.2e}  "
+            f"tps={tps:.0f}  total={total_tok_b:.1f}B tok")
+        loss_acc = 0.0; tok_log = 0; t_log = time.time()
+
+    if step % SAVE_EVERY == 0:
+        save_ckpt(step)
+
+save_ckpt(TOTAL_STEPS)
+log("=== Phase 1.6 v2b COMPLETE ===")
+dist.destroy_process_group()