使用 Hugging Face Transformers 和 Amazon SageMaker 部署 GPT-J 6B 進行推理

距今大約 6 個月前，EleutherAI 釋出了 GPT-J 6B，這是 OpenAI GPT-3 的一個開源替代品。GPT-J 6B 是 EleutherAI 的 GPT-NEO 系列的 60 億引數後繼者，該系列是基於 GPT 架構用於文字生成的 transformer 語言模型家族。

EleutherAI 的主要目標是訓練一個與 GPT-3 規模相當的模型，並以開放許可證向公眾提供。

在過去的 6 個月裡，`GPT-J` 引起了研究人員、資料科學家甚至軟體開發人員的極大興趣，但將其部署到生產環境中用於實際用例和產品仍然非常具有挑戰性。

有一些託管解決方案可將 `GPT-J` 用於生產工作負載，例如 Hugging Face 推理 API，或者使用 EleutherAI 的 6b playground 進行實驗，但關於如何輕鬆地將其部署到您自己環境中的示例較少。

在這篇博文中，您將學習如何使用 Amazon SageMaker 和 Hugging Face 推理工具包，通過幾行程式碼輕鬆部署 `GPT-J`，從而使用常規大小的 GPU 例項（NVIDIA T4，約 500 美元/月）實現可擴充套件、可靠和安全的即時推理。

但在我們深入探討之前，我想解釋一下為什麼將 `GPT-J` 部署到生產環境中具有挑戰性。

---

背景

這個 60 億引數模型的權重佔用約 24GB 記憶體。要以 float32 格式載入它，至少需要 2 倍模型大小的 CPU 記憶體：1 倍用於初始權重，另 1 倍用於載入檢查點。因此，對於 `GPT-J`，僅載入模型就需要至少 48GB 的 CPU 記憶體。

為了讓模型更易於使用，EleutherAI 也提供了 float16 權重，並且 `transformers` 有了新的選項來減少載入大型語言模型時的記憶體佔用。綜合所有這些，載入模型大約需要 12.1GB 的 CPU 記憶體。

from transformers import GPTJForCausalLM
import torch

model = GPTJForCausalLM.from_pretrained(
    "EleutherAI/gpt-j-6B",
        revision="float16",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
)

這個例子的一個問題是，模型載入到記憶體並準備好使用需要很長時間。在我的實驗中，在 `P3.2xlarge` AWS EC2 例項上使用上面的程式碼片段載入模型需要 `3 分 32 秒`（模型未儲存在磁碟上）。透過將模型預先儲存在磁碟上，可以縮短這個時間，載入時間減少到 `1 分 23 秒`，這對於需要考慮可擴充套件性和可靠性的生產工作負載來說仍然很長。

例如，Amazon SageMaker 對請求響應有 60 秒的限制，這意味著模型載入和預測執行必須在 60 秒內完成，在我看來，這對於保持模型/端點的可擴充套件性和工作負載的可靠性非常有意義。如果您有更長的預測任務，可以使用批次轉換。

在 Transformers 中，使用 `from_pretrained` 方法載入的模型遵循 PyTorch 的推薦做法，對於 BERT 來說大約需要 `1.97 秒` [參考]。PyTorch 提供了另一種儲存和載入模型的方法，即使用 `torch.save(model, PATH)` 和 `torch.load(PATH)`。

“用這種方式儲存模型將使用 Python 的 pickle 模組儲存整個模組。這種方法的缺點是，序列化的資料與儲存模型時使用的特定類和確切的目錄結構繫結在一起。”

這意味著，當我們使用 `transformers==4.13.2` 儲存模型時，在嘗試使用 `transformers==4.15.0` 載入時可能會不相容。然而，以這種方式載入模型可將載入時間減少約 **12 倍**，對於 BERT 來說，時間降至 `0.166 秒`。

將此方法應用於 `GPT-J` 意味著我們可以將載入時間從 `1 分 23 秒` 減少到 `7.7 秒`，速度提高了約 10.5 倍。

教程

透過這種儲存和載入模型的方法，我們為 `GPT-J` 實現了與生產場景相容的模型載入效能。但我們需要記住，我們需要對齊

在使用 `torch.save(model,PATH)` 儲存模型和使用 `torch.load(PATH)` 載入模型時，對齊 PyTorch 和 Transformers 的版本，以避免不相容性。

使用 `torch.save` 儲存 `GPT-J`

為了建立我們與 `torch.load()` 相容的模型檔案，我們使用 Transformers 和 `from_pretrained` 方法載入 `GPT-J`，然後使用 `torch.save()` 儲存它。

from transformers import AutoTokenizer,GPTJForCausalLM
import torch

# load fp 16 model
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
# save model with torch.save
torch.save(model, "gptj.pt")

現在我們能夠使用 `torch.load()` 載入我們的 `GPT-J` 模型來執行預測。

from transformers import pipeline
import torch

# load model
model = torch.load("gptj.pt")
# load tokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

# create pipeline
gen = pipeline("text-generation",model=model,tokenizer=tokenizer,device=0)

# run prediction
gen("My Name is philipp")
#[{'generated_text': 'My Name is philipp k. and I live just outside of Detroit....

---

為 Amazon SageMaker 即時端點建立 `model.tar.gz`

既然我們可以快速載入模型並進行推理，讓我們把它部署到 Amazon SageMaker 上。

有兩種方式可以部署 transformers 到 Amazon SageMaker。您可以直接“從 Hugging Face Hub 部署模型”，或者“使用儲存在 S3 上的 `model_data` 部署模型”。由於我們不使用預設的 Transformers 方法，我們需要選擇第二種方案，並使用儲存在 S3 上的模型來部署我們的端點。

為此，我們需要建立一個 `model.tar.gz` 檔案，其中包含我們的模型權重和推理所需的其他檔案，例如 `tokenizer.json`。

我們提供了已上傳且可公開訪問的 `model.tar.gz` 檔案，這些檔案可以與 `HuggingFaceModel` 一起使用，以將 `GPT-J` 部署到 Amazon SageMaker。

請參閱“將 `GPT-J` 部署為 Amazon SageMaker 端點”部分了解如何使用它們。

如果您仍然想或需要建立自己的 `model.tar.gz`，例如出於合規性準則，您可以使用輔助指令碼 convert_gpt.py 來實現此目的，該指令碼會建立 `model.tar.gz` 並將其上傳到 S3。

# clone directory
git clone https://github.com/philschmid/amazon-sagemaker-gpt-j-sample.git

# change directory to amazon-sagemaker-gpt-j-sample
cd amazon-sagemaker-gpt-j-sample

# create and upload model.tar.gz
pip3 install -r requirements.txt
python3 convert_gptj.py --bucket_name {model_storage}

`convert_gpt.py` 應該會打印出一個類似於此的 S3 URI。`s3://hf-sagemaker-inference/gpt-j/model.tar.gz`。

將 `GPT-J` 部署為 Amazon SageMaker 端點

為了部署我們的 Amazon SageMaker 端點，我們將使用 Amazon SageMaker Python SDK 和 `HuggingFaceModel` 類。

下面的程式碼片段使用了 `get_execution_role`，它僅在 Amazon SageMaker Notebook 例項或 Studio 中可用。如果您想在這些環境之外部署模型，請檢視文件。

`model_uri` 定義了我們的 `GPT-J` 模型檔案的位置。我們將使用由我們提供的公開可用的檔案。

from sagemaker.huggingface import HuggingFaceModel
import sagemaker

# IAM role with permissions to create endpoint
role = sagemaker.get_execution_role()

# public S3 URI to gpt-j artifact
model_uri="s3://huggingface-sagemaker-models/transformers/4.12.3/pytorch/1.9.1/gpt-j/model.tar.gz"

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    model_data=model_uri,
    transformers_version='4.12.3',
    pytorch_version='1.9.1',
    py_version='py38',
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1, # number of instances
    instance_type='ml.g4dn.xlarge' #'ml.p3.2xlarge' # ec2 instance type
)

如果您想使用自己的 `model.tar.gz`，只需將 `model_uri` 替換為您的 S3 URI 即可。

部署過程大約需要 3-5 分鐘。

執行預測

我們可以使用由 `.deploy` 方法建立的 `predictor` 例項來執行預測。為了向我們的端點發送請求，我們使用 `predictor.predict` 和我們的 `inputs`。

predictor.predict({
    "inputs": "Can you please let us know more details about your "
})

如果您想使用額外的 `kwargs` (如 `min_length`) 來自定義您的預測，請檢視下面的“使用最佳實踐”部分。

使用最佳實踐

在使用生成模型時，大多數情況下您希望配置或自定義您的預測以滿足您的需求，例如使用束搜尋、配置生成序列的最大或最小長度，或調整溫度以減少重複。Transformers 庫提供了不同的策略和 `kwargs` 來實現這一點，Hugging Face 推理工具包透過請求負載中的 `parameters` 屬性提供了相同的功能。下面您可以找到如何生成不帶引數的文字、使用束搜尋以及使用自定義配置的示例。如果您想了解不同的解碼策略，請檢視這篇部落格文章。

預設請求

這是一個使用 `greedy` 搜尋的預設請求示例。

第一次請求後的推理時間：`3 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your "
})

束搜尋請求

這是一個使用 `beam` 搜尋（5 個束）的請求示例。

首次請求後的推理時間：`3.3 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your ",
  "parameters" : {
    "num_beams": 5,
  }
})

引數化請求

這是一個使用自定義引數的請求示例，例如使用 `min_length` 生成至少 512 個 token。

第一次請求後的推理時間：`38 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your ",
  "parameters" : {
    "max_length": 512,
    "temperature": 0.9,
  }
})

少樣本示例（高階）

這是一個如何使用 `eos_token_id` 在特定標記處停止生成的示例，例如 `\n`、`.` 或 `###` 用於少樣本預測。下面是一個為關鍵詞生成推文的少樣本示例。

首次請求後的推理時間：`15-45 秒`

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

end_sequence="###"
temperature=4
max_generated_token_length=25
prompt= """key: markets
tweet: Take feedback from nature and markets, not from people.
###
key: children
tweet: Maybe we die so we can come back as children.
###
key: startups
tweet: Startups shouldn’t worry about how to put out fires, they should worry about how to start them.
###
key: hugging face
tweet:"""

predictor.predict({
    'inputs': prompt,
  "parameters" : {
    "max_length": int(len(prompt) + max_generated_token_length),
    "temperature": float(temperature),
    "eos_token_id": int(tokenizer.convert_tokens_to_ids(end_sequence)),
    "return_full_text":False
  }
})

---

要刪除您的端點，您可以執行。

predictor.delete_endpoint()

結論

我們成功地使用 Amazon SageMaker 部署了 `GPT-J`，這是一個由 EleutherAI 建立的 60 億引數語言模型。我們將模型載入時間從 3.5 分鐘減少到 8 秒，以便能夠執行可擴充套件、可靠的推理。

請記住，使用 `torch.save()` 和 `torch.load()` 可能會產生不相容問題。如果您想了解更多關於擴充套件 Amazon SageMaker 端點的資訊，請檢視我的另一篇部落格文章：“MLOps：使用 Hub 和 SageMaker Pipelines 實現端到端的 Hugging Face Transformers”。

---

感謝閱讀！如果您有任何問題，請隨時透過 Github 或在 論壇上與我聯絡。您也可以在 Twitter 或 LinkedIn 上與我聯絡。