mirror of
https://github.com/zenlm/zen-identity.git
synced 2026-10-30 14:57:57 +00:00
zeekay
63ce2c4931
- Update pyproject.toml URLs to zen-identity repo - Fix ZEN_MODELS → ZEN_CODER_MODELS imports in benchmark.py and trainer.py - Simplify train_4b.py to match pattern of flash/max scripts - Rewrite train_4b_simple.py for MLX-only training - Remove hardcoded local paths, use HuggingFace dataset zenlm/zen-identity - Clean up __init__.py documentation
115 lines
3.2 KiB
Python
115 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Train zen-coder-4b (Qwen3-Coder-4B based, 4B params)
|
|
|
|
Usage:
|
|
python trainer/train_4b.py
|
|
python trainer/train_4b.py --lora-rank 64 --epochs 3
|
|
"""
|
|
|
|
import argparse
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, BitsAndBytesConfig
|
|
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
|
from datasets import load_dataset
|
|
|
|
BASE_MODEL = "Qwen/Qwen3-Coder-4B"
|
|
OUTPUT_DIR = "./output/zen-coder-4b-lora"
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--epochs", type=int, default=3)
|
|
parser.add_argument("--batch-size", type=int, default=4)
|
|
parser.add_argument("--lr", type=float, default=2e-4)
|
|
parser.add_argument("--lora-rank", type=int, default=64)
|
|
parser.add_argument("--max-seq-length", type=int, default=4096)
|
|
args = parser.parse_args()
|
|
|
|
print("=" * 60)
|
|
print("zen-coder-4b Training")
|
|
print(f"Base: {BASE_MODEL}")
|
|
print(f"Architecture: 4B dense")
|
|
print("=" * 60)
|
|
|
|
# 4-bit quantization
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
)
|
|
|
|
print("Loading model...")
|
|
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
BASE_MODEL,
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
model = prepare_model_for_kbit_training(model)
|
|
|
|
lora_config = LoraConfig(
|
|
r=args.lora_rank,
|
|
lora_alpha=args.lora_rank * 2,
|
|
lora_dropout=0.05,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
|
|
print("Loading dataset...")
|
|
dataset = load_dataset("zenlm/zen-identity", split="train")
|
|
|
|
def format_example(example):
|
|
# Qwen3 chat format
|
|
text = f"<|im_start|>user\n{example['instruction']}<|im_end|>\n<|im_start|>assistant\n{example['output']}<|im_end|>"
|
|
return {"text": text}
|
|
|
|
dataset = dataset.map(format_example)
|
|
|
|
def tokenize(examples):
|
|
return tokenizer(
|
|
examples["text"],
|
|
truncation=True,
|
|
max_length=args.max_seq_length,
|
|
padding="max_length",
|
|
)
|
|
|
|
tokenized = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
|
|
|
|
training_args = TrainingArguments(
|
|
output_dir=OUTPUT_DIR,
|
|
num_train_epochs=args.epochs,
|
|
per_device_train_batch_size=args.batch_size,
|
|
gradient_accumulation_steps=4,
|
|
learning_rate=args.lr,
|
|
warmup_ratio=0.03,
|
|
logging_steps=10,
|
|
save_steps=100,
|
|
bf16=True,
|
|
optim="paged_adamw_8bit",
|
|
report_to="none",
|
|
)
|
|
|
|
trainer = Trainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=tokenized,
|
|
)
|
|
|
|
print("Training...")
|
|
trainer.train()
|
|
|
|
print(f"Saving to {OUTPUT_DIR}")
|
|
model.save_pretrained(OUTPUT_DIR)
|
|
tokenizer.save_pretrained(OUTPUT_DIR)
|
|
print("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|