yi_lu/Tokenizer_Swap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean tokenizer swap migration

Browse Source
This commit is contained in:
Codex
2026-06-18 10:10:57 +00:00
commit 5d1e9c4bc3
37 changed files with 9574 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
model_training/README.md Normal file
View File

@@ -0,0 +1,64 @@
# Model Training
This folder owns full-parameter training recipes for the final experiments.
It does not build datasets and does not perform tokenizer remapping. It consumes artifacts produced by `dataset_building/` and `model_building/`.
## Training Implementations
```text
train_dsv4_tokenized_full_sft.py
train_cpt_packed_full.py
```
`train_dsv4_tokenized_full_sft.py` trains on pre-tokenized chat JSONL. Prompt tokens are masked with `-100`; assistant tokens are optimized.
`train_cpt_packed_full.py` trains next-token prediction over packed CPT blocks.
## Final Recipe Scripts
```text
run_sft1m_remap_v2_5epoch.sh
run_sft1m_remap_v2_then_v4_noupsample_5epoch_bsz16.sh
run_cpt1b_seed42_train_eval.sh
run_cpt5b_seed42_train_eval.sh
run_cpt5b_then_sft1m_5epoch.sh
```
## Common Environment Variables
The run scripts are configurable through environment variables:
```text
ROOT repo root, default /ssd/yi/Tokenizer_Swap
NPROC number of GPUs/processes, default 8
MODEL input checkpoint
DATA packed CPT dataset directory
TRAIN tokenized SFT train file
EVAL tokenized SFT validation file
OUT output checkpoint directory
```
Example:
```bash
ROOT=/ssd/yi/Tokenizer_Swap \
MODEL=/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2 \
DATA=/ssd/yi/Tokenizer_Swap/dataset_building/generated/cpt_packed_5b_seq8192_seed42_stratified \
OUT=/ssd/yi/Tokenizer_Swap/model_training/checkpoints/cpt5b \
bash model_training/run_cpt5b_seed42_train_eval.sh
```
## Output
Generated checkpoints go under:
```text
model_training/checkpoints/
```
This directory is ignored by git. Do not commit model weights, optimizer states, or partial checkpoints.
## Output Contract
Trained checkpoints are consumed by `evaluation_reporting/` as `MODEL`.

21
model_training/run_cpt1b_seed42_train_eval.sh Executable file
View File

@@ -0,0 +1,21 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
DATA=${DATA:-$ROOT/dataset_building/generated/cpt_packed_1b_seq8192_seed42_stratified}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_1b_seed42_seq8192_bsz3acc3}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_cpt_packed_full.py" \
--model "$MODEL" \
--data "$DATA" \
--out "$OUT" \
--epochs 1 \
--batch-size 3 \
--grad-accum 3 \
--lr 2e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

21
model_training/run_cpt5b_seed42_train_eval.sh Executable file
View File

@@ -0,0 +1,21 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
DATA=${DATA:-$ROOT/dataset_building/generated/cpt_packed_5b_seq8192_seed42_stratified}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_5b_seed42_seq8192_bsz3acc3}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_cpt_packed_full.py" \
--model "$MODEL" \
--data "$DATA" \
--out "$OUT" \
--epochs 1 \
--batch-size 3 \
--grad-accum 3 \
--lr 2e-5 \
--eval-steps 2000 \
--save-steps 2000 \
--num-workers 2

24
model_training/run_cpt5b_then_sft1m_5epoch.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_5b_seed42_seq8192_bsz3acc3}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt5b_then_sft1m_5epoch}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 300 \
--num-workers 2

24
model_training/run_sft1m_remap_v2_5epoch.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_alt_sources_1m_fixed_eval_20260607/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_alt_sources_1m_fixed_eval_20260607/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_5epoch_bsz8_accum16}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

24
model_training/run_sft1m_remap_v2_then_v4_noupsample_5epoch_bsz16.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_5epoch_bsz8_accum16}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_then_v4_noupsample_5epoch_bsz16}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

247
model_training/train_cpt_packed_full.py Normal file
View File

@@ -0,0 +1,247 @@
#!/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()

303
model_training/train_dsv4_tokenized_full_sft.py Executable file
View File

@@ -0,0 +1,303 @@
#!/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()