#!/usr/bin/env python3 # train_lm_adamw_8bit.py # # Classic Hugging Face Trainer script for causal LM fine-tuning on: # ./train.txt and ./valid.txt # # Extra: # - Replaces literal "\\n" with real newline "\n" inside each sample. # - Ensures each text sample ends with literal "". # - Prints “what the model sees” before training. import os import json import random import torch from datasets import load_dataset from transformers import ( AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, set_seed, ) # ---------------------------- # Paths / knobs # ---------------------------- TOKENIZER_PATH = "./checkpoint-55000" POLICY_CKPT = "./checkpoint-55000" TRAIN_TXT = "./train.txt" VALID_TXT = "./valid.txt" OUT_DIR = "./runs/lm_finetune_adamw8bit" LOG_DIR = os.path.join(OUT_DIR, "tb") SEED = 42 MAX_LENGTH = 512 EOS_STR = "" # Preview controls PREVIEW_TEXT_SAMPLES = 8 PREVIEW_TOKEN_TRUNC = 256 PREVIEW_BATCH_ITEMS = 2 WRITE_PREVIEW_JSONL = True # Training knobs MAX_STEPS = 50000 PER_DEVICE_TRAIN_BATCH_SIZE = 16 PER_DEVICE_EVAL_BATCH_SIZE = 64 GRADIENT_ACCUMULATION_STEPS = 1 LEARNING_RATE = 5e-5 WEIGHT_DECAY = 0.0 WARMUP_RATIO = 0.01 LR_SCHEDULER_TYPE = "cosine" DATALOADER_NUM_WORKERS = 16 def preprocess_text(text: str) -> str: """ 1) Convert literal '\\n' into a real newline '\n' inside the sample. 2) Append literal EOS_STR if missing. """ t = text.rstrip("\n") t = t.replace("\\n", "\n") if not t.rstrip().endswith(EOS_STR): t = t + EOS_STR return t def safe_makedirs(path: str) -> None: os.makedirs(path, exist_ok=True) def tokens_preview(tokenizer, ids, max_tokens=80): toks = tokenizer.convert_ids_to_tokens(ids) if len(toks) > max_tokens: return toks[:max_tokens] + ["…"] return toks def dump_jsonl(path: str, rows): with open(path, "w", encoding="utf-8") as f: for r in rows: f.write(json.dumps(r, ensure_ascii=False) + "\n") class CausalLMDataCollator: """ Pads input_ids and labels for causal language modeling. """ def __init__(self, tokenizer): self.tokenizer = tokenizer def __call__(self, features): batch = self.tokenizer.pad( {"input_ids": [f["input_ids"] for f in features]}, padding=True, return_tensors="pt", ) labels = [] max_len = batch["input_ids"].shape[1] for f in features: ids = f["labels"] ids = ids + [self.tokenizer.pad_token_id] * (max_len - len(ids)) labels.append(ids[:max_len]) batch["labels"] = torch.tensor(labels, dtype=torch.long) return batch def main(): set_seed(SEED) torch.backends.cuda.matmul.allow_tf32 = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("CUDA available:", torch.cuda.is_available()) if torch.cuda.is_available(): print("GPU:", torch.cuda.get_device_name(0)) # ---------------------------- # Load tokenizer/model # ---------------------------- tokenizer = AutoTokenizer.from_pretrained( TOKENIZER_PATH, use_fast=True, trust_remote_code=True, return_token_type_ids=False, ) if tokenizer.pad_token is None: if tokenizer.eos_token is not None: tokenizer.pad_token = tokenizer.eos_token else: tokenizer.add_special_tokens({"pad_token": "<|pad|>"}) model = AutoModelForCausalLM.from_pretrained( POLICY_CKPT, torch_dtype=torch.float32, device_map=None, trust_remote_code=True, ) if len(tokenizer) != model.get_input_embeddings().num_embeddings: model.resize_token_embeddings(len(tokenizer)) model.config.use_cache = False model.to(device) # ---------------------------- # Load datasets # ---------------------------- raw = load_dataset( "text", data_files={"train": TRAIN_TXT, "validation": VALID_TXT}, ) raw = raw.filter(lambda x: len(x["text"].strip()) > 0) def preprocess_map(batch): processed = [preprocess_text(t) for t in batch["text"]] return {"text": processed} raw = raw.map( preprocess_map, batched=True, desc="Preprocess: literal \\n -> newline + ensure ", ) # ---------------------------- # ---------------------------- def tokenize_batch(batch): texts = batch["text"] enc = tokenizer( texts, truncation=True, max_length=MAX_LENGTH, padding=False, ) enc["labels"] = [ids[:] for ids in enc["input_ids"]] return enc tok = raw.map( tokenize_batch, batched=True, remove_columns=[], desc="Tokenizing", ) drop_cols = [] if "text" in tok["train"].column_names: drop_cols.append("text") train_ds = tok["train"].remove_columns(drop_cols) if drop_cols else tok["train"] valid_ds = tok["validation"].remove_columns(drop_cols) if drop_cols else tok["validation"] # ---------------------------- # Preview: sample-level # ---------------------------- safe_makedirs(OUT_DIR) train_len = len(tok["train"]) k = min(PREVIEW_TEXT_SAMPLES, train_len) random.seed(SEED) preview_indices = random.sample(range(train_len), k=k) if train_len > 0 else [] preview_rows = [] print("\n=== PREVIEW: samples after preprocess + tokenization + labels ===") for idx in preview_indices: ex = tok["train"][idx] text = ex["text"] input_ids = ex["input_ids"] labels = ex["labels"] label_ids = labels row = { "split": "train", "index": idx, "text": text, "len_ids": len(input_ids), "input_ids_head": input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)], "tokens_head": tokens_preview( tokenizer, input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)], PREVIEW_TOKEN_TRUNC, ), "decoded_full": tokenizer.decode(input_ids, skip_special_tokens=False), "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False), } preview_rows.append(row) print(f"\n--- sample idx={idx} ---") print("TEXT postprocessed:") print(text) print(f"len(input_ids) = {len(input_ids)}") print("TOKENS head:") print(row["tokens_head"]) print("DECODED PROMPT:") print(row["decoded_prompt"]) print("DECODED TARGET:") print(row["decoded_labels"]) if WRITE_PREVIEW_JSONL: dump_path = os.path.join(OUT_DIR, "preview_samples.jsonl") dump_jsonl(dump_path, preview_rows) print(f"\nWrote sample previews to: {dump_path}") # ---------------------------- # Collator + Trainer # ---------------------------- data_collator = CausalLMDataCollator(tokenizer=tokenizer) args = TrainingArguments( output_dir=OUT_DIR, overwrite_output_dir=True, max_steps=50000, per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH_SIZE, per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH_SIZE, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY, warmup_ratio=WARMUP_RATIO, lr_scheduler_type=LR_SCHEDULER_TYPE, logging_dir=LOG_DIR, logging_steps=1, report_to=["tensorboard"], eval_strategy="epoch", save_strategy="epoch", dataloader_num_workers=DATALOADER_NUM_WORKERS, remove_unused_columns=False, load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, optim="adamw_torch", ) trainer = Trainer( model=model, args=args, train_dataset=train_ds, eval_dataset=valid_ds, tokenizer=tokenizer, data_collator=data_collator, ) # ---------------------------- # Preview: collated batch # ---------------------------- print("\n=== PREVIEW: a collated training batch after padding + labels ===") n_take = min(2, len(train_ds)) features = [train_ds[i] for i in range(n_take)] first_batch = data_collator(features) input_ids_b = first_batch["input_ids"].detach().cpu() labels_b = first_batch["labels"].detach().cpu() attn_b = first_batch["attention_mask"].detach().cpu() print("Batch keys:", list(first_batch.keys())) print("input_ids shape:", tuple(input_ids_b.shape)) print("labels shape:", tuple(labels_b.shape)) print("attention_mask shape:", tuple(attn_b.shape)) batch_preview_rows = [] n_show = min(PREVIEW_BATCH_ITEMS, input_ids_b.shape[0]) for i in range(n_show): ids = input_ids_b[i].tolist() labels = labels_b[i].tolist() label_ids = labels row = { "batch_item": i, "input_len": int(attn_b[i].sum().item()), "input_ids_head": ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)], "tokens_head": tokens_preview( tokenizer, ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)], PREVIEW_TOKEN_TRUNC, ), "decoded_input": tokenizer.decode(ids, skip_special_tokens=False), "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False), } batch_preview_rows.append(row) print(f"\n--- batch item {i} ---") print("TOKENS head:") print(row["tokens_head"]) print("DECODED INPUT, specials kept:") print(row["decoded_input"]) print("DECODED LABELS, only learned tokens:") print(row["decoded_labels"]) if WRITE_PREVIEW_JSONL: dump_path = os.path.join(OUT_DIR, "preview_first_batch.jsonl") dump_jsonl(dump_path, batch_preview_rows) print(f"\nWrote batch preview to: {dump_path}") # ---------------------------- # Train / save # ---------------------------- train_result = trainer.train() final_dir = os.path.join(OUT_DIR, "final") trainer.save_model(final_dir) tokenizer.save_pretrained(final_dir) metrics = train_result.metrics metrics["train_samples"] = len(train_ds) trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() eval_metrics = trainer.evaluate() eval_metrics["eval_samples"] = len(valid_ds) trainer.log_metrics("eval", eval_metrics) trainer.save_metrics("eval", eval_metrics) print("\nDone.") print(f"Final model saved to: {final_dir}") print(f"TensorBoard: tensorboard --logdir {OUT_DIR}") if __name__ == "__main__": main()