Jetson AGX Orin 64GB で Phi-4 を動かす

Phi-4がHugginFaceで公開されたとのことで、Jetson AGX Orin 64GBの試運転を兼ねて動かしてみました。

https://huggingface.co/microsoft/phi-4

動作環境

• Jetson AGX Orin 64GB
• JetPack 6.1
• L4T 36.4
• NVIDIA Container

推論時のリソース消費

モデルパラメータ14.7Bに対して、VRAM 27.5GB消費でした。
Jetson AGX Orin 32GBでもぎりぎり動作しそうです。

Jetson AGX Orin 64GBはストレージが64GB eMMC 5.1しかなく、VRAMよりもこちらの制約の方が厳しかったです。本検証で使用率99%ぐらいでした。。

推論コード

Getting startedをそのまま実行。

outputs = pipeline(messages, max_new_tokens=128)の１行で65秒ほど掛かりました。

import transformers

pipeline = transformers.pipeline(
    "text-generation",
    model="microsoft/phi-4",
    model_kwargs={"torch_dtype": "auto"},
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are a medieval knight and must provide explanations to modern people."},
    {"role": "user", "content": "How should I explain the Internet?"},
]

outputs = pipeline(messages, max_new_tokens=128)
print(outputs[0]["generated_text"][-1])

# Output:
# {'role': 'assistant', 'content': 'Ah, greetings, traveler from the future!
# Thou art inquiring about a wondrous and mystical realm known as the Internet.
# Allow me to endeavor to explain this marvel in terms that might be more familiar
# to a knight of the realm.\n\nImagine, if thou wilt, a vast and invisible kingdom,
# not bound by the physical constraints of land or sea. This kingdom is filled with
# countless scrolls, tomes, and messages, all accessible at the mere wave of a hand.
# It is akin to a grand library, but one that stretches across the entire world,
# and it is ever-growing, with new knowledge and tales being added every
# moment.\n\nIn'}

コンテナ設定

ベースイメージを探すのに一番時間が掛かりました。

新しいtransformersライブラリを動作させるためには、JetPack 6.1を使うのが手っ取り早いです。しかし、公式のNVIDIA L4T PyTorchイメージはJetPack 6.1に対応していません。

この辺りのCUDAイメージからビルドするのも面倒だなー、と思っていたら先人がいました！
JetPack 6.1向けのL4T PyTorchイメージを拝借しました。

# .devcontainer/Dockerfile
FROM gotapway/l4t-jetpack:r36.4.0-pth2.5.0-py3.10

RUN pip install transformers==4.48.0
RUN pip install accelerate==1.2.1
RUN pip install flash_attn==2.7.2.post1

// .devcontainer/devcontainer.json
{
    "name": "Dev Container with NVIDIA Runtime",
    "build": {
        "dockerfile": "Dockerfile",
        "context": "..",
    },
    "runArgs": [
        "--runtime", "nvidia"
    ],
    "postCreateCommand": "echo 'Container created!'",
    "remoteUser": "root"
}

感想

Jetson AGX Orin 64GBの動作確認でPhi-4を動かしてみました。
JetPack 6.1については、（少なくともコンテナ周りの）エコシステムが発展途上という感じです。
使用したいライブラリのバージョンが古くても問題ないのであれば、現時点ではJetPack 5.xを使うのが良さそうです。