248 lines
10 KiB
Python
248 lines
10 KiB
Python
#!/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()
|