Skip to main content

Speculative decoding engine with local and remote target model execution

Project description

MobiTree 실행 가이드 (환경 구축 + UI 사용)

이 문서는 처음 설치하는 사람도 복사/붙여넣기만으로 실행할 수 있게 작성했습니다.

  • 대상: Ubuntu/Linux + NVIDIA GPU 환경
  • 범위: 드라이버/CUDA 확인, venv 구성, 의존성 설치, Target 실행, UI 실행, UI 사용법

1) 사전 준비

1-1. 필수 소프트웨어

  • NVIDIA Driver 설치
  • Python 3.10 이상
  • git
  • (선택) tmux 또는 screen (서버 장시간 실행 시 편리)

1-2. 드라이버/CUDA 확인 (복붙)

nvidia-smi
python3 --version

정상이라면:

  • nvidia-smi가 GPU/Driver 정보를 출력
  • Python 버전이 출력

참고:

  • 이 저장소는 torch==2.7.0+cu128 기준 requirements.txt를 사용합니다.
  • 시스템에 CUDA Toolkit(nvcc)이 없어도, PyTorch wheel + NVIDIA driver 조합으로 동작할 수 있습니다.
  • requirements.txt에 PyTorch cu128 인덱스가 포함되어 있으므로 pip install -r requirements.txt만 실행하면 됩니다.

2) 프로젝트 받기

git clone git@github.com:PJChoi1/MobiTree.git MobiTree
cd MobiTree
git checkout seperate

이미 코드가 있다면 이 단계는 생략하세요.


3) 가상환경(venv) + 의존성 설치

아래 블록을 한 번에 그대로 실행하세요.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
pip install -r requirements.txt

설치 확인:

python -c "import torch; print('torch:', torch.__version__); print('cuda:', torch.cuda.is_available())"

cuda: True가 나오면 GPU 인식이 정상입니다.


4) 실행 방법 (가장 많이 쓰는 방법)

아래 순서대로 진행하세요. (터미널 A: UI)

4-1. 1단계: UI 실행

터미널 A:

source .venv/bin/activate
python3 -m uvicorn chat_ui.main:app --host 0.0.0.0 --port 8000

브라우저 접속:

4-2. 2단계: UI 접속 확인

브라우저에서 UI가 열리면 상단 Server 영역과 Add Server 버튼이 보이는지 확인하세요.

4-3. 3단계: UI에서 Server 추가

브라우저 UI에서:

  1. Add Server 클릭
  2. Server Name 입력 (예: icnl-server)
  3. IP Address 입력 (예: 163.152.163.152)
  4. Port 입력 (예: 26001)
  5. Add Server 클릭

5) UI 사용법 (처음 사용자용)

5-1. 서버/모델/Quantization 선택

이미 등록된 서버를 기준으로 아래 순서로 선택하세요.

  1. Server 드롭다운에서 사용할 서버 선택
  2. Server Model에서 target 모델 선택
  3. Draft Model에서 draft 모델 선택
  4. Server Q/Draft Q에서 quantization(4bit/8bit) 선택
  5. Start를 눌러 런타임 시작 후 질문 전송

5-2. 주요 버튼/설정값 의미

  • Start: 채팅 런타임 시작
  • Stop: 채팅 런타임 중지 + target 모델 unload 요청
  • Profile LLM: reference cache/profile 갱신
  • Token Source Coloring: 토큰 출처별 색상 표시
  • Algorithm: 디코딩 전략 선택 (MobiTree, Server-Only, Server-Only-AR, OPT-Tree, Fixed-tree)
  • Mode: 실행 모드 선택 (Chat은 대화, Benchmark는 평가 중심)
  • Max New Tokens: 한 번 응답에서 생성할 최대 토큰 수
  • Dataset: Benchmark 모드에서 사용할 평가 데이터셋 선택

5-3. UI에서 add한 server name으로 Target 실행

핵심: 4-3에서 추가한 Server Name과 Target 실행 시 --server-name 값을 같게 맞추세요.

예를 들어 4-3에서 Server Name=icnl-server로 추가했다면, 터미널 B에서:

source .venv/bin/activate
./run_target.sh --host 0.0.0.0 --port 26001 --device-map auto --load-in-8bit --lazy-load --enable-auto-target-profile --server-name icnl-server

설명:

  • --host/--port: 4-3에서 입력한 IP Address/Port와 동일해야 연결됩니다.
  • --server-name: 4-3에서 입력한 Server Name과 동일하게 사용하세요.
  • --lazy-load: 최초 접속 시 모델 로드
  • --enable-auto-target-profile: target profile 없으면 자동 생성

6) 캐시/프로파일 파일 위치

  • Target profile: data/profile/profile_target_<server-name>_<model>_tq-<quant>.json
  • Reference cache: data/reference/ref_<server-name>_<base>_<device>_<draft>_tq-<...>_dq-<...>_mt_bench_<metric>_<mode>_*.json

파일명이 맞아야 재사용됩니다.


7) autodraft Python API

autodraft 는 위 MobiTree 런타임을 감싸는 얇은 Python wrapper 입니다. 기존 shell script (run_target.sh, run_mt_bench_sd.sh 등) 와 CLI 동작은 그대로 유지되고, 추가로 다음 두 가지 사용법이 지원됩니다.

7-1. Repo clone 후 바로 import

git clone https://github.com/PJChoi1/MobiTree.git
cd MobiTree
# 터미널 A: target server (Python only, ./run_target.sh 대체)
python examples/local_target.py
# 터미널 B: draft side (실제 generated text)
python examples/local_draft.py

또는 Python 코드에서:

# 터미널 A — target server
from autodraft import serve_target

serve_target(
    host="0.0.0.0",
    port=26001,
    target_model="meta-llama/Llama-3.2-1B-Instruct",
    quantization="8bit",
    server_name="autodraft",
    hf_token=None,  # 또는 "hf_xxx"
)  # 이 호출은 서버 loop라 영원히 block 됩니다.

# 터미널 B — draft side
from autodraft import Autodraft

engine = Autodraft(
    draft_model="meta-llama/Llama-3.2-1B-Instruct",
    target_model="meta-llama/Llama-3.2-1B-Instruct",
    remote=False,
    target_host="127.0.0.1",
    target_port=26001,
)
result = engine.run("input text", server_name="autodraft")
print(result["generated_text"])

server_name 은 양쪽이 일치해야 profile/reference cache 가 재사용됩니다 (default: "autodraft").

주의:

  • pip install 없이 import 하려면 Python 의 sys.path 에 repo root 가 포함되어야 합니다. examples/basic_usage.py 는 자기 자신의 부모 디렉토리를 sys.path 에 끼워넣는 4줄짜리 bootstrap 이 들어 있어 python examples/basic_usage.py 로 그대로 실행됩니다.
  • 본인 코드에서 pip install 없이 쓰려면 다음 중 하나를 사용하세요:
    # (a) repo root 에서 -m 으로 실행
    python -m examples.basic_usage
    
    # (b) PYTHONPATH 에 repo path 를 추가
    export PYTHONPATH=/path/to/MobiTree:$PYTHONPATH
    python my_script.py
    
    # (c) 본인 스크립트 첫 줄에 sys.path bootstrap 을 추가 (examples 와 동일한 패턴)
    
  • MobiTree 의 architecture 상 draft 측은 항상 target server 와 socket 으로 통신하기 때문에, remote=False 로 설정해도 같은 머신 (127.0.0.1:26001 기본값) 에 target server 가 떠 있어야 합니다. 먼저 5-3 의 ./run_target.sh 를 실행해 두세요.
  • gated 모델 (meta-llama/* 등) 을 사용하려면 HF access token 이 필요합니다. 두 가지 방법이 있습니다:
    # (a) 생성자에 직접 전달
    engine = Autodraft(draft_model=..., target_model=..., hf_token="hf_xxx")
    
    # (b) 환경변수로 설정 (HF_TOKEN 또는 HUGGING_FACE_HUB_TOKEN)
    # bash:
    #   export HF_TOKEN=hf_xxx
    
    hf_token 인자는 repr 출력에서 '***' 로 마스킹됩니다. 이 token 은 run_local 안에서 HF_TOKEN / HUGGING_FACE_HUB_TOKEN env var 로 export 되어 run_target.sh 로 띄운 target server 도 같은 token 을 사용할 수 있게 됩니다 (target 쪽은 server 프로세스의 env 에 자체적으로 token 을 export 하면 됩니다).

7-2. Editable install

pip install -e .
# 4/8-bit quantization 도 쓰려면:
pip install -e ".[quant]"
# 테스트/lint/build 까지:
pip install -e ".[dev]"

이렇게 하면 repo 밖 어디서든 import autodraft 가 가능합니다.

7-3. PyPI 설치 (배포 후)

pip install autodraft-sd

Distribution 이름은 autodraft-sd 입니다. PyPI 에 autodraft 이름이 이미 사용 중이라 충돌을 피하기 위해 -sd (speculative decoding) 접미사를 붙였습니다. import 이름은 그대로 autodraft 이므로 코드에서는 from autodraft import Autodraft 를 쓰면 됩니다 (Pillow 가 pip install pillow 인데 from PIL 로 import 하는 것과 같은 패턴).

PyPI wheel 에는 wrapper API (autodraft/) 와 함께 speculative decoding runtime (evaluation/, opt_classic/) 도 포함됩니다. 따라서 source checkout 없이 pip install autodraft-sd 한 번으로 engine.run(...)serve_target(...) 모두 동작합니다 (단, GPU 와 적절한 PyTorch wheel 이 필요 — 7-5 참조).

profile/reference cache 디렉토리는 기본적으로 ~/.cache/autodraft/data 에 생성됩니다. 다른 위치를 쓰려면:

export MOBITREE_DATA_DIR=/path/to/your/cache

chat_ui/, data/ (벤치마크 데이터셋), result/ 등 연구용 자산은 wheel 에 포함되지 않습니다. 그것들은 GitHub source checkout 으로만 접근하세요.

4/8-bit quantization (bitsandbytes 의존) 은 별도 extra 입니다:

pip install "autodraft-sd[quant]"

7-4. Remote target 사용

서버 측은 기존 방식 그대로 띄웁니다:

./run_target.sh --host 0.0.0.0 --port 26001 \
                --base-model-path meta-llama/Llama-2-70b-chat-hf \
                --draft-model-path meta-llama/Llama-2-7b-chat-hf \
                --load-in-8bit --lazy-load --enable-auto-target-profile \
                --server-name my-server

클라이언트 측:

from autodraft import Autodraft

engine = Autodraft(
    draft_model="...",
    target_model="...",
    remote=True,
    target_host="TARGET_SERVER_IP",
    target_port=26001,
)

result = engine.run("input text", proactive=True, cs=None)

7-5. PyTorch 설치에 대한 안내

autodraft 의 base wheel 은 PyTorch 를 의존성으로 가지지 않습니다 (autodraft/ 자체는 stdlib 만 사용). in-process 실행에 필요한 torch / transformers / accelerate / fastchat 등은 [local] extra 또는 기존 requirements.txt 로 설치하세요.

특히 PyTorch 는 CUDA 버전마다 wheel 이 다르므로, 본 repo 의 requirements.txt 에 명시된 torch==2.7.0+cu128 과 같이 환경에 맞는 wheel 을 직접 설치하는 것을 권장합니다.


8) PyPI 배포 (Trusted Publishing)

PyPI 배포는 모든 push 에 대해 자동으로 수행되지 않습니다. 다음 두 가지가 모두 만족할 때만 배포됩니다.

  1. pyproject.tomlversion 이 올라가 있다.
  2. 그 version 에 매칭되는 git tag (vX.Y.Z) 가 push 된다.

8-1. Trusted Publishing 일회성 설정

  1. PyPI 계정 생성 (없다면)
  2. https://pypi.org/manage/account/publishing/ 에서 Add a new pending publisher 선택
  3. 다음 값을 입력:
    • PyPI Project Name: autodraft-sd (또는 pyproject.tomlname 값)
    • Owner: GitHub repository owner (예: PJChoi1)
    • Repository name: MobiTree
    • Workflow name: publish.yml
    • Environment name: pypi
  4. GitHub repository 의 Settings → Environments 에서 pypi environment 를 생성

이렇게 하면 첫 배포부터 API token 없이 OIDC 기반으로 publish 가 됩니다. PYPI_API_TOKEN secret 은 사용하지 않습니다.

8-2. 새 버전 배포

# pyproject.toml 의 version 을 0.1.0 → 0.1.1 등으로 올린 뒤
git add pyproject.toml
git commit -m "release: 0.1.1"
git tag v0.1.1
git push origin v0.1.1

publish.yml workflow 가 build → twine check → PyPI publish 를 자동으로 수행합니다. workflow 는 시작 시 git tag 와 pyproject.toml 의 version 일치 여부를 검증하므로, 둘이 어긋나면 배포가 거부됩니다.

PyPI 는 같은 version 을 다시 업로드할 수 없습니다. 코드를 수정할 때마다 자동 배포하지 마시고, version 을 올리고 tag 를 push 할 때만 배포하세요.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

autodraft_sd-0.1.0.tar.gz (496.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

autodraft_sd-0.1.0-py3-none-any.whl (546.4 kB view details)

Uploaded Python 3

File details

Details for the file autodraft_sd-0.1.0.tar.gz.

File metadata

  • Download URL: autodraft_sd-0.1.0.tar.gz
  • Upload date:
  • Size: 496.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for autodraft_sd-0.1.0.tar.gz
Algorithm Hash digest
SHA256 80cd3323786b94a48a04fa9fed4d9ac084df7869a46a7f7f28892e135d8d548a
MD5 fb420f6705ce1d088dd40e50c7b36500
BLAKE2b-256 0943d5302d4d5275751ac35b5e716f6e4ca4c1815cf26cb0cc783c189623a9cf

See more details on using hashes here.

Provenance

The following attestation bundles were made for autodraft_sd-0.1.0.tar.gz:

Publisher: publish.yml on PJChoi1/MobiTree

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file autodraft_sd-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: autodraft_sd-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 546.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for autodraft_sd-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 b60f1e91710ad8a1cabcef7885cc73cb5c6721d2b776d26bc67565007c094e58
MD5 69488b67698b06fa016a3c5eb555d1bd
BLAKE2b-256 18e3ed3ecd1e5b314a0804ddabba4490b293ceba7d670dd9ee932c2cd2742792

See more details on using hashes here.

Provenance

The following attestation bundles were made for autodraft_sd-0.1.0-py3-none-any.whl:

Publisher: publish.yml on PJChoi1/MobiTree

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page