#!/usr/bin/env python3 import argparse import gzip import json import math import os import time from pathlib import Path 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 def open_text(path: Path): if path.suffix == ".gz": return gzip.open(path, "rt", encoding="utf-8") return path.open("r", encoding="utf-8") class TokenizedSFTDataset(Dataset): def __init__(self, path: str, max_length: int): self.rows = [] self.max_length = max_length with open_text(Path(path)) as f: for line in f: if line.strip(): row = json.loads(line) self.rows.append( { "input_ids": row["input_ids"][:max_length], "labels": row["labels"][:max_length], "id": row.get("id"), } ) def __len__(self): return len(self.rows) def __getitem__(self, idx): return self.rows[idx] def collate(batch, pad_id): max_len = max(len(x["input_ids"]) for x in batch) input_ids, labels, attention_mask = [], [], [] for item in batch: ids = item["input_ids"] lab = item["labels"] pad = max_len - len(ids) input_ids.append(ids + [pad_id] * pad) labels.append(lab + [-100] * pad) attention_mask.append([1] * len(ids) + [0] * pad) return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.long), } def set_trainable_full(model): for p in model.parameters(): p.requires_grad_(True) return ["all_parameters"] @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) if isinstance(model, DDP) else model(**batch) ntok = (batch["labels"] != -100).sum().float() 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("--train", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/dsv4_chat_tokenized_fixed100k_20260605/train_dsv4_chat_tokenized.jsonl.gz") p.add_argument("--eval", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/dsv4_chat_tokenized_fixed100k_20260605/validation_dsv4_chat_tokenized.jsonl.gz") p.add_argument("--out", default="/ssd/yi/Tokenizer_Swap/model_training/checkpoints/dsv4_chat_full_sft_fixed100k_20260605") p.add_argument("--max-length", type=int, default=2048) p.add_argument("--epochs", type=float, default=3.0) p.add_argument("--batch-size", type=int, default=2) p.add_argument("--grad-accum", type=int, default=8) p.add_argument("--lr", type=float, default=5e-5) p.add_argument("--warmup-ratio", type=float, default=0.03) p.add_argument("--eval-steps", type=int, default=200) p.add_argument("--save-steps", type=int, default=0) p.add_argument("--max-steps", type=int, default=0) p.add_argument("--eval-max-batches", type=int, default=0) 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) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id 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() trainable_names = set_trainable_full(model) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) total_params = sum(p.numel() for p in model.parameters()) if is_dist(): model = DDP(model, device_ids=[local_rank()], output_device=local_rank(), find_unused_parameters=False) train_ds = TokenizedSFTDataset(args.train, args.max_length) eval_ds = TokenizedSFTDataset(args.eval, args.max_length) 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=lambda b: collate(b, pad_id), ) 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=lambda b: collate(b, pad_id), ) 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((p for p in model.parameters() if p.requires_grad), lr=args.lr, weight_decay=0.01) sched = get_cosine_schedule_with_warmup(optim, warmup_steps, total_steps) meta = { "model": args.model, "train": args.train, "eval": args.eval, "out": args.out, "max_length": args.max_length, "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_examples": args.batch_size * args.grad_accum * int(os.environ.get("WORLD_SIZE", "1")), "lr": args.lr, "train_examples": len(train_ds), "eval_examples": len(eval_ds), "total_steps": total_steps, "warmup_steps": warmup_steps, "trainable_names": trainable_names, "trainable_params": trainable_params, "total_params": total_params, "trainable_fraction": trainable_params / total_params, } 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, max_batches=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") 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()} out = model(**batch) loss = out.loss / args.grad_accum loss.backward() ntok = int((batch["labels"] != -100).sum().item()) accum_loss += float(out.loss.item()) * ntok accum_tokens += ntok if (batch_idx + 1) % args.grad_accum == 0: torch.nn.utils.clip_grad_norm_((p for p in model.parameters() if p.requires_grad), 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, max_batches=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()