Ambiakshi SLM Forge Studio

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer

# 1. Load Google Gemma 4 in Unquantized FP16
model_id = "google/gemma-4"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)
model.gradient_checkpointing_enable()

# 2. Configure Low-Rank Adaptation (LoRA)
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_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)

# 3. Training Arguments for 8,192-Token Sequences
training_args = TrainingArguments(
    output_dir="./gemma4-slm-checkpoint",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_steps=500,
    fp16=True,
    logging_steps=10,
    save_strategy="steps",
    save_steps=100
)

# 4. Initialize SFTTrainer with 8k Context Window
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    max_seq_length=8192,
    args=training_args
)
trainer.train()