File size: 3,824 Bytes
c1f5657 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | #!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: CPU-Only Training Script (Gradient Checkpointing + LoRA)
# WARNING: This will be VERY SLOW on CPU (~10-100x slower than GPU)
# For 7B models on CPU: only for small experiments or 1B models
# Recommended: use google/gemma-2b or microsoft/phi-2 for CPU training
# ============================================================
import os, warnings, sys
warnings.filterwarnings('ignore')
from transformers import (
AutoTokenizer, AutoModelForCausalLM,
TrainingArguments, set_seed
)
from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
from datasets import load_dataset
import torch
MODEL_ID = 'microsoft/phi-2' # 2.7B = smallest model that trains on CPU RAM
HF_REPO = 'shaikhsalman/zabaanai-v2-cpu'
DATA_PATH = '/app/data/processed/zabaanai_v2_sft_pakistan.jsonl'
OUTPUT_DIR = '/app/output/zabaanai-v2-cpu'
LORA_RANK = 16
SEQ_LEN = 512 # reduced for CPU RAM
BS = 1
GRAD_ACCUM = 32 # effective batch = 32
EPOCHS = 2
LR = 1e-3 # LoRA default lr
if not os.path.exists(DATA_PATH):
print(f'Data not found: {DATA_PATH} — run 01_curate_sft_data.py first!')
sys.exit(1)
set_seed(42)
print('Loading tokenizer...')
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
print('Loading dataset...')
ds = load_dataset('json', data_files=DATA_PATH, split='train')
ds = ds.train_test_split(test_size=0.02, seed=42)
print(f' Train: {len(ds[\"train\"])} | Eval: {len(ds[\"test\"])}')
def preprocess(example):
text = example.get('text', '')
enc = tokenizer(text, max_length=SEQ_LEN, truncation=True,
padding='max_length', return_tensors=None)
enc['labels'] = enc['input_ids'].copy()
return enc
print('Tokenizing...')
tokenized = ds.map(preprocess, remove_columns=ds.column_names, num_proc=2)
print('Loading model with LoRA...')
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map='cpu',
torch_dtype=torch.float16,
offload_folder='/tmp/offload',
trust_remote_code=True,
)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False
lora_config = LoraConfig(
r=LORA_RANK, lora_alpha=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=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
from trl import SFTTrainer
trainer = SFTTrainer(
model=model,
args=TrainingArguments(
output_dir=OUTPUT_DIR,
per_device_train_batch_size=BS,
per_device_eval_batch_size=BS,
gradient_accumulation_steps=GRAD_ACCUM,
num_train_epochs=EPOCHS,
learning_rate=LR,
bf16=False, fp16=True,
logging_dir=f'{OUTPUT_DIR}/logs',
logging_strategy='steps', logging_steps=10,
save_strategy='steps', save_steps=200,
eval_steps=200, eval_strategy='steps',
save_total_limit=2,
report_to=['none'],
dataloader_num_workers=1,
seed=42,
max_grad_norm=1.0,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={'use_reentrant': False},
),
train_dataset=tokenized['train'],
eval_dataset=tokenized['test'],
data_collator=None,
max_seq_length=SEQ_LEN,
tokenizer=tokenizer,
dataset_text_field='text',
)
print('Starting CPU training (WARNING: this will take many hours/days)...')
print(f' Effective batch: {BS * GRAD_ACCUM}, steps: ~{len(tokenized[\"train\"]) // GRAD_ACCUM * EPOCHS}')
trainer.train()
trainer.save_model(f'{OUTPUT_DIR}/final')
print('CPU training complete!')
|