app

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+import requests
+
+
+# 定数定義
+API_BASE_URL = "https://huggingface.co/api/models/"
+MULTIPLIERS = {
+    "I8": 1,
+    "U8": 1,
+    "F8_E5M2": 1,
+    "F8_E4M3": 1,
+    "I16": 2,
+    "U16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "I32": 4,
+    "U32": 4,
+    "F32": 4,
+    "I64": 8,
+    "U64": 8,
+    "F64": 8,
+}
+PRECISION_KEYS = list(MULTIPLIERS.keys())
+COMMENT = {
+    "I8": "INT8の処理能力が高いGPUを選定してください。",
+    "U8": "INT8の処理能力が高いGPUを選定してください。",
+    "F8_E5M2": "FP8の処理能力が高いGPUを選定してください。",
+    "F8_E4M3": "FP8の処理能力が高いGPUを選定してください。",
+    "I16": "INT16の処理能力が高いGPUを選定してください。",
+    "U16": "INT16の処理能力が高いGPUを選定してください。",
+    "F16": "FP16の処理能力が高いGPUを選定してください。",
+    "BF16": "BF16の処理能力が高いGPUを選定してください。",
+    "I32": "INT32の処理能力が高いGPUを選定してください。",
+    "U32": "INT32の処理能力が高いGPUを選定してください。",
+    "F32": "FP32の処理能力が高いGPUを選定してください。",
+    "I64": "INT64の処理能力が高いGPUを選定してください。",
+    "U64": "INT64の処理能力が高いGPUを選定してください。",
+    "F64": "FP64の処理能力が高いGPUを選定してください。",
+}
+
+def get_model_api_info(model_name: str) -> dict:
+    """
+    Hugging Face Hub APIから指定モデルの情報を取得する。
+
+    :param model_name: モデル名（例: "meta-llama/Llama-3.3-70B-Instruct"）
+    :return: モデル情報の辞書。取得に失敗した場合はNoneを返す。
+    """
+    api_url = f"{API_BASE_URL}{model_name}"
+    try:
+        response = requests.get(api_url)
+        response.raise_for_status()
+        return response.json()
+    except requests.RequestException:
+        return None
+
+def sum_precision_sizes(parameters: dict) -> int:
+    """
+    safetensors内の各精度のパラメータ数に対して、定数の乗数をかけた合計バイト数を算出する。
+
+    :param parameters: 各精度のパラメータ数を含む辞書
+    :return: 合計バイト数
+    """
+    # 各データ型とその乗数の対応表
+    total_bytes = 0
+    for precision, multiplier in MULTIPLIERS.items():
+        count = parameters.get(precision, 0)
+        total_bytes += count * multiplier
+    return total_bytes
+
+def estimate_gpu_memory(model_name: str) -> str:
+    """
+    指定したモデル名からAPI情報を取得し、safetensors内の各精度パラメータサイズの合算値から
+    GPUメモリ必要量を概算する。
+
+    :param model_name: モデル名
+    :return: GPUメモリ必要量などの情報を含むメッセージ文字列
+    """
+    result_lines = []
+
+    model_info = get_model_api_info(model_name)
+    if model_info is None:
+        result_lines.append(f"エラー: モデル '{model_name}' の情報が取得できませんでした。")
+        return "\n".join(result_lines)
+
+    # リポジトリ全体のファイルサイズ（参考）
+    usedStorage = model_info.get("usedStorage")
+    usedStorage_gb = usedStorage / (1024 ** 3) if usedStorage else None
+
+    safetensors = model_info.get("safetensors")
+
+    # safetensorsの情報が見つからない場合
+    if not safetensors or "parameters" not in safetensors:
+        result_lines.append("safetensorsの情報が見つかりません。")
+        if usedStorage_gb is not None:
+            result_lines.append("")
+            result_lines.append(f"参考までに、該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GBです。")
+            result_lines.append("")
+            result_lines.append("これを全てモデルのデータとして仮定した場合、推論及び、トレーニングに必要な概算のGPUメモリサイズは")
+            result_lines.append(f"【推論】約 {usedStorage_gb * 1.3:.2f} GB")
+            result_lines.append(f"【トレーニング（LoRA）】約 {usedStorage_gb * 1.8:.2f} GB")
+            result_lines.append(f"【トレーニング（フルパラメータ）】約 {usedStorage_gb * 4:.2f} GB")
+            result_lines.append("となります。")
+        return "\n".join(result_lines)
+
+    # safetensorsの情報がある場合
+    parameters = safetensors["parameters"]
+    estimated_bytes = sum_precision_sizes(parameters)
+    estimated_gb = estimated_bytes / (1024 ** 3)
+
+    result_lines.append(f"モデル '{model_name}' のsafetensors情報より、各パラメータ数は")
+
+    # 各精度ごとのパラメータ数の出力
+    max_precision = None
+    max_count = 0
+    for precision in PRECISION_KEYS:
+        count = parameters.get(precision, 0)
+        if count > max_count:
+            max_precision = precision
+            max_count = count
+        if precision in parameters:
+            result_lines.append(f"【{precision}】 {parameters[precision]:,}")
+
+    result_lines.append(f"これらを合算するとモデルのデータサイズは約 {estimated_gb:.2f} GB です。")
+    result_lines.append("")
+    result_lines.append("推論及び、トレーニングに必要な概算のGPUメモリサイズは")
+    result_lines.append(f"【推論】約 {estimated_gb * 1.3:.2f} GB")
+    result_lines.append(f"【トレーニング（LoRA）】約 {estimated_gb * 1.8:.2f} GB")
+    result_lines.append(f"【トレーニング（フルパラメータ）】約 {estimated_gb * 4:.2f} GB")
+    result_lines.append("となります。")
+
+    # GPU選定のコメント
+    if max_precision is not None:
+        comment_message = COMMENT.get(max_precision, "")
+        result_lines.append(comment_message)
+
+    if usedStorage_gb is not None:
+        result_lines.append("")
+        result_lines.append(f"参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB")
+
+    return "\n".join(result_lines)
+
+# Gradio インターフェースの定義
+iface = gr.Interface(
+    fn=estimate_gpu_memory,
+    inputs=gr.Textbox(label="モデル名 (例: meta-llama/Llama-3.3-70B-Instruct)"),
+    outputs="text",
+    title="Model memory estimator",
+    description=("Hugging Face Hub APIから取得したsafetensorsの情報をもとに、すべてのパラメータのデータサイズを合計し、その合計値から必要なGPUメモリサイズを概算（GB単位）で計算します。")
+)
+
+iface.launch(server_name="0.0.0.0")