Views
No views yet
1import torch
2from transformers import (
3 pipeline,
4 AutoModelForCausalLM,
5 AutoTokenizer,
6)
7from datasets import load_dataset
8from trl import SFTTrainer, SFTConfig
9
10# Add this to monitor MPS memory usage
11def print_mps_memory():
12 if torch.backends.mps.is_available():
13 print(f"MPS allocated: {torch.mps.current_allocated_memory() / 1024**3:.2f} GB")
14 print(f"MPS cached: {torch.mps.driver_allocated_memory() / 1024**3:.2f} GB")
15
16# Call this periodically during training
17print_mps_memory()
18
19# Check if MPS is available
20if torch.backends.mps.is_available():
21 device = torch.device("mps")
22 print("MPS device found.")
23else:
24 device = torch.device("cpu")
25 print("MPS device not found, using CPU.")
26
27tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
28tokenizer.pad_token = tokenizer.eos_token
29tokenizer.padding_side = "right"
30
31model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
32model = model.to(device) # Move model to MPS
33
34ds = load_dataset("gofilipa/aclu_transgender")
35
36# Limit dataset size for testing
37train_dataset = ds['train'].select(range(min(2000, len(ds['train'])))) # Use only first 2000 samples
38
39# Clear memory first
40if torch.backends.mps.is_available():
41 torch.mps.empty_cache()
42
43# Reduce training parameters for lower memory usage
44training_params = SFTConfig(
45 output_dir="../checkpoints",
46 per_device_train_batch_size=1, # Keep at 1
47 per_device_eval_batch_size=1,
48 gradient_accumulation_steps=2, # Reduce from 4 to 2
49 num_train_epochs=3, # slowly increased as memory allows, from 1-3
50 learning_rate=2e-4,
51 weight_decay=0.001,
52 dataset_text_field="text", # Fixed: removed [:400]
53 report_to="none",
54 bf16=False,
55 fp16=False,
56 dataloader_pin_memory=False,
57 remove_unused_columns=False,
58 max_seq_length=512, # Add this to limit sequence length
59 gradient_checkpointing=True, # Add this to save memory
60)
61
62trainer = SFTTrainer(
63 model = model,
64 train_dataset = train_dataset,
65 processing_class = tokenizer,
66 args = training_params
67)
68
69trainer.train()
70