#!/usr/bin/env python3 import argparse import json import math import os import time from bisect import bisect_right from contextlib import nullcontext from pathlib import Path import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, Dataset, DistributedSampler from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup def is_dist(): return int(os.environ.get("WORLD_SIZE", "1")) > 1 def rank(): return int(os.environ.get("RANK", "0")) def local_rank(): return int(os.environ.get("LOCAL_RANK", "0")) def is_main(): return rank() == 0 class PackedBlockDataset(Dataset): def __init__(self, data_dir, split): self.data_dir = Path(data_dir) manifest = json.loads((self.data_dir / "manifest.json").read_text(encoding="utf-8")) self.seq_len = int(manifest.get("seq_len", 8192)) shards = manifest[f"{split}_shards"] if not shards: raise ValueError(f"no {split} shards in {self.data_dir}") self.paths = [self.data_dir / x["path"] for x in shards] self.lengths = [int(x["blocks"]) for x in shards] self.cum = np.cumsum(self.lengths).tolist() self._arrays = [None] * len(self.paths) def __len__(self): return self.cum[-1] def _array(self, shard_idx): arr = self._arrays[shard_idx] if arr is None: arr = np.load(self.paths[shard_idx], mmap_mode="r") self._arrays[shard_idx] = arr return arr def __getitem__(self, idx): shard_idx = bisect_right(self.cum, idx) prev = 0 if shard_idx == 0 else self.cum[shard_idx - 1] row_idx = idx - prev ids = np.asarray(self._array(shard_idx)[row_idx], dtype=np.int64) return torch.from_numpy(ids) def collate(batch): input_ids = torch.stack(batch, dim=0).long() return {"input_ids": input_ids, "labels": input_ids.clone()} @torch.no_grad() def evaluate(model, loader, device, max_batches=0): model.eval() total_loss = torch.tensor(0.0, device=device) total_tokens = torch.tensor(0.0, device=device) batches = 0 module = model.module if isinstance(model, DDP) else model for batch in loader: batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} out = module(**batch) ntok = batch["labels"].numel() total_loss += out.loss.float() * ntok total_tokens += ntok batches += 1 if max_batches and batches >= max_batches: break if is_dist(): dist.all_reduce(total_loss, op=dist.ReduceOp.SUM) dist.all_reduce(total_tokens, op=dist.ReduceOp.SUM) loss = (total_loss / total_tokens.clamp_min(1)).item() model.train() return {"loss": loss, "ppl": math.exp(min(loss, 20)), "tokens": int(total_tokens.item()), "batches": batches} def parse_args(): p = argparse.ArgumentParser() p.add_argument("--model", default="/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2") p.add_argument("--data", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/cpt_packed_1b_seq8192_20260614") p.add_argument("--out", default="/ssd/yi/Tokenizer_Swap/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_1b_seq8192_20260614") p.add_argument("--epochs", type=float, default=1.0) p.add_argument("--batch-size", type=int, default=2) p.add_argument("--grad-accum", type=int, default=4) p.add_argument("--lr", type=float, default=2e-5) p.add_argument("--warmup-ratio", type=float, default=0.03) p.add_argument("--eval-steps", type=int, default=100) p.add_argument("--save-steps", type=int, default=500) p.add_argument("--max-steps", type=int, default=0) p.add_argument("--eval-max-batches", type=int, default=32) p.add_argument("--num-workers", type=int, default=2) p.add_argument("--gradient-checkpointing", action="store_true") return p.parse_args() def main(): args = parse_args() if is_dist(): dist.init_process_group(backend="nccl") torch.cuda.set_device(local_rank()) device = torch.device("cuda", local_rank()) else: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") out_dir = Path(args.out) if is_main(): out_dir.mkdir(parents=True, exist_ok=True) tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device) model.config.use_cache = False if args.gradient_checkpointing: model.gradient_checkpointing_enable() for p in model.parameters(): p.requires_grad_(True) if is_dist(): model = DDP(model, device_ids=[local_rank()], output_device=local_rank(), find_unused_parameters=False) train_ds = PackedBlockDataset(args.data, "train") eval_ds = PackedBlockDataset(args.data, "eval") train_sampler = DistributedSampler(train_ds, shuffle=True) if is_dist() else None eval_sampler = DistributedSampler(eval_ds, shuffle=False) if is_dist() else None train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=train_sampler is None, sampler=train_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate) eval_loader = DataLoader(eval_ds, batch_size=args.batch_size, shuffle=False, sampler=eval_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate) steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum) total_steps = int(math.ceil(steps_per_epoch * args.epochs)) if args.max_steps > 0: total_steps = min(total_steps, args.max_steps) warmup_steps = int(total_steps * args.warmup_ratio) optim = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.1) sched = get_cosine_schedule_with_warmup(optim, warmup_steps, total_steps) meta = { "model": args.model, "data": args.data, "out": args.out, "epochs": args.epochs, "batch_size_per_rank": args.batch_size, "grad_accum": args.grad_accum, "world_size": int(os.environ.get("WORLD_SIZE", "1")), "effective_batch_tokens": args.batch_size * args.grad_accum * int(os.environ.get("WORLD_SIZE", "1")) * train_ds.seq_len, "lr": args.lr, "train_blocks": len(train_ds), "eval_blocks": len(eval_ds), "total_steps": total_steps, "warmup_steps": warmup_steps, } if is_main(): (out_dir / "train_config.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(meta, ensure_ascii=False, indent=2), flush=True) log_path = out_dir / "train_log.jsonl" start = time.time() step = 0 accum_loss = 0.0 accum_tokens = 0 model.train() optim.zero_grad(set_to_none=True) initial_eval = evaluate(model, eval_loader, device, args.eval_max_batches) if is_main(): with log_path.open("a", encoding="utf-8") as f: f.write(json.dumps({"event": "initial_eval", **initial_eval}, ensure_ascii=False) + "\n") print(json.dumps({"event": "initial_eval", **initial_eval}, ensure_ascii=False), flush=True) epoch = 0 while step < total_steps: if train_sampler is not None: train_sampler.set_epoch(epoch) epoch += 1 for batch_idx, batch in enumerate(train_loader): batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} sync_grad = (batch_idx + 1) % args.grad_accum == 0 sync_context = model.no_sync() if isinstance(model, DDP) and not sync_grad else nullcontext() with sync_context: out = model(**batch) loss = out.loss / args.grad_accum loss.backward() ntok = batch["labels"].numel() accum_loss += float(out.loss.item()) * ntok accum_tokens += ntok if (batch_idx + 1) % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optim.step() sched.step() optim.zero_grad(set_to_none=True) step += 1 if is_main() and (step == 1 or step % 10 == 0): rec = {"event": "train", "step": step, "loss": accum_loss / max(accum_tokens, 1), "tokens": accum_tokens, "lr": sched.get_last_lr()[0], "elapsed_sec": time.time() - start} print(json.dumps(rec, ensure_ascii=False), flush=True) with log_path.open("a", encoding="utf-8") as f: f.write(json.dumps(rec, ensure_ascii=False) + "\n") accum_loss = 0.0 accum_tokens = 0 if args.eval_steps and step % args.eval_steps == 0: rec = {"event": "eval", "step": step, **evaluate(model, eval_loader, device, args.eval_max_batches)} if is_main(): print(json.dumps(rec, ensure_ascii=False), flush=True) with log_path.open("a", encoding="utf-8") as f: f.write(json.dumps(rec, ensure_ascii=False) + "\n") if is_main() and args.save_steps and step % args.save_steps == 0: ckpt = out_dir / f"checkpoint-{step}" (model.module if isinstance(model, DDP) else model).save_pretrained(ckpt, safe_serialization=True) tokenizer.save_pretrained(ckpt) if step >= total_steps: break final_eval = evaluate(model, eval_loader, device, args.eval_max_batches) if is_main(): module = model.module if isinstance(model, DDP) else model module.save_pretrained(out_dir, safe_serialization=True) tokenizer.save_pretrained(out_dir) (out_dir / "DONE").write_text("ok\n", encoding="utf-8") with log_path.open("a", encoding="utf-8") as f: f.write(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False) + "\n") print(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False), flush=True) if is_dist(): dist.destroy_process_group() if __name__ == "__main__": main()