ML Framework

article technology medium #ml-framework#pytorch#jax#huggingface#deepspeed#fsdp#peft#llm-training
Created: 2026-05-28 Updated:

ML フレームワークを 10 セクション (DL コア / HuggingFace / 分散学習 / FT / MLOps / HPO / 古典 ML / データ / コンパイル / 落とし穴) で俯瞰、BUILD/TRAIN/FINE-TUNE/EXPORT スコープ。

ML Framework

ML フレームワークは BUILD/TRAIN/FINE-TUNE/EXPORT を担う 9 層 (DL コア+HuggingFace+分散学習+FT+MLOps+HPO+古典 ML+データ+コンパイル)。分散アルゴリズム (tech-48) の実装層で、物理基盤 (tech-71) と推論 (tech-69)・アプリ層 (tech-72) を繋ぐ中核。

Deep Learning Core

PyTorch (2.3.x、Meta、BSD-3) は LLM の事実上標準、HF Hub 上位 >95% が PyTorch 重み。torch.compile (2.0+) は TorchDynamoAOTAutogradInductor (Triton/C++) で最大 2× 高速化、mode=default/reduce-overhead/max-autotunetorch.distributed+DDP/FSDP/FSDP2 (per-parameter DTensor+torch.compile composable+fully_shard()、tech-48)、torch.export、ExecuTorch、flex_attention (2.5 想定、未確認)。TensorFlow/Keras (TF 2.16.x+Keras 3.x) は tf.function+XLA JIT、Mirrored/MultiWorkerMirrored/TPUStrategy、TFLite/TFX/SavedModel。LLM 主流からは 2022 頃流出、エッジ・既存本番は現役。Keras 3 はマルチバックエンド、採用限定的 (未確認)。JAX (0.4.x、DeepMind) は jit/grad/vmap/pmapshard_map (pjit 後継)、lax.scanPallas (Triton 相当 DSL、安定性未確認)、不変テンソル。NN は Flax/NNX (NNX は mutable 推奨)、EquinoxHaiku (Flax 統合)、OptaxOrbax。Gemini 学習で使用。レガシー: MXNet/Theano/Caffe/CNTK/Chainer (いずれも 2017–2023 に停止・統合・アーカイブ)。

HuggingFace Stack

モデル配布・学習・FT のデファクト統合基盤。transformers (4.40.x+): AutoModel/AutoTokenizer/AutoProcessorgenerate()Trainer、Hub 70 万+ モデル (v5 未確認)。accelerate (0.29.x+): Accelerator で DDP/FSDP/DeepSpeed 透過切替+mixed precision。peft (0.10.x+): LoRA (Hu 2021、rank-r W=BAW=BA 凍結)、QLoRA (Dettmers 2023、4-bit NF4+fp16 LoRA)、DoRA (2024)、IA3/Prefix/Prompt/P-tuning v2、LoftQtrl (0.8.x+): SFTTrainer/DPOTrainer (Rafailov 2023)/PPOTrainer/GRPOTrainer (DeepSeek-R1)/RewardTrainer/BCOTrainer/ORPOTrainerdiffusers (0.27.x+) は SD/SDXL/Flux/Consistency+LoRA/DreamBooth。datasets (Arrow+streaming+Hub)、tokenizers (Rust BPE/WordPiece/SentencePiece)、evaluate (BLEU/ROUGE/Perplexity)、safetensors (安全シリアライズ)。optimumexporters.onnx/optimum-intel (OpenVINO)/optimum-nvidia (TensorRT)/BetterTransformer のハードウェア最適化統合 API。

Distributed Training

tech-48 の実装層。DeepSpeed (0.14.x+、MS) は ZeRO 1/2/3+ZeRO-Offload (CPU/NVMe)+ZeRO-Infinity+MoE+1-bit Adam/LAMB。ZeRO++ (2023、通信 4× 削減、採用率未確認)。HF accelerate から呼出、Mistral/Falcon 等で実績。Megatron-LM/Megatron-Core (NVIDIA) は TP/PP/SP+Distributed Optimizer+Flash Attention+Transformer Engine (FP8)、Core はモジュラー再実装、GPT-NeoX/BLOOM/Llama 2 等 70B+ で標準。Megatron-DeepSpeed は BLOOM 176B 実績 (以降メンテ分散、未確認)。FSDP/FSDP2 は Meta Llama 2/3 公式採用。Colossal-AI は 1D/2D/2.5D/3D 並列+Auto-parallelism+Gemini (動的 ZeRO)、産業採用限定的 (未確認)。Ray Train は分散+elastic+Ray Tune 連携。PyTorch Lightning+Fabric (2.x) は Trainer で定型削減、Fabric は最小ラッパー、Lightning AI がクラウド実行提供。

Fine-tuning

Axolotl (Apache 2.0) は YAML 1 ファイルでフル/LoRA/QLoRA/RLHF、transformers+DeepSpeed/FSDP ラッパー。LLaMA-Factory (0.8.x+) は 100+ モデル統一 FT+WebUI/CLI+SFT/DPO/PPO/ORPO/GRPO+GPTQ/AWQ 量化、中国語圏で広く採用。Unsloth (Apache 2.0+Pro) は Triton 手書きカーネルで QLoRA を 2–4× 高速・50–80% メモリ削減 (公称値、独立検証未確認)、Llama/Mistral/Gemma/Phi。torchtune (0.2.x+、PyTorch 公式) はレシピベース+LoRA/QLoRA/フル FT+torch.compile+FSDP2 統合、安定だが網羅性不足 (未確認)。SWIFT (Alibaba/ModelScope、Qwen/InternLM)、LMFlow (HKUST 研究)。選定軸は YAML 駆動 → 機能網羅 → 速度極限 → 公式安定 → 中国語圏。

Experiment Tracking & MLOps

W&B (SaaS+オンプレ) は run/artifact/sweep/registry/LLM prompt logging 統合で最広範採用、wandb.init() で全メトリクス記録。MLflow (2.x、Databricks) は experiment/run/registry/serving でエンタープライズに厚い。TensorBoard は軽量。Neptune.ai (チーム協働)、Comet ML (LLMops)、ClearML (E2E+タスクキュー、セルフホスト)、Aim (ローカルファースト)、DVC (データ/モデル版管理 git-like)。コンテナ・オーケストレーションは tech-66

Hyperparameter Tuning

Optuna (3.x、PFN、MIT) は TPE/CMA-ES/QMC サンプラー+pruner で early stopping、PyTorch エコの OSS 標準。Ray Tune は分散 HPO、Optuna/HyperOpt バックエンド指定可。W&B Sweeps は内蔵で Bayesian/Grid/Random。Hyperopt (TPE 最古参、メンテ低下)、SigOpt (Intel 買収後 OSS、商用減)、Vizier (google-vizier OSS)、Ax (Meta、BO+bandit、BoTorch)。選定軸は探索空間構造+並列度 (TPE は逐次依存、CMA-ES は集団並列向き)。

Classical ML

タブラー領域は LLM 台頭後も独立した生態系。scikit-learn (1.4.x+) は統一 fit/predict/transform+Pipeline/ColumnTransformer+100+ モデル+joblib 並列、タブラー前処理・評価の永続標準。XGBoost (2.x) は GPU HistTreeMethod+Dask/Spark+DART+分位点回帰で Kaggle 最頻出。LightGBM (4.x、MS) は leaf-wise+GOSS+EFB+カテゴリ特徴ネイティブ+GPU で大規模カテゴリカル優位。CatBoost (Yandex) は ordered TS エンコード+target leakage 対策組込で産業データに強い。statsmodels は p 値/信頼区間、ARIMA/VAR、GLM で統計推論。タブラー回帰・分類で依然優位。

Data Processing

NumPy 2.0 (2024、C 拡張 ABI 変更で互換性注意、状況未確認)+SciPy (1.12.x+) が基盤。pandas (2.x、PyArrow+nullable dtypes) vs Polars (Rust、0.20.x+、遅延評価+並列+Arrow ネイティブ、CSV/Parquet で 5–20× 高速、ベンチ詳細未確認)。Apache Arrow は HF datasets/Polars/Spark の共通 I/O。DuckDB (0.10.x+) は SQL on Parquet/Arrow をインプロセス実行、pandas 代替 ETL が拡大。学習データロードは PyTorch DataLoader (num_workers+pin_memory+prefetch_factor が GPU 利用率に直結) と tf.data (TPU 相性最良)。WebDataset は tar ストリーミング+S3/GCS 直結で LAION 等の配布標準。FFCV (beton+OS ページキャッシュ) は ImageNet 5–10× 高速化 (未確認)。Mosaic StreamingDataset (Databricks) は MDS+シャーディングで LLM 事前学習スケール対応。NVIDIA DALI は CPU 前処理ボトルネックを GPU/CPU 混合で解消。MLX (Apple、NumPy ライク+lazy+Metal+統合メモリ、バージョン未確認) は Mac ローカル LLM プロトで台頭。

Compilation & Optimization

torch.compile/Inductor (PyTorch 2.0+) は Dynamo→AOTAutograd→Inductor で Triton 生成、max-autotune 自動チューニング (コスト大)、graph break で JIT フォールバック。XLA (Google、TF/JAX プライマリ) は TPU 必須、torch_xla で PyTorch も TPU 実行可。Apache TVM は CPU/GPU/ARM/RISC-V にコンパイル、Relax IR (2023~) で動的形状対応、AutoTVM/Ansor、MLIR 相互運用。MLIR (LLVM サブ) は複数 dialect で TF/IREE/Torch-MLIR/Triton の共通基盤。ONNX は交換形式、ONNX Runtime (MS、1.17.x+) は CPU/CUDA/TensorRT/OpenVINO 複数 EP、ORTTrainingModule で学習対応。OpenVINO (Intel) は INT8 量化+グラフ融合でエッジに強い。TensorRT は推論最適化 (tech-69)、学習側は Transformer Engine (FP8) として DeepSpeed/Megatron と統合。Triton-lang は Inductor 経由は本番、standalone API は別 (未確認)。GPU コンピュート API は tech-70MLC-LLM (TVM) は WebGPU/Metal/CUDA/OpenCL クロスプラットフォーム、ブラウザ内 LLM 推論候補。

2026 フロンティアと選定

未確認: PyTorch 2.5/2.6 flex_attention+torch.export 安定化+FSDP2 デフォルト化の時期・機能、JAX Pallas 安定化、Keras 3 産業採用率、HF transformers v5 機能・時期、DeepSpeed ZeRO++ 産業採用率、Modular Mojo/MAX 本番運用、Triton-lang standalone API 安定性、NumPy 2.0 互換、Unsloth 2–4× 独立検証、torchtune 網羅性比較根拠。選定: LLM 事前学習 (10B+) → PyTorch+FSDP2 または DeepSpeed ZeRO-3+Megatron-Core (NVIDIA GPU)。FT → transformers+peft+trl+Axolotl/LLaMA-Factory。高速 QLoRA → Unsloth。TPU → JAX+Flax/NNX+Optax。Apple M-series → MLX。タブラー → scikit-learn+XGBoost/LightGBM。HPO → Optuna/Ray Tune。実験追跡 → W&B/MLflow。エクスポート → ONNX Runtime/OpenVINO/ExecuTorch。スコープは BUILD/TRAIN/FINE-TUNE/EXPORT、推論本番は tech-69、アプリ層は tech-72、アルゴリズム理論は tech-48 に分離。

ピットフォール

  • P1 JAX を XLA 知識なしに採用: shape 静的解析失敗でコンパイルエラー多発、jit 再コンパイルでスループット激減。コンパイル境界・静的 shape 設計必須、入門は Equinox。
  • P2 FSDP1 と torch.compile 混在: graph break 大量で速度向上なし。PyTorch 2.x 以降は FSDP2 + torch.compile 公式推奨。
  • P3 torch.compile 初回ウォームアップ無視: 1 回目含む計測で誤結論。定常スループットで評価、mode="reduce-overhead" は CUDA graph で毎回コンパイル不要。
  • P4 DeepSpeed と FSDP 二重適用: HF Trainer + DeepSpeed config + FSDP フラグ同時有効で劣化。accelerate は片方のみ。
  • P5 QLoRA で base dtype 誤り: bf16 base に NF4 量化すべきところ fp32 でメモリ爆発。load_in_4bit=True+bnb_4bit_compute_dtype=torch.bfloat16+bnb_4bit_quant_type="nf4" 明示。
  • P6 DataLoader ボトルネックを GPU 不足と誤認: GPU 利用率 30–50% 停滞、num_workers=0 のまま。num_workers >= CPU コア/GPU 数+pin_memory=True+prefetch_factor=2、FFCV/DALI/WebDataset 検討。
  • P7 実験追跡なし HPO: どの run が最良か不明・再現不可能。W&B か MLflow を初期導入。
  • P8 ONNX shape/op カバレッジ過信: torch.onnx.export 成功でも ORT で unsupported op、動的 shape 実行時失敗。エクスポート直後に onnxruntime でサンプル検証、opset と target EP の op サポートを事前確認。

Local graph