evalstate HF Staff commited on
Commit
e78dff9
·
verified ·
1 Parent(s): c867455

Upload demo_train.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. demo_train.py +66 -0
demo_train.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = [
3
+ # "trl>=0.12.0",
4
+ # "peft>=0.7.0",
5
+ # "transformers>=4.36.0",
6
+ # "datasets>=2.14.0",
7
+ # ]
8
+ # ///
9
+
10
+ from datasets import load_dataset
11
+ from peft import LoraConfig
12
+ from trl import SFTTrainer, SFTConfig
13
+ import os
14
+
15
+ print("🚀 Starting quick demo training...")
16
+
17
+ # Load a tiny subset of data
18
+ dataset = load_dataset("trl-lib/Capybara", split="train[:50]")
19
+ print(f"✅ Dataset loaded: {len(dataset)} examples")
20
+
21
+ # Training configuration
22
+ config = SFTConfig(
23
+ # CRITICAL: Hub settings
24
+ output_dir="demo-qwen-sft",
25
+ push_to_hub=True,
26
+ hub_model_id="evalstate/demo-qwen-sft",
27
+
28
+ # Quick demo settings
29
+ max_steps=10, # Just 10 steps for quick demo
30
+ per_device_train_batch_size=2,
31
+ learning_rate=2e-5,
32
+
33
+ # Logging
34
+ logging_steps=2,
35
+ save_strategy="no", # Don't save checkpoints for quick demo
36
+
37
+ # Optimization
38
+ warmup_steps=2,
39
+ lr_scheduler_type="constant",
40
+ )
41
+
42
+ # LoRA configuration (reduces memory)
43
+ peft_config = LoraConfig(
44
+ r=16,
45
+ lora_alpha=32,
46
+ lora_dropout=0.05,
47
+ bias="none",
48
+ task_type="CAUSAL_LM",
49
+ target_modules=["q_proj", "v_proj"],
50
+ )
51
+
52
+ # Initialize and train
53
+ trainer = SFTTrainer(
54
+ model="Qwen/Qwen2.5-0.5B",
55
+ train_dataset=dataset,
56
+ args=config,
57
+ peft_config=peft_config,
58
+ )
59
+
60
+ print("🏃 Training for 10 steps...")
61
+ trainer.train()
62
+
63
+ print("💾 Pushing to Hub...")
64
+ trainer.push_to_hub()
65
+
66
+ print("✅ Complete! Model at: https://huggingface.co/evalstate/demo-qwen-sft")