entrenar 0.7.8

Training & Optimization library with autograd, LoRA, quantization, and model merging
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
//! Golden trace baselines for pruning module (PMAT QA Phase 6)
//!
//! This module defines expected performance characteristics and trace baselines
//! for the pruning operations. These serve as regression tests for performance
//! and can be verified with Renacer tracing when enabled.
//!
//! # Toyota Way: Hansei (Reflection)
//! Golden traces capture expected behavior for continuous improvement.

#![allow(unreachable_pub)]

use serde::{Deserialize, Serialize};

/// Maximum allowed latency for importance computation per layer (milliseconds)
const IMPORTANCE_LATENCY_MAX_MS: u64 = 5000;
/// Maximum allowed syscall/span count during calibration
const CALIBRATION_SYSCALL_BUDGET: u64 = 5000;

/// Performance assertion for a pruning operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PruningAssertion {
    /// Name of the assertion
    pub name: &'static str,
    /// Type of assertion (latency, memory, etc.)
    pub assertion_type: AssertionType,
    /// Maximum allowed value
    pub max_value: u64,
    /// Whether violation causes test failure
    pub fail_on_violation: bool,
    /// Whether this assertion is enabled
    pub enabled: bool,
}

/// Type of performance assertion.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AssertionType {
    /// Maximum duration in milliseconds
    Latency,
    /// Maximum memory usage in bytes
    Memory,
    /// Maximum number of spans/calls
    SpanCount,
    /// Anti-pattern detection threshold (0.0-1.0 as percentage * 100)
    AntiPattern,
}

/// Golden trace baseline for a pruning schedule.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScheduleGoldenTrace {
    /// Schedule type name
    pub schedule_type: &'static str,
    /// Expected steps where pruning should trigger
    pub expected_prune_steps: Vec<usize>,
    /// Expected sparsity at each step (step, sparsity)
    pub expected_sparsity_curve: Vec<(usize, f32)>,
}

/// Golden trace baseline for a pruning configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConfigGoldenTrace {
    /// Configuration identifier
    pub config_id: &'static str,
    /// Expected calibration requirement
    pub requires_calibration: bool,
    /// Expected validation result (true = valid)
    pub expected_valid: bool,
}

/// Collection of all golden trace baselines for pruning.
pub struct PruningGoldenTraces;

impl PruningGoldenTraces {
    /// Get performance assertions for pruning operations.
    pub fn performance_assertions() -> Vec<PruningAssertion> {
        vec![
            // Importance computation latency
            PruningAssertion {
                name: "pruning_importance_latency",
                assertion_type: AssertionType::Latency,
                max_value: IMPORTANCE_LATENCY_MAX_MS, // <5s per layer
                fail_on_violation: true,
                enabled: true,
            },
            // Mask generation latency
            PruningAssertion {
                name: "pruning_mask_generation_latency",
                assertion_type: AssertionType::Latency,
                max_value: 1000, // <1s
                fail_on_violation: true,
                enabled: true,
            },
            // Memory budget
            PruningAssertion {
                name: "pruning_memory_budget",
                assertion_type: AssertionType::Memory,
                max_value: 2_147_483_648, // 2GB
                fail_on_violation: true,
                enabled: true,
            },
            // Calibration syscall budget
            PruningAssertion {
                name: "calibration_syscall_budget",
                assertion_type: AssertionType::SpanCount,
                max_value: CALIBRATION_SYSCALL_BUDGET,
                fail_on_violation: false, // Warning only
                enabled: true,
            },
            // Redundant computation detection
            PruningAssertion {
                name: "detect_redundant_computation",
                assertion_type: AssertionType::AntiPattern,
                max_value: 70, // 70% threshold
                fail_on_violation: false,
                enabled: true,
            },
            // Memory thrashing detection
            PruningAssertion {
                name: "detect_memory_thrashing",
                assertion_type: AssertionType::AntiPattern,
                max_value: 80, // 80% threshold
                fail_on_violation: false,
                enabled: true,
            },
        ]
    }

    /// Get golden traces for schedule types.
    pub fn schedule_traces() -> Vec<ScheduleGoldenTrace> {
        vec![
            // OneShot schedule at step 1000
            ScheduleGoldenTrace {
                schedule_type: "oneshot_1000",
                expected_prune_steps: vec![1000],
                expected_sparsity_curve: vec![(0, 0.0), (999, 0.0), (1000, 1.0), (1001, 1.0)],
            },
            // Gradual schedule 0-100 with freq=10, sparsity 0.0->0.5
            ScheduleGoldenTrace {
                schedule_type: "gradual_0_100_50pct",
                expected_prune_steps: vec![0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
                expected_sparsity_curve: vec![
                    (0, 0.0),
                    (25, 0.125),
                    (50, 0.25),
                    (75, 0.375),
                    (100, 0.5),
                ],
            },
            // Cubic schedule 0-100, sparsity 0.0->0.5
            // Formula: s_t = s_f * (1 - (1 - t/T)^3)
            ScheduleGoldenTrace {
                schedule_type: "cubic_0_100_50pct",
                expected_prune_steps: (0..=100).collect(),
                expected_sparsity_curve: vec![
                    (0, 0.0),
                    (25, 0.2890625), // 0.5 * (1 - 0.75^3) = 0.5 * 0.578125
                    (50, 0.4375),    // 0.5 * (1 - 0.5^3) = 0.5 * 0.875
                    (75, 0.4921875), // 0.5 * (1 - 0.25^3) = 0.5 * 0.984375
                    (100, 0.5),
                ],
            },
        ]
    }

    /// Get golden traces for configurations.
    pub fn config_traces() -> Vec<ConfigGoldenTrace> {
        vec![
            ConfigGoldenTrace {
                config_id: "default",
                requires_calibration: false, // Magnitude doesn't require calibration
                expected_valid: true,
            },
            ConfigGoldenTrace {
                config_id: "wanda_nm24",
                requires_calibration: true,
                expected_valid: true,
            },
            ConfigGoldenTrace {
                config_id: "sparsegpt_unstructured",
                requires_calibration: true,
                expected_valid: true,
            },
            ConfigGoldenTrace {
                config_id: "invalid_nm_n_gte_m",
                requires_calibration: false,
                expected_valid: false,
            },
        ]
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::prune::{PruneMethod, PruningConfig, PruningSchedule, SparsityPatternConfig};

    // =========================================================================
    // Golden Trace Verification Tests
    // =========================================================================

    #[test]
    fn test_oneshot_schedule_matches_golden_trace() {
        // TEST_ID: GOLD-001
        // Verify OneShot schedule behavior matches golden trace
        let schedule = PruningSchedule::OneShot { step: 1000 };
        let golden = PruningGoldenTraces::schedule_traces()
            .into_iter()
            .find(|t| t.schedule_type == "oneshot_1000")
            .expect("Golden trace not found");

        for (step, expected) in &golden.expected_sparsity_curve {
            let actual = schedule.sparsity_at_step(*step);
            assert!(
                (actual - expected).abs() < 1e-6,
                "GOLD-001 FALSIFIED: OneShot at step {step} expected {expected}, got {actual}"
            );
        }

        // Verify prune steps
        for step in &golden.expected_prune_steps {
            assert!(
                schedule.should_prune_at_step(*step),
                "GOLD-001 FALSIFIED: OneShot should prune at step {step}"
            );
        }
    }

    #[test]
    fn test_gradual_schedule_matches_golden_trace() {
        // TEST_ID: GOLD-002
        // Verify Gradual schedule behavior matches golden trace
        let schedule = PruningSchedule::Gradual {
            start_step: 0,
            end_step: 100,
            initial_sparsity: 0.0,
            final_sparsity: 0.5,
            frequency: 10,
        };
        let golden = PruningGoldenTraces::schedule_traces()
            .into_iter()
            .find(|t| t.schedule_type == "gradual_0_100_50pct")
            .expect("Golden trace not found");

        for (step, expected) in &golden.expected_sparsity_curve {
            let actual = schedule.sparsity_at_step(*step);
            assert!(
                (actual - expected).abs() < 1e-6,
                "GOLD-002 FALSIFIED: Gradual at step {step} expected {expected}, got {actual}"
            );
        }

        // Verify prune steps count matches
        assert_eq!(
            schedule.num_pruning_steps(),
            golden.expected_prune_steps.len(),
            "GOLD-002 FALSIFIED: Gradual num_pruning_steps mismatch"
        );
    }

    #[test]
    fn test_cubic_schedule_matches_golden_trace() {
        // TEST_ID: GOLD-003
        // Verify Cubic schedule behavior matches golden trace
        let schedule = PruningSchedule::Cubic { start_step: 0, end_step: 100, final_sparsity: 0.5 };
        let golden = PruningGoldenTraces::schedule_traces()
            .into_iter()
            .find(|t| t.schedule_type == "cubic_0_100_50pct")
            .expect("Golden trace not found");

        for (step, expected) in &golden.expected_sparsity_curve {
            let actual = schedule.sparsity_at_step(*step);
            assert!(
                (actual - expected).abs() < 1e-6,
                "GOLD-003 FALSIFIED: Cubic at step {step} expected {expected}, got {actual}"
            );
        }
    }

    #[test]
    fn test_config_calibration_matches_golden_trace() {
        // TEST_ID: GOLD-004
        // Verify config calibration requirements match golden traces
        let configs = [
            ("default", PruningConfig::default()),
            (
                "wanda_nm24",
                PruningConfig::new()
                    .with_method(PruneMethod::Wanda)
                    .with_pattern(SparsityPatternConfig::nm_2_4()),
            ),
            (
                "sparsegpt_unstructured",
                PruningConfig::new()
                    .with_method(PruneMethod::SparseGpt)
                    .with_pattern(SparsityPatternConfig::Unstructured),
            ),
        ];

        for (id, config) in &configs {
            let golden = PruningGoldenTraces::config_traces()
                .into_iter()
                .find(|t| t.config_id == *id)
                .expect("Golden trace not found");

            assert_eq!(
                config.requires_calibration(),
                golden.requires_calibration,
                "GOLD-004 FALSIFIED: Config {id} calibration requirement mismatch"
            );
        }
    }

    #[test]
    fn test_config_validation_matches_golden_trace() {
        // TEST_ID: GOLD-005
        // Verify config validation matches golden traces
        let valid_config = PruningConfig::default();
        assert!(
            valid_config.validate().is_ok(),
            "GOLD-005 FALSIFIED: Default config should be valid"
        );

        let invalid_config =
            PruningConfig::new().with_pattern(SparsityPatternConfig::NM { n: 5, m: 4 });
        assert!(
            invalid_config.validate().is_err(),
            "GOLD-005 FALSIFIED: Config with n >= m should be invalid"
        );
    }

    #[test]
    fn test_performance_assertions_defined() {
        // TEST_ID: GOLD-010
        // Verify all expected performance assertions are defined
        let assertions = PruningGoldenTraces::performance_assertions();

        let expected_names = [
            "pruning_importance_latency",
            "pruning_mask_generation_latency",
            "pruning_memory_budget",
            "calibration_syscall_budget",
            "detect_redundant_computation",
            "detect_memory_thrashing",
        ];

        for name in &expected_names {
            assert!(
                assertions.iter().any(|a| a.name == *name),
                "GOLD-010 FALSIFIED: Missing assertion: {name}"
            );
        }
    }

    #[test]
    fn test_latency_assertions_reasonable() {
        // TEST_ID: GOLD-011
        // Verify latency assertions have reasonable values
        let assertions = PruningGoldenTraces::performance_assertions();

        for assertion in &assertions {
            if assertion.assertion_type == AssertionType::Latency {
                assert!(
                    assertion.max_value <= 60000, // Max 60s
                    "GOLD-011 FALSIFIED: Latency {} unreasonably high: {}ms",
                    assertion.name,
                    assertion.max_value
                );
                assert!(
                    assertion.max_value >= 100, // Min 100ms
                    "GOLD-011 FALSIFIED: Latency {} unreasonably low: {}ms",
                    assertion.name,
                    assertion.max_value
                );
            }
        }
    }

    #[test]
    fn test_memory_assertion_reasonable() {
        // TEST_ID: GOLD-012
        // Verify memory assertion has reasonable value
        let assertions = PruningGoldenTraces::performance_assertions();

        let memory_assertion = assertions
            .iter()
            .find(|a| a.assertion_type == AssertionType::Memory)
            .expect("Memory assertion not found");

        // Should be between 100MB and 16GB
        assert!(
            memory_assertion.max_value >= 100_000_000,
            "GOLD-012 FALSIFIED: Memory budget too low: {}",
            memory_assertion.max_value
        );
        assert!(
            memory_assertion.max_value <= 17_179_869_184, // 16GB
            "GOLD-012 FALSIFIED: Memory budget too high: {}",
            memory_assertion.max_value
        );
    }

    #[test]
    fn test_golden_trace_serialization() {
        // TEST_ID: GOLD-020
        // Verify golden traces can be serialized
        let schedule_traces = PruningGoldenTraces::schedule_traces();
        let json = serde_json::to_string(&schedule_traces)
            .expect("GOLD-020 FALSIFIED: Failed to serialize schedule traces");
        assert!(
            json.contains("oneshot_1000"),
            "GOLD-020 FALSIFIED: Serialized trace should contain schedule type"
        );

        let assertions = PruningGoldenTraces::performance_assertions();
        let json = serde_json::to_string(&assertions)
            .expect("GOLD-020 FALSIFIED: Failed to serialize assertions");
        assert!(
            json.contains("pruning_importance_latency"),
            "GOLD-020 FALSIFIED: Serialized assertion should contain name"
        );
    }
}