69 lines
1.8 KiB
Python
69 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Full fine-tune gemma-3-270m to reconstruct prompts from model outputs."""
|
|
|
|
from datasets import load_from_disk
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
|
|
|
|
# --- Configuration ---
|
|
MODEL_NAME = "google/gemma-3-270m"
|
|
DATASET_PATH = "tokenized-dataset-plain"
|
|
OUTPUT_DIR = "checkpoints-270m-ft"
|
|
MAX_SEQ_LENGTH = 2048
|
|
|
|
# --- Training ---
|
|
BATCH_SIZE = 2
|
|
GRAD_ACCUM = 8
|
|
LEARNING_RATE = 5e-5
|
|
NUM_EPOCHS = 1
|
|
WARMUP_STEPS = 100
|
|
LOGGING_STEPS = 1
|
|
SAVE_STEPS = 100
|
|
|
|
|
|
def main():
|
|
print("Loading pre-tokenized dataset...")
|
|
dataset = load_from_disk(DATASET_PATH)
|
|
print(f"Training examples: {len(dataset)}")
|
|
|
|
print(f"Loading model: {MODEL_NAME}")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_NAME,
|
|
torch_dtype="bfloat16",
|
|
device_map="auto",
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
|
|
trainer = Trainer(
|
|
model=model,
|
|
processing_class=tokenizer,
|
|
train_dataset=dataset,
|
|
data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True),
|
|
args=TrainingArguments(
|
|
output_dir=OUTPUT_DIR,
|
|
per_device_train_batch_size=BATCH_SIZE,
|
|
gradient_accumulation_steps=GRAD_ACCUM,
|
|
learning_rate=LEARNING_RATE,
|
|
num_train_epochs=NUM_EPOCHS,
|
|
warmup_steps=WARMUP_STEPS,
|
|
logging_steps=LOGGING_STEPS,
|
|
save_steps=SAVE_STEPS,
|
|
bf16=True,
|
|
seed=42,
|
|
report_to="wandb",
|
|
logging_strategy="steps",
|
|
gradient_checkpointing=True,
|
|
),
|
|
)
|
|
|
|
print("Training...")
|
|
trainer.train()
|
|
|
|
print("Saving model...")
|
|
trainer.save_model(f"{OUTPUT_DIR}/final")
|
|
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
|
|
print("Done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|