Tokenizer_Swap/model_training/train_dsv4_tokenized_full_sft.py

#!/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()