sokuto_cli/commands/
quantize.rs

1//! # ONNX INT8 PTQ 量子化コマンド
2//!
3//! FP32 ONNX モデルに対して動的 INT8 量子化 (Post-Training Quantization) を適用し、
4//! デシジョンヘッド (OptionGatherLayer, out_proj) の FP32 精度を維持しつつ
5//! モデルサイズ半減と低メモリ推論を実現する。
6
7use std::fs;
8use std::path::{Path, PathBuf};
9use std::process::Command;
10
11use serde::{Deserialize, Serialize};
12
13use sokuto_runtime::engine::{InferenceEngine, SessionConfig};
14use sokuto_runtime::tokenizer::JevTokenizer;
15
16use crate::commands::serve::validate_model_dir;
17
18/// 量子化実行オプション構造体。
19#[derive(Debug, Clone)]
20pub struct QuantizeArgs {
21    /// 元モデルファイル群が存在するディレクトリパス。
22    pub model_dir: PathBuf,
23    /// 量子化後成果物バンドルの出力先ディレクトリパス。
24    pub output_dir: PathBuf,
25    /// チャンネル単位での重み量子化フラグ。
26    pub per_channel: bool,
27    /// FP32 モデルとのパリティ検証を実行するフラグ。
28    pub verify: bool,
29}
30
31/// 量子化結果メタデータ。
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct QuantizeMetadata {
34    /// タイムスタンプ。
35    pub timestamp: String,
36    /// 量子化結果詳細。
37    pub quantize_result: QuantizeResultDetails,
38    /// チャンネル単位量子化フラグ。
39    pub per_channel: bool,
40}
41
42/// 量子化結果の詳細サマリー。
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct QuantizeResultDetails {
45    /// 入力モデルパス。
46    pub input_model_path: String,
47    /// 出力モデルパス。
48    pub output_model_path: String,
49    /// バンドル出力先ディレクトリ。
50    pub bundle_dir: String,
51    /// 元のファイルサイズ (バイト)。
52    pub original_size_bytes: u64,
53    /// 量子化後のファイルサイズ (バイト)。
54    pub quantized_size_bytes: u64,
55    /// 圧縮率。
56    pub compression_ratio: f64,
57    /// パリティ検証合否。
58    pub parity_passed: bool,
59    /// Top-1 決定一致率。
60    pub top1_agreement_rate: f64,
61    /// 最大絶対誤差。
62    pub max_abs_error: f64,
63    /// 平均絶対誤差。
64    pub mean_abs_error: f64,
65}
66
67/// ONNX INT8 PTQ 量子化コマンドを実行する。
68pub fn run_quantize(args: QuantizeArgs) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
69    // 1. 入力モデルディレクトリの存在と必須ファイルチェック
70    if let Err(err) = validate_model_dir(&args.model_dir) {
71        eprintln!("エラー: {err}");
72        return Err(err.into());
73    }
74
75    eprintln!(
76        "ONNX INT8 PTQ 量子化を開始します (入力: {}, 出力先: {})...",
77        args.model_dir.display(),
78        args.output_dir.display()
79    );
80
81    // 2. 出力先ディレクトリの確保
82    fs::create_dir_all(&args.output_dir)?;
83
84    // 3. train 側の Python 量子化スクリプトを実行 (uv run)
85    let project_root = find_project_root(&args.model_dir)?;
86    let train_dir = project_root.join("train");
87
88    if !train_dir.exists() {
89        return Err(format!(
90            "train ディレクトリが見つかりません: {}。",
91            train_dir.display()
92        )
93        .into());
94    }
95
96    let model_dir_abs = fs::canonicalize(&args.model_dir)?;
97    let output_dir_abs = if args.output_dir.is_absolute() {
98        args.output_dir.clone()
99    } else {
100        std::env::current_dir()?.join(&args.output_dir)
101    };
102
103    let quantize_script = train_dir.join("export").join("quantize.py");
104    if !quantize_script.exists() {
105        return Err(format!(
106            "量子化スクリプトが見つかりません: {}。",
107            quantize_script.display()
108        )
109        .into());
110    }
111
112    let mut cmd = Command::new("uv");
113    cmd.arg("run")
114        .arg("--project")
115        .arg(&train_dir)
116        .arg("python")
117        .arg(&quantize_script)
118        .arg("--model-dir")
119        .arg(&model_dir_abs)
120        .arg("--output-dir")
121        .arg(&output_dir_abs);
122
123    if args.per_channel {
124        cmd.arg("--per-channel");
125    } else {
126        cmd.arg("--no-per-channel");
127    }
128
129    if args.verify {
130        cmd.arg("--verify");
131    } else {
132        cmd.arg("--no-verify");
133    }
134
135    cmd.current_dir(&project_root);
136
137    eprintln!("Python 量子化パイプラインを起動中...");
138    let output = match cmd.output() {
139        Ok(out) => out,
140        Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
141            return Err(
142                "外部コマンド 'uv' の実行に失敗しました。uv がインストールされていないか、PATH に通っていません。Python 環境と uv を確認してください。".into()
143            );
144        }
145        Err(err) => {
146            return Err(format!("量子化サブプロセスの起動に失敗しました: {err}。").into());
147        }
148    };
149
150    if !output.status.success() {
151        let stderr = String::from_utf8_lossy(&output.stderr);
152        let stdout = String::from_utf8_lossy(&output.stdout);
153        eprintln!("量子化スクリプトの実行に失敗しました:\nSTDERR:\n{stderr}\nSTDOUT:\n{stdout}");
154        return Err("量子化サブプロセスの実行に失敗しました。".into());
155    }
156
157    // 4. 出力先成果物バンドルの完全性検証
158    let quantized_onnx = args.output_dir.join("model.onnx");
159    let quantized_tok = args.output_dir.join("tokenizer.json");
160
161    if !quantized_onnx.exists() || !quantized_tok.exists() {
162        return Err(format!(
163            "量子化成果物ファイルが不足しています: {} または {}。",
164            quantized_onnx.display(),
165            quantized_tok.display()
166        )
167        .into());
168    }
169
170    // 5. Rust ランタイムによるロード可能性の検証 (スモークテスト)
171    eprintln!("Rust ランタイムによる量子化モデルの整合性チェックを実行中...");
172    let test_engine = InferenceEngine::new(&quantized_onnx, SessionConfig::cpu_only())?;
173    let _test_tok = JevTokenizer::from_file(&quantized_tok)?;
174    drop(test_engine);
175
176    // 6. メタデータの読み込みとサマリー表示
177    let meta_file = args.output_dir.join("quantize_metadata.json");
178    if meta_file.exists() {
179        let meta_str = fs::read_to_string(&meta_file)?;
180        if let Ok(meta) = serde_json::from_str::<QuantizeMetadata>(&meta_str) {
181            let res = &meta.quantize_result;
182            println!();
183            println!("==========================================================================");
184            println!(" Sokuto ONNX INT8 PTQ 量子化サマリー");
185            println!("==========================================================================");
186            println!("成果物ディレクトリ : {}", args.output_dir.display());
187            println!(
188                "ファイルサイズ     : {:.1} MB -> {:.1} MB (削減率: {:.1}%)",
189                res.original_size_bytes as f64 / 1_048_576.0,
190                res.quantized_size_bytes as f64 / 1_048_576.0,
191                res.compression_ratio * 100.0
192            );
193            println!(
194                "Top-1 決定一致率   : {:.1}% ({})",
195                res.top1_agreement_rate * 100.0,
196                if res.parity_passed {
197                    "合格"
198                } else {
199                    "警告"
200                }
201            );
202            println!("最大絶対誤差 (Max) : {:.4}", res.max_abs_error);
203            println!("平均絶対誤差 (Mean): {:.4}", res.mean_abs_error);
204            println!("==========================================================================");
205            println!();
206        }
207    } else {
208        println!("量子化が正常に完了しました: {}", args.output_dir.display());
209    }
210
211    Ok(())
212}
213
214/// プロジェクトルートディレクトリを探索する。
215fn find_project_root(
216    start_dir: &Path,
217) -> Result<PathBuf, Box<dyn std::error::Error + Send + Sync>> {
218    let mut current = if start_dir.is_absolute() {
219        start_dir.to_path_buf()
220    } else {
221        std::env::current_dir()?.join(start_dir)
222    };
223
224    loop {
225        if current.join("Cargo.toml").exists() && current.join("crates").exists() {
226            return Ok(current);
227        }
228        if let Some(parent) = current.parent() {
229            current = parent.to_path_buf();
230        } else {
231            break;
232        }
233    }
234
235    // カレントディレクトリ起点でも探索
236    let cwd = std::env::current_dir()?;
237    let mut current = cwd;
238    loop {
239        if current.join("Cargo.toml").exists() && current.join("crates").exists() {
240            return Ok(current);
241        }
242        if let Some(parent) = current.parent() {
243            current = parent.to_path_buf();
244        } else {
245            break;
246        }
247    }
248
249    Err("プロジェクトルートディレクトリ (Cargo.toml, crates/ が存在する場所) を特定できませんでした。".into())
250}
sokuto_cli/commands/quantize.rs

sokuto_cli/commands/
quantize.rs
