Nemotron Nano를 NIM으로 추론하기

Nemotron Nano를 NIM으로 추론하기

Nemotron Nano는 엔비디아의 NIM(NVIDIA Inference Microservice) 으로 바로 서빙할 수 있어요. OpenAI 호환 API로 접근할 수 있어서 기존 OpenAI SDK를 그대로 쓸 수 있다는 게 큰 장점이에요.

출처: NVIDIA NIM: llama-3.1-nemotron-nano-8b-v1

OpenAI 호환 API로 호출하기

from openai import OpenAI

client = OpenAI(
  base_url = "https://integrate.api.nvidia.com/v1",
  api_key = "$NVIDIA_API_KEY"
)

completion = client.chat.completions.create(
  model="nvidia/llama-3.1-nemotron-nano-8b-v1",
  messages=[{"role":"system","content":"You are a helpful assistant."},
            {"role":"user","content":"Explain how a transformer neural network works."}],
  temperature=0.6,
  top_p=0.95,
  max_tokens=4096,
  stream=False
)
print(completion.choices[0].message)

Docker로 로컬 NIM 띄우기

NGC에 로그인하고 로컬 캐시를 만들어 컨테이너를 실행해요.

export NGC_API_KEY=<PASTE_API_KEY_HERE>
export LOCAL_NIM_CACHE=~/.cache/nim
mkdir -p "$LOCAL_NIM_CACHE"
docker run -it --rm     --gpus all     --shm-size=16GB     -e NGC_API_KEY     -v "$LOCAL_NIM_CACHE:/opt/nim/.cache"     -u $(id -u)     -p 8000:8000     nvcr.io/nim/nvidia/llama-3.1-nemotron-nano-8b-v1:latest

테스트는 /v1/chat/completions 엔드포인트로 curl 하면 돼요.

curl -X 'POST' 'http://0.0.0.0:8000/v1/chat/completions'   -H 'accept: application/json' -H 'Content-Type: application/json'   -d '{"model": "nvidia/llama-3.1-nemotron-nano-8b-v1",
       "messages": [{"role":"user","content":"Explain how a transformer neural network works."}],
       "max_tokens": 64}'

더 알아보기