felafax

FELAFAX -Google Cloud TPUでLlama3.1を30％低コストとスケールでシームレスに調整します！

FELAFAXは、 XLAランタイムを使用して、継続的なトレーニングおよび微調整オープンソースLLMのフレームワークです。必要なランタイムのセットアップを処理し、開始するためにJupyterノートブックをボックス外に提供します。

• 使いやすい。
• トレーニングのすべての側面を簡単に構成できます（ML研究者とハッカー向けに設計）。
• 8つのコアを備えた単一のTPU VMから、6000 TPUコア（ 1000X ）を含むTPU POD全体に簡単に拡張できます！

Felafaxでの私たちの目標は、インフラを構築して、非NVIDIAハードウェア（TPU、AWSトレインウム、AMD GPU、およびIntel GPU）でAIワークロードを簡単に実行できるようにすることです。

無料でFinetune

データセットを追加し、[すべて実行]をクリックすると、Google Colabで無料のTPUリソースで実行できます。

Felafaxサポート	無料ノートブック
llama 3.1（1b、3b）	▶️ Google Colab TPUで無料で開始します

現在サポートされているモデル

• llama-3.1JAX実装 $$ { color {red} new！} $$

  • パフォーマンスを向上させるために、PytorchからJaxに変換されました
  • 1b、3b、8b、70b、405bの全精度およびLoraトレーニングサポート。
  • Jaxのハードウェアが最適化されたXLAバックエンドを介して、多様なハードウェア（TPU、AWS Trainium、Nvidia、AMD）を効率的に実行します
  • 複数のアクセラレータにシャーディングして、より大きなコンテキストの長さとデータセットを処理するためのシームレスにスケール

• llama-3/3.1 pytorch xla

  • LORAおよび全精度トレーニングサポート
  • CODEPOINTER

Felafax CLI経由で微調整を実行します $$ { color {red} new！} $$

Felafax CLIを使用してモデルを微調整して、いくつかの簡単な手順で始めましょう。

ステップ1。CLIをインストールして認証します

CLIをインストールすることから始めます。

pip install pipx
pipx install felafax-cli

次に、認証トークンを生成します。

• Felafax.aiにアクセスして、アカウントに作成/サインインしてください。
• トークンページに移動し、新しいトークンを作成します。

最後に、トークンを使用してCLIセッションを認証します。

felafax-cli auth login --token < your_token >

ステップ2。微調整設定を設定します

まず、微調整用のデフォルトの構成ファイルを生成します。このコマンドは、デフォルトのハイパーパラメーター値を使用して、現在のディレクトリにconfig.ymlファイルを生成します。

felafax-cli tune init-config

次に、ハイパーパラメーターで構成ファイルを更新します。

• ハギングフェイスノブ：

  • HuggingfaceトークンとリポジトリIDを提供して、微調整されたモデルをアップロードします。

• データセットパイプラインとトレーニングパラメーション：

  • 微調整データセットに使用するように、 batch_size 、 max_seq_lengthを調整します。
  • Trainigをデータセット全体で実行する場合は、num_stepsをnullに設定します。 num_stepsが数に設定されている場合、指定されたステップ数の後にトレーニングが停止します。
  • learning_rateとlora_rankを設定して、微調整に使用します。
  • eval_intervalは、評価間のステップ数です。

ステップ3。微調整ランを開始します

フォローコマンドを実行して、微調整できるベースモデルのリストを表示します。今のところ、llama-3.1のすべてのバリエーションをサポートしています。

felafax-cli tune start --help

これで、上記のリストから選択したモデルを使用して微調整プロセスを開始し、Huggingfaceからデータセット名（ yahma/alpaca-cleanedなど）を開始できます。

felafax-cli tune start --model < your_selected_model > --config ./config.yml --hf-dataset-id < your_hf_dataset_name >

あなたを開始するための例コマンド：

felafax-cli tune start --model llama3-2-1b --config ./config.yml --hf-dataset-id yahma/alpaca-cleaned

微調整ジョブを開始した後、 Felafax CLIはTPUをスピンアップし、トレーニングを実行し、微調整されたモデルをHuggingface Hubにアップロードします。

その他の便利なコマンド

微調整ジョブを監視します

リアルタイムログをストリーミングして、微調整するジョブの進捗状況を監視できます。

 # Use `<job_name>` with the job namethat you get after starting the fine-tuning.
felafax-cli tune logs --job-id < job_name > -f

微調整されたモデルをリストします

微調整が完了した後、すべての微調整されたモデルをリストできます。

felafax-cli model list

微調整されたモデルとチャットします（もう一度TPUで実行されます！）：

インタラクティブな端末セッションを開始して、微調整されたモデルとチャットできます。

 # Replace `<model_id>` with model id from `model list` command you ran above.
felafax-cli model chat --model-id < model_id > 

ヘルプを使用して、より多くのコマンドを探索してください！

CLIは3つの主要なコマンドグループに分かれています。

• tune ：微調整ジョブを開始/停止します。
• model ：微調整されたモデルを管理および対話する。
• files ：YourDatasetファイルをアップロード/表示します。

--helpフラグを使用して、コマンドグループの詳細を発見してください。

felafax-cli tune --help

AMD 405B微調整ラン

最近、Pytorchの代わりにJaxを使用して、8xamd Mi300X GPUのLlama3.1 405bモデルを微調整しました。 JaxのAdvanced Sharding APIにより、素晴らしいパフォーマンスを実現できました。セットアップと使用したシャードのトリックについて学ぶために、ブログ投稿をご覧ください。

BFLOAT16精度のすべてのモデル重量とロラパラメーター、およびLORAランク8および16のロラアルファでLORA微調整を行いました。

• モデルサイズ： Llamaモデルの重みは、約800GBのVRAMを占めています。
• Lora Weights + Optimizer状態：約400GBのVRAM。
• 総VRAM使用量：総VRAMの77％、約1200GB。
• 制約： 405Bモデルのサイズが大きいため、バッチサイズとシーケンスの長さのスペースは限られていました。使用されたバッチサイズは16で、シーケンスの長さは64でした。
• トレーニング速度： 〜35トークン/秒
• メモリ効率：一貫して約70％
• スケーリング： Jaxを使用すると、スケーリングは8 GPUでほぼ線形でした。

GPU使用率とVRAM使用グラフは、以下にあります。ただし、モデルフロップ利用（MFU）を計算する必要があります。注：インフラストラクチャとVRAMの制約により、JITコンパイルされたバージョンの405Bモデルのバージョンを実行できませんでした（これをさらに調査する必要があります）。トレーニング全体の実行はJax Eaverモードで実行されたため、パフォーマンスの改善には大きな可能性があります。

• GPU利用：
• VRAMの利用：
• ROCM-SMIデータはここにあります。

クレジット：

• Google DeepmindのGemma Repo。
• JaxのLlama Modelsでの素晴らしい作業のためのEasylmとEleutherai
• Heegyukimが行ったPytorch XLA FSDPおよびSPMDテスト。
• Pytorch-XLAレポの例。

接触

ご質問がある場合は、[email protected]までお問い合わせください。