初始化项目，由ModelHub XC社区提供模型
Model: iamrahulreddy/Quintus Source: Original Platform
This commit is contained in:
ModelHub XC
89
sft/chat.py
Normal file
89
sft/chat.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
|
||||
import sys
|
||||
import argparse
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Quintus Interactive Chat")
|
||||
parser.add_argument("--model_path", type=str, default="iamrahulreddy/Quintus", help="Model repo ID or local weights directory")
|
||||
parser.add_argument("--trust_remote_code", action="store_true", help="Allow custom code from the model repository.")
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = args.model_path
|
||||
print(f"Loading Quintus from {model_path}...")
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=args.trust_remote_code)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
trust_remote_code=args.trust_remote_code
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading model: {e}")
|
||||
print(f"Ensure '{model_path}' exists and contains the model weights.")
|
||||
sys.exit(1)
|
||||
|
||||
# Defining stopping criteria
|
||||
stop_tokens = ["<|endoftext|>", "<|im_end|>"]
|
||||
eos_token_ids = [tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else []
|
||||
for token in stop_tokens:
|
||||
t_id = tokenizer.convert_tokens_to_ids(token)
|
||||
if t_id is not None and t_id not in eos_token_ids:
|
||||
eos_token_ids.append(t_id)
|
||||
|
||||
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
|
||||
conversation_history = [
|
||||
{"role": "system", "content": "You are Quintus, a highly capable AI assistant created by Muskula Rahul. You are helpful, precise, and logically sound."}
|
||||
]
|
||||
|
||||
print()
|
||||
print("Quintus Chat (type 'quit' to exit)")
|
||||
print()
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("You: ").strip()
|
||||
if user_input.lower() in ["quit", "exit"]:
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
conversation_history.append({"role": "user", "content": user_input})
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
conversation_history,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
||||
|
||||
print("Quintus: ", end="", flush=True)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=512,
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
do_sample=True,
|
||||
streamer=streamer,
|
||||
pad_token_id=tokenizer.eos_token_id,
|
||||
eos_token_id=eos_token_ids
|
||||
)
|
||||
|
||||
# Extract response for history
|
||||
generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
|
||||
assistant_response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
|
||||
conversation_history.append({"role": "assistant", "content": assistant_response})
|
||||
print()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nGoodbye!")
|
||||
break
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
267
sft/evaluate.py
Normal file
267
sft/evaluate.py
Normal file
@@ -0,0 +1,267 @@
|
||||
# Automated EvalPlus runner for HumanEval and MBPP benchmarks.
|
||||
# Using the vLLM backend in greedy mode.
|
||||
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
import time
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
MODELS = [
|
||||
{
|
||||
"name": "Quintus-1.7B",
|
||||
"id": "iamrahulreddy/Quintus",
|
||||
"is_local": False
|
||||
},
|
||||
{
|
||||
"name": "Qwen3-1.7B-Instruct",
|
||||
"id": "Qwen/Qwen3-1.7B",
|
||||
"is_local": False
|
||||
},
|
||||
{
|
||||
"name": "Qwen3-1.7B-Base",
|
||||
"id": "Qwen/Qwen3-1.7B-Base",
|
||||
"is_local": False
|
||||
}
|
||||
]
|
||||
|
||||
DATASETS = [
|
||||
"humaneval", "mbpp", # EvalPlus benchmarks
|
||||
"gsm8k", "winogrande", # lm-eval fast benchmarks
|
||||
"arc_challenge", "boolq", "piqa"
|
||||
]
|
||||
EVALPLUS_DATASETS = {"humaneval", "mbpp"}
|
||||
|
||||
LM_EVAL_SHOTS = {
|
||||
"gsm8k": "10",
|
||||
"winogrande": "5",
|
||||
"arc_challenge": "25",
|
||||
"boolq": "0",
|
||||
"piqa": "0"
|
||||
}
|
||||
|
||||
HF_TOKEN = os.environ.get("HF_TOKEN")
|
||||
TRUST_REMOTE_CODE = os.environ.get("QUINTUS_TRUST_REMOTE_CODE", "").strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
def extract_lm_eval_score(results_dir: Path, task: str) -> str:
|
||||
"""Finds and extracts the primary score from JSON files outputted by lm-evaluation-harness."""
|
||||
for json_path in sorted(results_dir.rglob("*.json"), reverse=True):
|
||||
try:
|
||||
with open(json_path, encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
task_results = data.get("results", {})
|
||||
for candidate in (task, f"leaderboard_{task}"):
|
||||
if candidate in task_results:
|
||||
task_data = task_results[candidate]
|
||||
# Try common metric names
|
||||
for metric in ["acc,none", "acc_norm,none", "exact_match,strict-match", "exact_match,none"]:
|
||||
if metric in task_data:
|
||||
return f"{task_data[metric]*100:.1f}"
|
||||
except Exception:
|
||||
continue
|
||||
return "N/A"
|
||||
|
||||
def is_noise(line: str) -> bool:
|
||||
l = line.strip()
|
||||
if not l:
|
||||
return False
|
||||
# Progress bar indicators & block characters
|
||||
if any(c in l for c in ["█", "━", "╸", "•", "━━━━━━━━"]):
|
||||
return True
|
||||
# vLLM, ray, flash_attn, huggingface setup/warnings logs
|
||||
noise_keywords = [
|
||||
"INFO ", "WARNING ", "DEBUG ", "ERROR ", "(EngineCore",
|
||||
"Loading safetensors", "Capturing CUDA graphs",
|
||||
"Codegen:", "Downloading dataset", "downloading dataset",
|
||||
"Initializing a decoder", "Unknown vLLM environment",
|
||||
"world_size=", "Using V2 Model Runner", "Model loading took",
|
||||
"Using FLASH_ATTN", "Using FlashAttention", "Kernel JIT monitor",
|
||||
"autotuner.py", "autotuning", "Autotuning", "loading weights",
|
||||
"Loading weights", "Failed to get device capability", "Sanitized code outputs",
|
||||
"Raw outputs will be saved", "init engine", "Dynamo bytecode",
|
||||
"Directly load the compiled graph", "Directly load AOT compilation", "torch.compile took"
|
||||
]
|
||||
if any(k.lower() in l.lower() for k in noise_keywords):
|
||||
return True
|
||||
# TQDM lines (e.g. 100%|... [00:17<00:00, 9.45it/s])
|
||||
if "%|" in l and ("it/s" in l or "s/it" in l):
|
||||
return True
|
||||
return False
|
||||
|
||||
def main():
|
||||
print("=" * 80)
|
||||
print(" EVALPLUS BENCHMARK RUNNER (HUMANEVAL & MBPP)")
|
||||
print("=" * 80)
|
||||
print(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print(f"Models to evaluate: {[m['name'] for m in MODELS]}")
|
||||
print(f"Datasets: {DATASETS}")
|
||||
print("=" * 80)
|
||||
|
||||
# Set optional HF token and runtime configuration.
|
||||
if HF_TOKEN:
|
||||
os.environ["HF_TOKEN"] = HF_TOKEN
|
||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
os.environ["VLLM_MAX_MODEL_LEN"] = "4096"
|
||||
|
||||
# Step 1: Pre-download and prepare model caches
|
||||
print("\n--- STAGE 1: WARMING UP MODEL WEIGHTS CACHE ---")
|
||||
|
||||
# Cache all models
|
||||
for model in MODELS:
|
||||
if model["is_local"]:
|
||||
continue
|
||||
print(f"\n[DOWNLOADING] Fetching cache for {model['name']} ({model['id']})...")
|
||||
try:
|
||||
snapshot_download(
|
||||
repo_id=model["id"],
|
||||
token=HF_TOKEN or None
|
||||
)
|
||||
print(f"[DOWNLOAD SUCCESS] {model['name']} is cached and ready.")
|
||||
except Exception as e:
|
||||
print(f"[DOWNLOAD WARNING] Could not pre-download model {model['name']} via snapshot_download: {e}")
|
||||
print("The evaluation run will attempt to download it directly during execution.")
|
||||
|
||||
print("\n--- STAGE 2: SEQUENTIAL EVALPLUS EVALUATION ---")
|
||||
results = []
|
||||
|
||||
# Run evaluations sequentially
|
||||
for model in MODELS:
|
||||
# Resolve path
|
||||
model_path = str(Path(model["id"]).resolve()) if model["is_local"] else model["id"]
|
||||
|
||||
for dataset in DATASETS:
|
||||
print(f"\n[STARTING] Evaluating {model['name']} on {dataset}...")
|
||||
print("-" * 60)
|
||||
|
||||
if dataset in EVALPLUS_DATASETS:
|
||||
cmd = [
|
||||
sys.executable, "-m", "evalplus.evaluate",
|
||||
"--model", model_path,
|
||||
"--dataset", dataset,
|
||||
"--backend", "vllm",
|
||||
"--greedy"
|
||||
]
|
||||
else:
|
||||
shots = LM_EVAL_SHOTS.get(dataset, "0")
|
||||
out_dir = Path("eval_results") / model["name"] / dataset
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_args = (
|
||||
f"pretrained={model_path},dtype=bfloat16,"
|
||||
f"trust_remote_code={str(TRUST_REMOTE_CODE).lower()},"
|
||||
"gpu_memory_utilization=0.9,max_model_len=4096"
|
||||
)
|
||||
cmd = [
|
||||
sys.executable, "-m", "lm_eval",
|
||||
"--model", "vllm",
|
||||
"--model_args", model_args,
|
||||
"--tasks", dataset,
|
||||
"--num_fewshot", shots,
|
||||
"--batch_size", "auto",
|
||||
"--output_path", str(out_dir),
|
||||
"--log_samples"
|
||||
]
|
||||
if dataset == "gsm8k":
|
||||
cmd.extend(["--gen_kwargs", "max_gen_toks=512"])
|
||||
|
||||
print(f"Running command: {' '.join(cmd)}")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Run the command and stream output
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1
|
||||
)
|
||||
|
||||
# Stream and capture output (filtering out vLLM and progress bar noise)
|
||||
stdout_text = ""
|
||||
for line in process.stdout:
|
||||
stdout_text += line
|
||||
if not is_noise(line):
|
||||
print(line, end="")
|
||||
|
||||
process.wait()
|
||||
duration = time.time() - start_time
|
||||
time.sleep(5) # Let OS/driver fully reclaim GPU VRAM before starting next subprocess
|
||||
|
||||
score_str = "N/A"
|
||||
if process.returncode == 0:
|
||||
print(f"[SUCCESS] Completed {model['name']} on {dataset} in {duration:.1f} seconds.")
|
||||
|
||||
# Parse scores
|
||||
if dataset in EVALPLUS_DATASETS:
|
||||
# Find all pass@1 scores
|
||||
matches = re.findall(r"pass@1:\s+([0-9.]+)", stdout_text)
|
||||
if len(matches) >= 2:
|
||||
val0 = float(matches[0])
|
||||
val1 = float(matches[1])
|
||||
if val0 <= 1.0: val0 *= 100
|
||||
if val1 <= 1.0: val1 *= 100
|
||||
score_str = f"Base: {val0:.1f} | Plus: {val1:.1f}"
|
||||
elif len(matches) == 1:
|
||||
val0 = float(matches[0])
|
||||
if val0 <= 1.0: val0 *= 100
|
||||
score_str = f"Base: {val0:.1f}"
|
||||
else:
|
||||
score_str = extract_lm_eval_score(out_dir, dataset)
|
||||
|
||||
results.append({
|
||||
"model": model["name"],
|
||||
"dataset": dataset,
|
||||
"status": "Success",
|
||||
"score": score_str,
|
||||
"duration": f"{duration/60:.1f} min"
|
||||
})
|
||||
else:
|
||||
print(f"[ERROR] command failed with exit code {process.returncode}")
|
||||
results.append({
|
||||
"model": model["name"],
|
||||
"dataset": dataset,
|
||||
"status": f"Failed ({process.returncode})",
|
||||
"score": "ERROR",
|
||||
"duration": f"{duration/60:.1f} min"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
duration = time.time() - start_time
|
||||
print(f"[ERROR] Failed to run benchmark: {e}")
|
||||
results.append({
|
||||
"model": model["name"],
|
||||
"dataset": dataset,
|
||||
"status": f"Error",
|
||||
"score": "ERROR",
|
||||
"duration": f"{duration/60:.1f} min"
|
||||
})
|
||||
print("-" * 60)
|
||||
|
||||
# Print and save summary report
|
||||
report_lines = []
|
||||
report_lines.append("\n" + "=" * 100)
|
||||
report_lines.append(" BENCHMARK RUN SUMMARY")
|
||||
report_lines.append("=" * 100)
|
||||
report_lines.append(f"| {'Model':<30} | {'Dataset':<15} | {'Score':<25} | {'Status':<10} | {'Time':<8} |")
|
||||
report_lines.append(f"|{'-'*32}|{'-'*17}|{'-'*27}|{'-'*12}|{'-'*10}|")
|
||||
for r in results:
|
||||
report_lines.append(f"| {r['model']:<30} | {r['dataset']:<15} | {r['score']:<25} | {r['status']:<10} | {r['duration']:<8} |")
|
||||
report_lines.append("=" * 100)
|
||||
|
||||
report_text = "\n".join(report_lines)
|
||||
print(report_text)
|
||||
print("\nNote: Results are saved in the default EvalPlus directory and eval_results/.")
|
||||
|
||||
# Save to file
|
||||
with open("qwen_quintus_scores.txt", "w", encoding="utf-8") as f:
|
||||
f.write(report_text + "\n")
|
||||
print("\n[SUCCESS] Final score report saved to 'qwen_quintus_scores.txt'")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
690
sft/train_sft.py
Normal file
690
sft/train_sft.py
Normal file
@@ -0,0 +1,690 @@
|
||||
# SFT Training and Downstream Evaluation Pipeline
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
import yaml
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup
|
||||
|
||||
# Load Configuration
|
||||
def load_config() -> dict:
|
||||
cfg_path = Path(__file__).resolve().parent / "config.yaml"
|
||||
if not cfg_path.exists():
|
||||
return {}
|
||||
with open(cfg_path, "r", encoding="utf-8") as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
|
||||
cfg = load_config()
|
||||
|
||||
# PROMPTS (50 PROMPTS)
|
||||
EASY_PROMPTS = [
|
||||
"What is the capital of Japan, and what is it known for?",
|
||||
"What does the term 'CPU' stand for, and what is its role in a computer?",
|
||||
"Name three mammals that live primarily in water.",
|
||||
"What is the difference between a virus and a bacterium?",
|
||||
"Convert 72 degrees Fahrenheit to Celsius.",
|
||||
"What is the purpose of a hash function?",
|
||||
"What does HTTP stand for and what is it used for?",
|
||||
"In which continent is the Amazon rainforest located?",
|
||||
"What is the difference between RAM and ROM?",
|
||||
"Name two programming languages commonly used for data science.",
|
||||
"What is the function of the mitochondria in a cell?",
|
||||
"What is a palindrome? Give two examples.",
|
||||
"What is the difference between a compiler and an interpreter?",
|
||||
"What unit is used to measure electrical resistance?",
|
||||
"Name the four blood types in the ABO system.",
|
||||
"What is the primary purpose of DNS in networking?",
|
||||
"What does it mean for a function to be 'pure' in programming?"
|
||||
]
|
||||
|
||||
MEDIUM_PROMPTS = [
|
||||
"Explain the difference between supervised and unsupervised learning with a concrete example of each.",
|
||||
"Write a Python function that takes a list of integers and returns all pairs that sum to a given target.",
|
||||
"Explain how TCP/IP ensures reliable data delivery over an unreliable network.",
|
||||
"What are the trade-offs between using a relational database and a document store for a user profile system?",
|
||||
"Describe how gradient descent works and explain the role of the learning rate.",
|
||||
"Write a SQL query that returns the top 5 customers by total order value, including customers with no orders.",
|
||||
"What is the CAP theorem and what does it imply for distributed system design?",
|
||||
"Explain the difference between process and thread, including when you would prefer one over the other.",
|
||||
"How does HTTPS prevent a man-in-the-middle attack? Walk through the handshake at a high level.",
|
||||
"Write a regex that validates an email address and annotate each part of the pattern.",
|
||||
"What is the difference between memoization and dynamic programming?",
|
||||
"Describe three ways to handle class imbalance in a machine learning dataset.",
|
||||
"Explain what a foreign key constraint does and give an example of why it matters.",
|
||||
"What is the difference between horizontal and vertical scaling, and when would you choose each?",
|
||||
"How does Python's garbage collector handle circular references?",
|
||||
"Explain the intuition behind the attention mechanism in Transformer models.",
|
||||
"What is a race condition? Write a minimal pseudocode example that demonstrates one."
|
||||
]
|
||||
|
||||
TOUGH_PROMPTS = [
|
||||
"Design a rate limiter for a public API that must handle 100k requests per second across multiple regions. Describe the data structures, algorithms, and infrastructure trade-offs involved.",
|
||||
"Explain why training very deep neural networks with sigmoid activations suffers from vanishing gradients. How do residual connections and normalization layers address this, and what are their respective limitations?",
|
||||
"A message queue is consuming events from an upstream producer faster than a downstream consumer can process them. The queue is filling up and the producer cannot be slowed down. Describe at least three architectural strategies to resolve this, with trade-offs.",
|
||||
"Given an undirected weighted graph, write Python code to find the minimum spanning tree using Kruskal's algorithm. Include the union-find data structure. Analyze time and space complexity.",
|
||||
"You are given two sorted arrays of size m and n. Find the median of the combined array in O(log(m+n)) time. Explain the approach before writing the code.",
|
||||
"Explain the difference between Byzantine fault tolerance and crash fault tolerance. In what scenario does the distinction become critical, and how does a consensus protocol like PBFT address Byzantine failures?",
|
||||
"A large language model fine-tuned on customer service data starts producing confident but factually wrong answers about product details. Propose a complete mitigation strategy covering training, inference, and deployment layers.",
|
||||
"Explain the mechanism behind speculative execution in modern CPUs and how it led to the Spectre vulnerability. What classes of software-level mitigations exist and what performance cost do they carry?",
|
||||
"Design a schema and indexing strategy for a social graph where you need to efficiently answer: (1) mutual friends between two users, (2) shortest path between two users, (3) top-k most influential accounts. Justify your choices.",
|
||||
"Implement a thread-safe LRU cache in Python with O(1) get and put operations. Explain why your synchronization approach is correct and where contention bottlenecks might appear under high concurrency.",
|
||||
"Explain the difference between weak, strong, and eventual consistency in distributed databases. Give a concrete example of a bug that arises when a developer assumes strong consistency but the system only guarantees eventual consistency.",
|
||||
"You are designing the storage layer for a time-series database that ingests 1 million data points per second and must support range queries going back 2 years. Describe compression strategies, write amplification concerns, and compaction trade-offs.",
|
||||
"Explain how LoRA (Low-Rank Adaptation) reduces the number of trainable parameters in fine-tuning. Derive why a weight update matrix can be approximated as a product of two low-rank matrices and discuss what is lost in this approximation.",
|
||||
"A binary tree is given where each node has a value. Write an algorithm to find the maximum path sum between any two nodes (not necessarily leaf nodes). Prove the correctness of your recurrence relation.",
|
||||
"Explain the economic concept of Goodhart's Law and give three examples of how it manifests in AI system evaluation.",
|
||||
"Describe the full lifecycle of a memory allocation in a system using jemalloc or tcmalloc. How do thread-local caches, size classes, and slab allocation interact, and what are the implications for long-running server processes?"
|
||||
]
|
||||
|
||||
ALL_PROMPTS = []
|
||||
for p in EASY_PROMPTS: ALL_PROMPTS.append({"text": p, "difficulty": "EASY"})
|
||||
for p in MEDIUM_PROMPTS: ALL_PROMPTS.append({"text": p, "difficulty": "MEDIUM"})
|
||||
for p in TOUGH_PROMPTS: ALL_PROMPTS.append({"text": p, "difficulty": "TOUGH"})
|
||||
|
||||
# UTILITIES AND DATASET LOADERS
|
||||
class SFTDataset(Dataset):
|
||||
def __init__(self, file_path: str, max_samples: int = -1):
|
||||
self.samples = []
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if 0 < max_samples <= len(self.samples):
|
||||
break
|
||||
self.samples.append(json.loads(line))
|
||||
print(f"Loaded {len(self.samples)} SFT samples from {file_path}")
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
return self.samples[idx]
|
||||
|
||||
def pack_sequences(samples: list[dict], pack_length: int, pad_token_id: int, eos_token_id: int) -> list[dict]:
|
||||
"""Sort and pack short samples into fixed-size bins (FFD packing) to accelerate training."""
|
||||
print(f"Packing sequences into {pack_length}-token bins...")
|
||||
# Sort samples by input_ids length descending
|
||||
indexed_samples = sorted(
|
||||
samples,
|
||||
key=lambda x: len(x["input_ids"]),
|
||||
reverse=True
|
||||
)
|
||||
|
||||
bins: list[list[dict]] = []
|
||||
bin_lengths: list[int] = []
|
||||
|
||||
for sample in indexed_samples:
|
||||
s_len = len(sample["input_ids"])
|
||||
if s_len > pack_length:
|
||||
sample["input_ids"] = sample["input_ids"][:pack_length]
|
||||
sample["loss_mask"] = sample["loss_mask"][:pack_length]
|
||||
s_len = pack_length
|
||||
|
||||
# Try to place sample into an existing bin
|
||||
placed = False
|
||||
for b_idx in range(len(bins)):
|
||||
needed = s_len + (1 if len(bins[b_idx]) > 0 else 0)
|
||||
if bin_lengths[b_idx] + needed <= pack_length:
|
||||
bins[b_idx].append(sample)
|
||||
bin_lengths[b_idx] += needed
|
||||
placed = True
|
||||
break
|
||||
|
||||
if not placed:
|
||||
bins.append([sample])
|
||||
bin_lengths.append(s_len)
|
||||
|
||||
# Convert packed bins to training formats
|
||||
packed_samples = []
|
||||
for b in bins:
|
||||
input_ids = []
|
||||
loss_mask = []
|
||||
for i, sample in enumerate(b):
|
||||
if i > 0:
|
||||
input_ids.append(eos_token_id)
|
||||
loss_mask.append(0) # Mask out the EOS separator token
|
||||
input_ids.extend(sample["input_ids"])
|
||||
loss_mask.extend(sample["loss_mask"])
|
||||
|
||||
real_len = len(input_ids)
|
||||
pad_len = pack_length - real_len
|
||||
if pad_len > 0:
|
||||
input_ids.extend([pad_token_id] * pad_len)
|
||||
loss_mask.extend([0] * pad_len)
|
||||
|
||||
packed_samples.append({
|
||||
"input_ids": torch.tensor(input_ids, dtype=torch.long),
|
||||
"loss_mask": torch.tensor(loss_mask, dtype=torch.long),
|
||||
"attention_mask": torch.cat([
|
||||
torch.ones(real_len, dtype=torch.long),
|
||||
torch.zeros(pad_len, dtype=torch.long)
|
||||
])
|
||||
})
|
||||
|
||||
utilization = sum(bin_lengths) / (len(bins) * pack_length)
|
||||
print(f"Packed {len(samples)} samples into {len(bins)} bins. Utilization: {utilization * 100:.2f}%")
|
||||
return packed_samples
|
||||
|
||||
def collate_sft(batch: list[dict], pad_token_id: int) -> dict:
|
||||
"""Collates batch for standard unpacked training, dynamically padding batch to max length."""
|
||||
max_len = max(len(s["input_ids"]) for s in batch)
|
||||
input_ids_list = []
|
||||
attention_mask_list = []
|
||||
labels_list = []
|
||||
|
||||
for s in batch:
|
||||
ids = s["input_ids"]
|
||||
mask = s["loss_mask"]
|
||||
pad_len = max_len - len(ids)
|
||||
|
||||
padded_ids = ids + [pad_token_id] * pad_len
|
||||
padded_labels = [ids[i] if mask[i] == 1 else -100 for i in range(len(ids))] + [-100] * pad_len
|
||||
|
||||
input_ids_list.append(torch.tensor(padded_ids, dtype=torch.long))
|
||||
attention_mask_list.append(torch.tensor([1] * len(ids) + [0] * pad_len, dtype=torch.long))
|
||||
labels_list.append(torch.tensor(padded_labels, dtype=torch.long))
|
||||
|
||||
return {
|
||||
"input_ids": torch.stack(input_ids_list),
|
||||
"attention_mask": torch.stack(attention_mask_list),
|
||||
"labels": torch.stack(labels_list)
|
||||
}
|
||||
|
||||
def collate_packed(batch: list[dict]) -> dict:
|
||||
"""Collates pre-packed sequence bins by simple stacking."""
|
||||
input_ids = torch.stack([item["input_ids"] for item in batch])
|
||||
attention_mask = torch.stack([item["attention_mask"] for item in batch])
|
||||
loss_mask = torch.stack([item["loss_mask"] for item in batch])
|
||||
|
||||
labels = input_ids.clone()
|
||||
labels = labels.masked_fill(loss_mask == 0, -100)
|
||||
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"labels": labels
|
||||
}
|
||||
|
||||
# PARSING AND MAIN LOGIC
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Clean SFT training and evaluation suite")
|
||||
parser.add_argument("--student_model", type=str, default=cfg.get("model", {}).get("student", "Qwen/Qwen3-1.7B-Base"))
|
||||
parser.add_argument("--tokenizer_model", type=str, default=cfg.get("model", {}).get("tokenizer", "Qwen/Qwen3-1.7B"))
|
||||
parser.add_argument("--data_repo", type=str, default=os.environ.get("QUINTUS_SFT_DATA_REPO"), help="HF dataset repo containing train_sft.jsonl. Optional when data/tokenized/train_sft.jsonl exists.")
|
||||
parser.add_argument("--token", type=str, default=None)
|
||||
parser.add_argument("--trust_remote_code", action="store_true", help="Allow custom code from model/tokenizer repositories.")
|
||||
|
||||
parser.add_argument("--num_epochs", type=int, default=1)
|
||||
parser.add_argument("--learning_rate", type=float, default=2e-5)
|
||||
parser.add_argument("--micro_batch_size", type=int, default=4)
|
||||
parser.add_argument("--grad_accum_steps", type=int, default=2)
|
||||
parser.add_argument("--max_seq_len", type=int, default=4096)
|
||||
parser.add_argument("--sequence_packing", action="store_true", default=True)
|
||||
parser.add_argument("--no_sequence_packing", action="store_false", dest="sequence_packing")
|
||||
|
||||
parser.add_argument("--output_dir", type=str, default="quintus_sft_output")
|
||||
|
||||
parser.add_argument("--run_prompt_suite", action="store_true", default=True)
|
||||
parser.add_argument("--no_prompt_suite", action="store_false", dest="run_prompt_suite")
|
||||
parser.add_argument("--run_gsm8k", action="store_true", default=True)
|
||||
parser.add_argument("--no_gsm8k", action="store_false", dest="run_gsm8k")
|
||||
parser.add_argument("--gsm8k_samples", type=int, default=100)
|
||||
|
||||
parser.add_argument("--optim", type=str, choices=["adamw", "adamw_8bit"], default="adamw")
|
||||
parser.add_argument("--gradient_checkpointing", action="store_true", default=False)
|
||||
parser.add_argument("--load_in_4bit", action="store_true", default=False)
|
||||
parser.add_argument("--use_lora", action="store_true", default=False)
|
||||
parser.add_argument("--lora_r", type=int, default=8)
|
||||
parser.add_argument("--lora_alpha", type=int, default=16)
|
||||
|
||||
parser.add_argument("--push_to_hub", action="store_true", default=False, help="Automatically push fine-tuned model to Hugging Face Hub after training")
|
||||
parser.add_argument("--hub_model_id", type=str, default="iamrahulreddy/Quintus", help="Target Hugging Face Hub repository ID")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
def download_hf_dataset(repo_id: str | None, token: str | None) -> str:
|
||||
print(f"Checking for tokenized dataset in local folders...")
|
||||
local_path = "data/tokenized/train_sft.jsonl"
|
||||
if os.path.exists(local_path):
|
||||
print(f"Found local dataset: {local_path}")
|
||||
return local_path
|
||||
|
||||
if not repo_id:
|
||||
raise ValueError(
|
||||
"No local SFT dataset found at data/tokenized/train_sft.jsonl. "
|
||||
"Pass --data_repo or set QUINTUS_SFT_DATA_REPO."
|
||||
)
|
||||
|
||||
print(f"Local file not found. Pulling from Hugging Face: {repo_id}...")
|
||||
from huggingface_hub import hf_hub_download
|
||||
os.makedirs("data/tokenized", exist_ok=True)
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename="train_sft.jsonl",
|
||||
repo_type="dataset",
|
||||
local_dir="data/tokenized",
|
||||
token=token
|
||||
)
|
||||
# Ensure correct local path layout
|
||||
if os.path.exists(downloaded) and downloaded != local_path:
|
||||
os.rename(downloaded, local_path)
|
||||
print(f"Dataset downloaded to: {local_path}")
|
||||
return local_path
|
||||
|
||||
# DOWNSTREAM EVALUATION CODE
|
||||
def run_prompt_suite(model, tokenizer, device, output_dir: str):
|
||||
print("\n" + "="*70)
|
||||
print("RUNNING QUALITATIVE PROMPT SUITE (50 Prompts)")
|
||||
print("="*70)
|
||||
|
||||
# Compile stop token IDs
|
||||
eos_token_ids = [tokenizer.eos_token_id]
|
||||
for token in ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]:
|
||||
t_id = tokenizer.convert_tokens_to_ids(token)
|
||||
if t_id is not None and t_id != tokenizer.unk_token_id:
|
||||
eos_token_ids.append(t_id)
|
||||
eos_token_ids = list(set(eos_token_ids))
|
||||
|
||||
# Initialize output file
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
out_path = os.path.join(output_dir, f"prompt_suite_eval_{timestamp}.txt")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
f.write("QUINTUS SFT POST-TRAINING PROMPT SUITE\n")
|
||||
f.write(f"Timestamp: {timestamp}\n")
|
||||
f.write("="*72 + "\n\n")
|
||||
f.flush()
|
||||
|
||||
# Set padding side to left for batch generation
|
||||
orig_padding_side = tokenizer.padding_side
|
||||
tokenizer.padding_side = "left"
|
||||
|
||||
batch_size = 16
|
||||
for i in range(0, len(ALL_PROMPTS), batch_size):
|
||||
batch_items = ALL_PROMPTS[i : i + batch_size]
|
||||
|
||||
# Format prompts
|
||||
formatted_prompts = []
|
||||
for item in batch_items:
|
||||
prompt_text = item["text"]
|
||||
if tokenizer.chat_template is not None:
|
||||
prompt_str = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": prompt_text}],
|
||||
tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
else:
|
||||
prompt_str = f"<|im_start|>user\n{prompt_text}<|im_end|>\n<|im_start|>assistant\n"
|
||||
formatted_prompts.append(prompt_str)
|
||||
|
||||
# Tokenize with padding
|
||||
inputs = tokenizer(formatted_prompts, padding=True, return_tensors="pt").to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=2048,
|
||||
do_sample=False, # Greedy for clean, reproducible comparison
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=eos_token_ids
|
||||
)
|
||||
|
||||
# Decode and write results in real-time
|
||||
for idx, item in enumerate(batch_items):
|
||||
input_len = inputs["input_ids"][idx].shape[0]
|
||||
gen_tokens = outputs[idx][input_len:]
|
||||
|
||||
# Slice at the first EOS token
|
||||
eos_indices = []
|
||||
for eos_id in eos_token_ids:
|
||||
indices = (gen_tokens == eos_id).nonzero(as_tuple=True)[0]
|
||||
if len(indices) > 0:
|
||||
eos_indices.append(indices[0].item())
|
||||
if eos_indices:
|
||||
gen_tokens = gen_tokens[:min(eos_indices)]
|
||||
|
||||
response = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
# Log progress
|
||||
global_idx = i + idx + 1
|
||||
print(f"[{global_idx:02d}/50] ({item['difficulty']}) Q: {item['text'][:40]}... -> Answered ({len(gen_tokens)} tokens)")
|
||||
|
||||
# Append directly to output file
|
||||
with open(out_path, "a", encoding="utf-8") as f:
|
||||
f.write(f"[{global_idx:02d}/50] {item['difficulty']}\n")
|
||||
f.write(f"Q: {item['text']}\n\n")
|
||||
f.write(f"Response:\n{response}\n")
|
||||
f.write("\n" + "-"*72 + "\n\n")
|
||||
f.flush()
|
||||
|
||||
# Restore original tokenizer settings
|
||||
tokenizer.padding_side = orig_padding_side
|
||||
print(f"\nPrompt suite evaluation complete. Saved report to: {out_path}\n")
|
||||
|
||||
def extract_gsm8k_answer(text: str) -> str | None:
|
||||
text = text.replace(",", "")
|
||||
match = re.findall(r"The answer is\s*:?\s*(-?\d+)", text, re.IGNORECASE)
|
||||
if match:
|
||||
return match[-1]
|
||||
match = re.findall(r"(-?\d+)", text)
|
||||
if match:
|
||||
return match[-1]
|
||||
return None
|
||||
|
||||
def run_gsm8k_eval(model, tokenizer, device, num_samples: int = 100):
|
||||
print("\n" + "="*70)
|
||||
print(f"RUNNING GSM8K MATH EVALUATION ({num_samples} Samples)")
|
||||
print("="*70)
|
||||
|
||||
from datasets import load_dataset
|
||||
try:
|
||||
dataset = load_dataset("openai/gsm8k", "main", split="test")
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not download GSM8K test set directly: {e}")
|
||||
return
|
||||
|
||||
dataset = dataset.shuffle(seed=42).select(range(min(num_samples, len(dataset))))
|
||||
|
||||
correct = 0
|
||||
total = 0
|
||||
|
||||
for idx, item in enumerate(dataset):
|
||||
question = item["question"]
|
||||
answer = item["answer"]
|
||||
|
||||
target_match = re.search(r"####\s*(-?\d+)", answer)
|
||||
if not target_match:
|
||||
continue
|
||||
target_val = target_match.group(1)
|
||||
|
||||
if tokenizer.chat_template is not None:
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": question + "\nShow your work and conclude with 'The answer is: <number>'."}],
|
||||
tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
else:
|
||||
prompt = f"<|im_start|>user\n{question}\nShow your work and conclude with 'The answer is: <number>'.<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=1024,
|
||||
do_sample=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
gen_tokens = outputs[0][inputs.input_ids.shape[1]:]
|
||||
generated_text = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
pred_val = extract_gsm8k_answer(generated_text)
|
||||
is_match = (pred_val == target_val)
|
||||
|
||||
if is_match:
|
||||
correct += 1
|
||||
total += 1
|
||||
|
||||
# Log sample output periodically
|
||||
if idx % 10 == 0:
|
||||
print(f"\n[GSM8K Sample {idx+1}]")
|
||||
print(f"Q: {question[:80]}...")
|
||||
print(f"A: {generated_text[:120]}... (Target: {target_val} | Pred: {pred_val})")
|
||||
print(f"Match: {is_match}")
|
||||
|
||||
accuracy = (correct / total * 100) if total > 0 else 0
|
||||
print("\n" + "="*70)
|
||||
print(f"GSM8K EVALUATION SUMMARY: {correct}/{total} Correct -> Accuracy: {accuracy:.2f}%")
|
||||
print("="*70 + "\n")
|
||||
|
||||
# TRAINING PIPELINE
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# Propagate HF token to environment for auto-authentication of downstream hub calls
|
||||
try:
|
||||
import huggingface_hub
|
||||
cached_token = huggingface_hub.get_token()
|
||||
except Exception:
|
||||
cached_token = None
|
||||
resolved_token = os.environ.get("HF_TOKEN") or cached_token or args.token
|
||||
if resolved_token:
|
||||
os.environ["HF_TOKEN"] = resolved_token
|
||||
args.token = resolved_token
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"SFT Environment initialized. Target device: {device}")
|
||||
|
||||
# 1. Pull dataset from HF
|
||||
try:
|
||||
dataset_file = download_hf_dataset(args.data_repo, args.token)
|
||||
except ValueError as exc:
|
||||
print(f"Error: {exc}")
|
||||
sys.exit(1)
|
||||
|
||||
# 2. Setup Tokenizer and Model
|
||||
print(f"Loading tokenizer: {args.tokenizer_model}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_model, trust_remote_code=args.trust_remote_code)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
# 4-bit configuration if requested
|
||||
bnb_config = None
|
||||
if args.load_in_4bit:
|
||||
from transformers import BitsAndBytesConfig
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_use_double_quant=True
|
||||
)
|
||||
print("Using 4-bit BitsAndBytes quantization.")
|
||||
|
||||
# Liger Kernel (skipped for 4-bit/PEFT as it can interfere with quantized layers)
|
||||
if not args.load_in_4bit:
|
||||
try:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_qwen3
|
||||
apply_liger_kernel_to_qwen3(
|
||||
rope=True,
|
||||
swiglu=True,
|
||||
rms_norm=True,
|
||||
cross_entropy=False,
|
||||
fused_linear_cross_entropy=False,
|
||||
)
|
||||
print("Liger Kernel optimizations applied successfully.")
|
||||
except ImportError:
|
||||
print("Liger Kernel not installed, skipping optimizations.")
|
||||
|
||||
attn_impl = "sdpa"
|
||||
if device.type == "cuda":
|
||||
try:
|
||||
import flash_attn
|
||||
attn_impl = "flash_attention_2"
|
||||
print("FlashAttention-2 enabled.")
|
||||
except ImportError:
|
||||
print("flash-attn not installed, falling back to SDPA.")
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.student_model,
|
||||
quantization_config=bnb_config,
|
||||
dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
attn_implementation=attn_impl
|
||||
)
|
||||
if not args.load_in_4bit:
|
||||
model = model.to(device)
|
||||
model.config.use_cache = False
|
||||
|
||||
# Wrap with LoRA if requested or required for 4-bit training
|
||||
if args.use_lora or args.load_in_4bit:
|
||||
try:
|
||||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||||
if args.load_in_4bit:
|
||||
model = prepare_model_for_kbit_training(model)
|
||||
|
||||
peft_config = LoraConfig(
|
||||
r=args.lora_r,
|
||||
lora_alpha=args.lora_alpha,
|
||||
target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
print("LoRA adapters successfully attached to target modules.")
|
||||
model.print_trainable_parameters()
|
||||
except ImportError:
|
||||
print("Error: peft not installed. Please run `!pip install -q peft` to use LoRA/QLoRA.")
|
||||
sys.exit(1)
|
||||
|
||||
if args.gradient_checkpointing:
|
||||
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
print("Gradient checkpointing enabled.")
|
||||
|
||||
# 3. Prepare dataset
|
||||
raw_dataset = SFTDataset(dataset_file)
|
||||
|
||||
if args.sequence_packing:
|
||||
packed_samples = pack_sequences(
|
||||
raw_dataset.samples,
|
||||
pack_length=args.max_seq_len,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
train_dataloader = DataLoader(
|
||||
packed_samples,
|
||||
batch_size=args.micro_batch_size,
|
||||
shuffle=True,
|
||||
collate_fn=collate_packed
|
||||
)
|
||||
else:
|
||||
train_dataloader = DataLoader(
|
||||
raw_dataset,
|
||||
batch_size=args.micro_batch_size,
|
||||
shuffle=True,
|
||||
collate_fn=lambda b: collate_sft(b, tokenizer.pad_token_id)
|
||||
)
|
||||
|
||||
# 4. Optimizer and scheduler setup
|
||||
if args.optim == "adamw_8bit":
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=args.learning_rate, weight_decay=0.1)
|
||||
print("Using BitsAndBytes 8-bit AdamW optimizer.")
|
||||
except ImportError:
|
||||
print("Warning: bitsandbytes not installed. Falling back to standard AdamW.")
|
||||
use_fused = (device.type == "cuda")
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0.1, fused=use_fused)
|
||||
else:
|
||||
use_fused = (device.type == "cuda")
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0.1, fused=use_fused)
|
||||
print(f"Using standard AdamW optimizer (fused={use_fused}).")
|
||||
steps_per_epoch = (len(train_dataloader) + args.grad_accum_steps - 1) // args.grad_accum_steps
|
||||
total_steps = steps_per_epoch * args.num_epochs
|
||||
warmup_steps = int(total_steps * 0.05)
|
||||
scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
|
||||
|
||||
# 5. Training Loop
|
||||
print("\n" + "="*70)
|
||||
print(f"STARTING SFT TRAINING (Epochs: {args.num_epochs} | Steps: {total_steps})")
|
||||
print("="*70)
|
||||
|
||||
model.train()
|
||||
step = 0
|
||||
total_tokens_processed = 0
|
||||
t0 = time.time()
|
||||
|
||||
for epoch in range(args.num_epochs):
|
||||
epoch_loss = 0.0
|
||||
for batch_idx, batch in enumerate(train_dataloader):
|
||||
input_ids = batch["input_ids"].to(device)
|
||||
attention_mask = batch["attention_mask"].to(device)
|
||||
labels = batch["labels"].to(device)
|
||||
|
||||
# Accumulate the number of active (non-padded) tokens processed
|
||||
total_tokens_processed += attention_mask.sum().item()
|
||||
|
||||
outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
|
||||
loss = outputs.loss / args.grad_accum_steps
|
||||
loss.backward()
|
||||
|
||||
epoch_loss += loss.item() * args.grad_accum_steps
|
||||
|
||||
if (batch_idx + 1) % args.grad_accum_steps == 0 or (batch_idx + 1) == len(train_dataloader):
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
step += 1
|
||||
|
||||
if step % 5 == 0 or step == total_steps:
|
||||
elapsed = time.time() - t0
|
||||
tokens_per_sec = total_tokens_processed / max(elapsed, 1e-5)
|
||||
print(
|
||||
f"Epoch {epoch+1}/{args.num_epochs} | "
|
||||
f"Step {step}/{total_steps} | "
|
||||
f"Loss: {loss.item() * args.grad_accum_steps:.4f} | "
|
||||
f"LR: {scheduler.get_last_lr()[0]:.2e} | "
|
||||
f"Tokens: {total_tokens_processed} | "
|
||||
f"Speed: {tokens_per_sec:.2f} tokens/s"
|
||||
)
|
||||
|
||||
# 6. Save model weights and tokenizer
|
||||
print(f"\nTraining complete in {time.time() - t0:.1f}s. Saving weights to: {args.output_dir}")
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
if hasattr(model, "merge_and_unload") and not args.load_in_4bit:
|
||||
print("Merging LoRA adapters into base weights...")
|
||||
try:
|
||||
merged_model = model.merge_and_unload()
|
||||
merged_model.save_pretrained(args.output_dir)
|
||||
print("Merged model weights saved successfully.")
|
||||
except Exception as e:
|
||||
print(f"Failed to merge and unload: {e}. Saving adapter weights only.")
|
||||
model.save_pretrained(args.output_dir)
|
||||
else:
|
||||
model.save_pretrained(args.output_dir)
|
||||
tokenizer.save_pretrained(args.output_dir)
|
||||
print("Weights and configuration saved successfully.")
|
||||
|
||||
# 7. SFT Downstream Evaluations
|
||||
model.eval()
|
||||
|
||||
if args.run_prompt_suite:
|
||||
run_prompt_suite(model, tokenizer, device, args.output_dir)
|
||||
|
||||
if args.run_gsm8k:
|
||||
run_gsm8k_eval(model, tokenizer, device, num_samples=args.gsm8k_samples)
|
||||
|
||||
if args.push_to_hub:
|
||||
print(f"\nUploading fine-tuned model and tokenizer to Hugging Face Hub: {args.hub_model_id}...")
|
||||
try:
|
||||
from huggingface_hub import create_repo, HfApi
|
||||
token_val = args.token or os.environ.get("HF_TOKEN")
|
||||
create_repo(repo_id=args.hub_model_id, token=token_val, exist_ok=True)
|
||||
|
||||
api = HfApi()
|
||||
api.upload_folder(
|
||||
folder_path=args.output_dir,
|
||||
repo_id=args.hub_model_id,
|
||||
repo_type="model",
|
||||
token=token_val
|
||||
)
|
||||
print("Successfully uploaded model and tokenizer to Hugging Face Hub!")
|
||||
except Exception as hub_err:
|
||||
print(f"Failed to push to Hub: {hub_err}")
|
||||
|
||||
print("Pipeline Execution Complete. Model is ready.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user