Neuron 模型快取

為什麼要使用快取？

問題：大型模型的 Neuron 編譯需要 30-60 分鐘
解決方案：在幾秒鐘內下載預編譯模型

快取系統將編譯好的 Neuron 模型儲存在 HuggingFace Hub 上，為您的團隊省去了重新編譯的時間。當您訓練或載入模型時，系統會自動檢查是否存在快取版本，然後再開始昂貴的編譯過程。

主要優勢

節省時間：在幾秒鐘內下載編譯好的模型，而不是花費數小時進行編譯
團隊協作：在團隊成員和例項之間共享編譯好的模型
降低成本：避免在雲實例上重複編譯的成本
自動操作：與現有程式碼透明地協同工作

快速入門

訓練

from optimum.neuron import NeuronTrainer

# Cache works automatically - no configuration needed
trainer = NeuronTrainer(model=model, args=training_args)
trainer.train()  # Downloads cached models if available

推理

from optimum.neuron import NeuronModelForCausalLM

# Cache works automatically 
model = NeuronModelForCausalLM.from_pretrained("model_id")

就這樣！對於支援的模型類，快取會自動工作。

支援的模型

模型類	快取支援	用例	備註
`NeuronTrainer`	✅ 完全支援	訓練	訓練期間自動下載和上傳
`NeuronModelForCausalLM`	✅ 完全支援	推理	自動下載用於推理
其他 `NeuronModelForXXX`	❌ 無	推理	使用不同的匯出機制，無快取整合

重要限制：像 NeuronModelForSequenceClassification、NeuronModelForQuestionAnswering 等模型使用不同的編譯路徑，該路徑不與快取系統整合。只有 NeuronModelForCausalLM 和訓練工作流支援快取。

工作原理

快取系統在兩個層面上執行以最小化編譯時間

快取優先順序（從快到慢）

本地快取 → 從 /var/tmp/neuron-compile-cache 即時訪問
Hub 快取 → 在幾秒鐘內從 HuggingFace Hub 下載
從頭編譯 → 大型模型需要 30-60 分鐘

快取內容：系統快取 NEFF 檔案（Neuron 可執行檔案格式）—— 即在 Neuron 核心上執行的已編譯二進位制產物，而不是原始模型檔案。

快取標識：每個快取的編譯都會根據以下因素獲得一個唯一的雜湊值

模型因素：架構、精度（fp16/bf16）、輸入形狀、任務型別
編譯因素：NeuronX 編譯器版本、核心數、最佳化標誌
環境因素：模型檢查點版本、Optimum Neuron 版本

這意味著即使對您的設定進行微小更改也可能需要重新編譯，但相同的配置將始終命中快取。

私有快取設定

預設的公共快取（aws-neuron/optimum-neuron-cache）對使用者是只讀的——您可以下載快取的模型，但不能上傳您自己的編譯結果。這個公共快取只包含由 Optimum 團隊為常見配置編譯的模型。

對於大多數用例，您需要建立一個私有快取倉庫，用來儲存您自己編譯的模型。

為什麼使用私有快取？

上傳您的編譯結果：儲存您編譯的模型以供團隊重複使用
私有模型：確保專有模型編譯結果的安全
團隊協作：在團隊成員和 CI/CD 之間共享編譯好的產物
自定義配置：快取具有您特定批次大小、序列長度等的模型

方法 1：CLI 設定（推薦）

# Create private cache repository
optimum-cli neuron cache create

# Set as default cache
optimum-cli neuron cache set your-org/your-cache-name

方法 2：環境變數

# Use for single training run
CUSTOM_CACHE_REPO="your-org/your-cache" python train.py

# Or export for session
export CUSTOM_CACHE_REPO="your-org/your-cache"

先決條件

登入：huggingface-cli login
快取倉庫的寫許可權

CLI 命令

# Create new cache repository
optimum-cli neuron cache create [-n NAME] [--public]

# Set default cache repository  
optimum-cli neuron cache set REPO_NAME

# Search for cached models
optimum-cli neuron cache lookup MODEL_ID

# Sync local cache with Hub
optimum-cli neuron cache synchronize

高階用法

在訓練迴圈中使用快取

如果您不使用 NeuronTrainer 類，您仍然可以在自定義訓練迴圈中利用快取系統。當您需要更多地控制訓練過程，或者在與自定義訓練框架整合時，同時又想從快取的編譯結果中受益，這會非常有用。

何時使用此方法

自定義訓練迴圈不符合 NeuronTrainer 模式
需要精細控制的高階最佳化場景

注意：對於大多數用例，NeuronTrainer 會自動處理快取，是推薦的方法。

from optimum.neuron.cache import hub_neuronx_cache, synchronize_hub_cache
from optimum.neuron.cache.entries import SingleModelCacheEntry
from optimum.neuron.cache.training import patch_neuron_cc_wrapper

# Create cache entry
cache_entry = SingleModelCacheEntry(model_id, task, config, neuron_config)

# The NeuronX compiler will use the Hugging Face Hub cache system
with patch_neuron_cc_wrapper():
    # The compiler will check the specified remote cache for pre-compiled NEFF files
    with hub_neuronx_cache(entry=cache_entry, cache_repo_id="my-org/cache"):
        model = training_loop()  # Will use specified cache

# Synchronize local cache with Hub
synchronize_hub_cache(cache_repo_id="my-org/cache")

快取查詢

推理快取包含一個登錄檔，讓您可以在嘗試編譯之前搜尋相容的預編譯模型。這對於您希望完全避免編譯的推理場景尤其有用。

optimum-cli neuron cache lookup meta-llama/Llama-2-7b-chat-hf

示例輸出

*** 1 entries found in cache ***
task: text-generation
batch_size: 1, sequence_length: 2048
num_cores: 24, precision: fp16
compiler_version: 2.12.54.0
checkpoint_revision: c1b0db933684edbfe29a06fa47eb19cc48025e93

重要：找到條目並不能保證快取命中。您的確切配置必須與快取的引數相匹配，包括編譯器版本和模型修訂版本。

CI/CD 整合

快取系統在自動化環境中無縫工作

環境變數：在 CI 工作流中使用 CUSTOM_CACHE_REPO 指定快取倉庫

# In your CI configuration
CUSTOM_CACHE_REPO="your-org/your-cache" python train.py

認證：確保您的 CI 環境有權訪問您的私有快取倉庫

設定具有適當讀/寫許可權的 HF_TOKEN 環境變數
對於 GitHub Actions，將其儲存為倉庫機密

最佳實踐:

為不同環境（開發/預釋出/生產）使用單獨的快取倉庫
在設定自動化工作流時，考慮快取倉庫的許可權
在長期執行的 CI 工作流中監控快取倉庫的大小

問題排查

“快取倉庫不存在”

Fix: Check repository name and login status
→ huggingface-cli login
→ Verify repo format: org/repo-name

“計算圖將被重新編譯”

Cause: No cached model matches your exact configuration
Fix: Use lookup to find compatible configurations
→ optimum-cli neuron cache lookup MODEL_ID

訓練期間快取未上傳

Cause: No write permissions to cache repository  
Fix: Verify access and authentication
→ huggingface-cli whoami
→ Check cache repo permissions

下載緩慢

Cause: Large compiled models (GBs) downloading
Fix: Ensure good internet connection
→ Monitor logs for download progress

清除損壞的本地快取

rm -rf /var/tmp/neuron-compile-cache/*

AWS Trainium 和 Inferentia