erfanzar/EasyDeL

解決的問題

EasyDeL 是一個用於大規模機器學習模型訓練、微調和推理服務的高性能框架。它彌合了 HuggingFace Transformers 等函式庫的易用性與基於 JAX 的系統極高性能之間的鴻溝，允許開發人員在不犧牲可擴展性的情況下，在 TPU 和 GPU 集群上部署 LLM、多模態模型和視覺模型。

工作原理

EasyDeL 構建在 Flax NNX 和 JAX 之上，提供模組化架構，支援廣泛的分散式訓練策略（DP, FSDP, TP, EP, SP）以及透過 Triton (GPU) 和 Pallas (TPU) 優化的內核。它包含用於 16 種專業訓練器（涵蓋 SFT、RLHF 和蒸餾）的統一 API，以及名為 eSurge 的專用推理引擎，該引擎實現了連續批處理（continuous batching）和分頁 KV 快取，用於高吞吐量推理。

適用對象

專為需要在 TPU/GPU 集群上大規模訓練或部署大型模型，但又希望擁有透明、可讀且易於修改和定制的程式碼庫的 ML 工程師和研究人員設計。

亮點

• 廣泛的模型支援：包括 Llama、Qwen、Mistral、DeepSeek 和 Gemma 在內的 70 多種架構。
• 企業級推理：eSurge 引擎提供相容 OpenAI 的 API、串流傳輸以及文本、圖像和影片的多模態推理服務。
• 多樣化的訓練套件：針對 DPO、ORPO、GRPO 和知識蒸餾的專業訓練器。
• 先進的注意力機制：實現了包括 Flash Attention 2 和 Ring Attention 在內的 10 多種機制。
• 多模態能力：支援視覺語言模型 (LLaVA, Qwen2-VL)、語音辨識 (Whisper) 和擴散模型。
• 統一管理：eLargeModel 類別簡化了大型模型的配置、構建和編排。