foxhunt/ml/tests/qat_device_consistency_test.rs

#[cfg(test)]
mod qat_device_consistency_tests {
    use candle_core::{Device, Tensor};
    use ml::memory_optimization::qat::*;
    use ml::tft::{QATTemporalFusionTransformer, TFTConfig, TemporalFusionTransformer};

    #[test]
    fn test_fake_quantize_device_consistency() {
        // Test that FakeQuantize respects input device
        let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

        let config = QATConfig::default();
        let mut observer = QuantizationObserver::new(config.clone(), device.clone());

        // Calibrate observer with some data
        let calibration_data = Tensor::randn(0f32, 1.0, (32, 64), &device).unwrap();
        observer.observe(&calibration_data).unwrap();

        // Simulate calibration complete
        for _ in 0..config.calibration_batches {
            let batch = Tensor::randn(0f32, 1.0, (32, 64), &device).unwrap();
            observer.observe(&batch).unwrap();
        }

        let fake_quant = FakeQuantize::from_observer(&observer).unwrap();

        // Create input tensor on CUDA (if available)
        let input = Tensor::randn(0f32, 1.0, (32, 64), &device).unwrap();

        // Forward pass should not crash with device mismatch
        let output = fake_quant.forward(&input).unwrap();

        // Verify output is on same device as input
        // Note: Device doesn't implement PartialEq, so we compare debug strings
        assert_eq!(
            format!("{:?}", input.device()),
            format!("{:?}", output.device()),
            "Output device mismatch: expected {:?}, got {:?}",
            input.device(),
            output.device()
        );
    }

    #[test]
    fn test_qat_tft_device_consistency() {
        // Test QATTemporalFusionTransformer device handling
        // Use CPU to avoid OOM on GPU during testing
        let device = Device::Cpu;

        // Create FP32 model with smaller batch size
        let config = TFTConfig::default();
        let fp32_model =
            TemporalFusionTransformer::new_with_device(config.clone(), device.clone()).unwrap();

        // Create QAT wrapper
        let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model).unwrap();

        // Create input tensors on same device with smaller batch size (4 instead of 32)
        let static_features = Tensor::randn(0f32, 1.0, (4, 5), &device).unwrap();
        let historical_ts = Tensor::randn(0f32, 1.0, (4, 50, 210), &device).unwrap();
        let future_ts = Tensor::randn(0f32, 1.0, (4, 10, 10), &device).unwrap();

        // Forward pass should not crash
        let output = qat_model
            .forward(&static_features, &historical_ts, &future_ts)
            .unwrap();

        // Verify output device
        assert_eq!(
            format!("{:?}", device),
            format!("{:?}", output.device()),
            "Output device mismatch"
        );
    }

    /// NEW TEST: CPU calibration → CUDA training (primary production scenario)
    /// This test validates the fix for the device mismatch bug that occurs when
    /// calibration happens on CPU but training happens on CUDA.
    #[test]
    fn test_fake_quantize_cpu_calibration_cuda_training() {
        // Skip if CUDA not available (CI environments)
        if Device::cuda_if_available(0).is_err() {
            println!("Skipping CUDA test: no CUDA device available");
            return;
        }

        let cpu_device = Device::Cpu;
        let cuda_device = Device::new_cuda(0).unwrap();

        // Phase 1: Calibration on CPU (common for large datasets)
        let config = QATConfig::default();
        let mut observer = QuantizationObserver::new(config.clone(), cpu_device.clone());

        println!(
            "✅ Calibrating on CPU with {} batches",
            config.calibration_batches
        );
        for i in 0..config.calibration_batches {
            let batch = Tensor::randn(0f32, 1.0, (8, 32), &cpu_device).unwrap();
            observer.observe(&batch).unwrap();

            if i % 20 == 0 {
                println!(
                    "   Calibrated {}/{} batches",
                    i + 1,
                    config.calibration_batches
                );
            }
        }

        assert!(observer.is_calibrated(), "Observer should be calibrated");
        println!("✅ Calibration complete");

        // Create FakeQuantize from calibrated observer (still on CPU)
        let fake_quant = FakeQuantize::from_observer(&observer).unwrap();
        println!("✅ FakeQuantize created from observer");

        // Phase 2: Training on CUDA (inference on GPU for speed)
        println!("✅ Testing forward pass with CUDA input");
        let cuda_input = Tensor::randn(0f32, 1.0, (8, 32), &cuda_device).unwrap();

        // This should NOT crash with device mismatch error
        let output = fake_quant.forward(&cuda_input).unwrap();

        // Verify output is on CUDA device (matches input)
        assert_eq!(
            format!("{:?}", cuda_device),
            format!("{:?}", output.device()),
            "Output should be on CUDA device to match input"
        );

        println!("✅ Device transition test passed: CPU calibration → CUDA training");
    }

    /// NEW TEST: CUDA calibration → CPU inference (edge case)
    /// Tests the opposite direction: calibrate on GPU, infer on CPU.
    #[test]
    fn test_fake_quantize_cuda_calibration_cpu_inference() {
        // Skip if CUDA not available
        if Device::cuda_if_available(0).is_err() {
            println!("Skipping CUDA test: no CUDA device available");
            return;
        }

        let cuda_device = Device::new_cuda(0).unwrap();
        let cpu_device = Device::Cpu;

        // Phase 1: Calibration on CUDA
        let config = QATConfig::default();
        let mut observer = QuantizationObserver::new(config.clone(), cuda_device.clone());

        println!(
            "✅ Calibrating on CUDA with {} batches",
            config.calibration_batches
        );
        for _ in 0..config.calibration_batches {
            let batch = Tensor::randn(0f32, 1.0, (8, 32), &cuda_device).unwrap();
            observer.observe(&batch).unwrap();
        }

        let fake_quant = FakeQuantize::from_observer(&observer).unwrap();

        // Phase 2: Inference on CPU
        println!("✅ Testing forward pass with CPU input");
        let cpu_input = Tensor::randn(0f32, 1.0, (8, 32), &cpu_device).unwrap();

        let output = fake_quant.forward(&cpu_input).unwrap();

        // Verify output is on CPU device (matches input)
        assert_eq!(
            format!("{:?}", cpu_device),
            format!("{:?}", output.device()),
            "Output should be on CPU device to match input"
        );

        println!("✅ Device transition test passed: CUDA calibration → CPU inference");
    }

    /// NEW TEST: Multiple device transitions in single session
    /// Stress test for device management robustness.
    #[test]
    fn test_fake_quantize_multiple_device_transitions() {
        // Skip if CUDA not available
        if Device::cuda_if_available(0).is_err() {
            println!("Skipping CUDA test: no CUDA device available");
            return;
        }

        let cpu_device = Device::Cpu;
        let cuda_device = Device::new_cuda(0).unwrap();

        // Calibrate on CPU
        let config = QATConfig::default();
        let mut observer = QuantizationObserver::new(config.clone(), cpu_device.clone());

        for _ in 0..config.calibration_batches {
            let batch = Tensor::randn(0f32, 1.0, (4, 16), &cpu_device).unwrap();
            observer.observe(&batch).unwrap();
        }

        let fake_quant = FakeQuantize::from_observer(&observer).unwrap();

        println!("✅ Testing multiple device transitions:");

        // Transition 1: CPU → CUDA
        println!("   1. CPU → CUDA");
        let cpu_input = Tensor::randn(0f32, 1.0, (4, 16), &cpu_device).unwrap();
        let output1 = fake_quant.forward(&cpu_input).unwrap();
        assert_eq!(
            format!("{:?}", cpu_device),
            format!("{:?}", output1.device())
        );

        // Transition 2: CUDA → CUDA (same device)
        println!("   2. CUDA → CUDA");
        let cuda_input = Tensor::randn(0f32, 1.0, (4, 16), &cuda_device).unwrap();
        let output2 = fake_quant.forward(&cuda_input).unwrap();
        assert_eq!(
            format!("{:?}", cuda_device),
            format!("{:?}", output2.device())
        );

        // Transition 3: CUDA → CPU (back to CPU)
        println!("   3. CUDA → CPU");
        let cpu_input2 = Tensor::randn(0f32, 1.0, (4, 16), &cpu_device).unwrap();
        let output3 = fake_quant.forward(&cpu_input2).unwrap();
        assert_eq!(
            format!("{:?}", cpu_device),
            format!("{:?}", output3.device())
        );

        // Transition 4: CPU → CUDA (repeat first transition)
        println!("   4. CPU → CUDA (repeat)");
        let cuda_input2 = Tensor::randn(0f32, 1.0, (4, 16), &cuda_device).unwrap();
        let output4 = fake_quant.forward(&cuda_input2).unwrap();
        assert_eq!(
            format!("{:?}", cuda_device),
            format!("{:?}", output4.device())
        );

        println!("✅ Multiple device transitions test passed");
    }

    /// NEW TEST: Validate quantization correctness across devices
    /// Ensures fake quantization produces same numerical results regardless of device.
    #[test]
    fn test_fake_quantize_numerical_consistency_across_devices() {
        // Skip if CUDA not available
        if Device::cuda_if_available(0).is_err() {
            println!("Skipping CUDA test: no CUDA device available");
            return;
        }

        let cpu_device = Device::Cpu;
        let cuda_device = Device::new_cuda(0).unwrap();

        // Create identical observers on both devices
        let config = QATConfig::default();
        let mut cpu_observer = QuantizationObserver::new(config.clone(), cpu_device.clone());
        let mut cuda_observer = QuantizationObserver::new(config.clone(), cuda_device.clone());

        // Calibrate with identical data (semantically, different devices)
        for _ in 0..config.calibration_batches {
            let cpu_batch = Tensor::randn(0f32, 1.0, (4, 8), &cpu_device).unwrap();
            let cuda_batch = cpu_batch.to_device(&cuda_device).unwrap();

            cpu_observer.observe(&cpu_batch).unwrap();
            cuda_observer.observe(&cuda_batch).unwrap();
        }

        let cpu_fake_quant = FakeQuantize::from_observer(&cpu_observer).unwrap();
        let cuda_fake_quant = FakeQuantize::from_observer(&cuda_observer).unwrap();

        // Create test input
        let test_input_cpu = Tensor::new(&[[1.5f32, 2.3, -1.2, 0.5]], &cpu_device).unwrap();
        let test_input_cuda = test_input_cpu.to_device(&cuda_device).unwrap();

        // Apply fake quantization on both devices
        let output_cpu = cpu_fake_quant.forward(&test_input_cpu).unwrap();
        let output_cuda = cuda_fake_quant.forward(&test_input_cuda).unwrap();

        // Move CUDA output to CPU for comparison
        let output_cuda_on_cpu = output_cuda.to_device(&cpu_device).unwrap();

        // Compare numerical results (should be identical or very close)
        let cpu_data = output_cpu.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        let cuda_data = output_cuda_on_cpu
            .flatten_all()
            .unwrap()
            .to_vec1::<f32>()
            .unwrap();

        assert_eq!(
            cpu_data.len(),
            cuda_data.len(),
            "Output lengths should match"
        );

        for (i, (cpu_val, cuda_val)) in cpu_data.iter().zip(cuda_data.iter()).enumerate() {
            let diff = (cpu_val - cuda_val).abs();
            assert!(
                diff < 1e-5,
                "Quantization result mismatch at index {}: CPU={}, CUDA={}, diff={}",
                i,
                cpu_val,
                cuda_val,
                diff
            );
        }

        println!("✅ Numerical consistency test passed across devices");
    }
}