ferrotensor_trainer/
demo.rs

1//! Simulated fine-tuning demo for the CLI (`ferrotensor-cli --demo`).
2//!
3//! Displays a training-style UI with a synthetic loss curve. Does **not** run
4//! GPU backward passes or update model weights.
5
6use crate::scaler::GradScaler;
7use ferrotensor_autograd::{CosineAnnealingLR, LRScheduler};
8use ferrotensor_backend::{ComputeEngine, WgpuContext};
9use std::sync::Arc;
10use std::time::Instant;
11
12/// Demo configuration for fine-tuning
13pub struct DemoConfig {
14    /// Model name shown in the banner.
15    pub model_name: String,
16    /// Displayed batch size.
17    pub batch_size: usize,
18    /// Starting learning rate of the simulated schedule.
19    pub learning_rate: f32,
20    /// Simulated epochs.
21    pub num_epochs: usize,
22    /// Displayed sequence length.
23    pub max_seq_len: usize,
24    /// Show the loss-scaler line.
25    pub use_fp16: bool,
26    /// Show LoRA settings.
27    pub use_lora: bool,
28    /// Displayed LoRA rank.
29    pub lora_rank: usize,
30}
31
32impl Default for DemoConfig {
33    fn default() -> Self {
34        Self {
35            model_name: "TinyLlama-1.1B".to_string(),
36            batch_size: 4,
37            learning_rate: 1e-4,
38            num_epochs: 3,
39            max_seq_len: 512,
40            use_fp16: true,
41            use_lora: true,
42            lora_rank: 8,
43        }
44    }
45}
46
47/// Print a styled banner for the demo
48pub fn print_banner() {
49    println!();
50    println!("======================================================================");
51    println!("                                                                      ");
52    println!("   FERROTENSOR - GPU-Accelerated ML Framework                               ");
53    println!("   Written in Pure Rust                                               ");
54    println!("                                                                      ");
55    println!("======================================================================");
56    println!();
57}
58
59/// Print GPU information
60pub fn print_gpu_info(ctx: &WgpuContext) {
61    let info = ctx.adapter.get_info();
62    println!("[GPU] Information:");
63    println!("      Name:    {}", info.name);
64    println!("      Backend: {:?}", info.backend);
65    println!("      Type:    {:?}", info.device_type);
66    println!();
67}
68
69/// Print configuration
70pub fn print_config(config: &DemoConfig) {
71    println!("[CONFIG] Training Configuration:");
72    println!("         Model:      {}", config.model_name);
73    println!("         Batch Size: {}", config.batch_size);
74    println!("         LR:         {:.0e}", config.learning_rate);
75    println!("         Epochs:     {}", config.num_epochs);
76    println!("         Seq Length: {}", config.max_seq_len);
77    println!(
78        "         FP16:       {}",
79        if config.use_fp16 { "Yes" } else { "No" }
80    );
81    println!(
82        "         LoRA:       {} (rank={})",
83        if config.use_lora { "Yes" } else { "No" },
84        config.lora_rank
85    );
86    println!();
87}
88
89/// Print training progress
90pub fn print_progress(
91    epoch: usize,
92    step: usize,
93    total_steps: usize,
94    loss: f32,
95    lr: f32,
96    tokens_per_sec: f32,
97) {
98    let progress = (step as f32 / total_steps as f32 * 100.0) as usize;
99    let bar_width = 30;
100    let filled = progress * bar_width / 100;
101    let bar: String = "=".repeat(filled) + &"-".repeat(bar_width - filled);
102
103    print!(
104        "\r[TRAIN] Epoch {}: [{}] {}% | Loss: {:.4} | LR: {:.2e} | {:.0} tok/s",
105        epoch, bar, progress, loss, lr, tokens_per_sec
106    );
107    std::io::Write::flush(&mut std::io::stdout()).ok();
108}
109
110/// Print epoch summary
111pub fn print_epoch_summary(epoch: usize, avg_loss: f32, duration_secs: f32) {
112    println!();
113    println!(
114        "        Epoch {} Complete | Avg Loss: {:.4} | Time: {:.1}s",
115        epoch, avg_loss, duration_secs
116    );
117}
118
119/// Sample training data for demo
120pub fn create_demo_data() -> Vec<(Vec<u32>, Vec<u32>)> {
121    let samples = [
122        "Ferrotensor is a GPU-accelerated ML framework written in Rust.",
123        "Flash Attention provides memory-efficient attention computation.",
124        "LoRA enables parameter-efficient fine-tuning of large models.",
125        "Mixed precision training uses FP16 for speed with FP32 for stability.",
126        "The Metal backend runs on Apple Silicon GPUs.",
127        "Transformers use self-attention to process sequences.",
128        "Gradient checkpointing trades compute for memory savings.",
129        "WGSL is the shader language for WebGPU compute kernels.",
130    ];
131
132    samples
133        .iter()
134        .map(|s| {
135            let tokens: Vec<u32> = s.bytes().map(|b| b as u32).collect();
136            let input = tokens[..tokens.len() - 1].to_vec();
137            let target = tokens[1..].to_vec();
138            (input, target)
139        })
140        .collect()
141}
142
143/// Run the demo training loop
144pub fn run_demo(config: DemoConfig) {
145    print_banner();
146
147    // Initialize GPU
148    println!("[INIT] Initializing GPU...");
149    let ctx = Arc::new(pollster::block_on(async { WgpuContext::new().await }));
150    let _engine = ComputeEngine::new(&ctx);
151    print_gpu_info(&ctx);
152
153    print_config(&config);
154
155    // Create demo data
156    println!("[DATA] Loading Training Data...");
157    let data = create_demo_data();
158    println!("       {} training samples loaded", data.len());
159    println!();
160
161    // Initialize training components
162    println!("[INIT] Initializing Training Components...");
163
164    // Create GradScaler for FP16
165    let scaler: Option<GradScaler> = if config.use_fp16 {
166        println!("       GradScaler initialized (FP16 training)");
167        Some(GradScaler::new())
168    } else {
169        None
170    };
171
172    // Create LR scheduler
173    let mut scheduler = CosineAnnealingLR::new(
174        config.learning_rate,
175        config.learning_rate * 0.01,
176        config.num_epochs * data.len() / config.batch_size,
177    );
178    println!("       CosineAnnealingLR scheduler ready");
179
180    // Create optimizer
181    let lr = scheduler.get_lr();
182    println!("       AdamW optimizer (lr={:.2e})", lr);
183    println!();
184
185    // Training loop
186    println!("[DEMO] Simulated training loop (no GPU weight updates)");
187    println!("{}", "-".repeat(70));
188
189    let total_steps = (data.len() / config.batch_size) * config.num_epochs;
190    let mut global_step = 0;
191
192    for epoch in 1..=config.num_epochs {
193        let epoch_start = Instant::now();
194        let mut epoch_loss = 0.0;
195        let mut num_batches = 0;
196
197        for batch_idx in (0..data.len()).step_by(config.batch_size) {
198            let step_start = Instant::now();
199
200            let batch_end = (batch_idx + config.batch_size).min(data.len());
201            let batch_samples = &data[batch_idx..batch_end];
202
203            // Simulated loss
204            let loss =
205                2.5 / (1.0 + (global_step as f32 * 0.1)) + 0.1 * (global_step as f32).sin().abs();
206
207            epoch_loss += loss;
208            num_batches += 1;
209            global_step += 1;
210
211            let batch_tokens = batch_samples.iter().map(|(i, _)| i.len()).sum::<usize>();
212            let step_time = step_start.elapsed().as_secs_f32().max(0.001);
213            let tokens_per_sec = batch_tokens as f32 / step_time;
214
215            scheduler.step();
216            let lr = scheduler.get_lr();
217
218            print_progress(epoch, global_step, total_steps, loss, lr, tokens_per_sec);
219
220            std::thread::sleep(std::time::Duration::from_millis(100));
221        }
222
223        let epoch_duration = epoch_start.elapsed().as_secs_f32();
224        let avg_loss = epoch_loss / num_batches as f32;
225        print_epoch_summary(epoch, avg_loss, epoch_duration);
226    }
227
228    println!("{}", "-".repeat(70));
229    println!();
230
231    // Print final summary
232    println!("[DONE] Demo complete (simulated — no weights were trained)");
233    println!();
234    println!("[STATS] Final Statistics:");
235    println!("        Total Steps:    {}", global_step);
236    println!("        Final Loss:     {:.4}", 0.15);
237    println!("        Final LR:       {:.2e}", scheduler.get_lr());
238    if config.use_fp16 {
239        if let Some(ref s) = scaler {
240            println!("        GradScaler:     scale={:.0}", s.scale_factor());
241        }
242    }
243    println!();
244
245    println!("[READY] Model Ready for Inference!");
246    println!();
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn test_demo_data() {
255        let data = create_demo_data();
256        assert!(!data.is_empty());
257    }
258}
ferrotensor_trainer/demo.rs

ferrotensor_trainer/
demo.rs
