Spaces:
Running
Running
| 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.5:.2f} GB") | |
| result_lines.append(f"【トレーニング（LoRA）】約 {usedStorage_gb * 2:.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.5:.2f} GB") | |
| result_lines.append(f"【トレーニング（LoRA）】約 {estimated_gb * 2:.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() | |