autodraft-sd 0.1.7

Speculative decoding engine with local and remote target model execution

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 [e-mail]: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

브라우저 접속:

• http://localhost: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 examples/target.py
# 터미널 B: draft side (실제 generated text)
python examples/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",
    target_host="127.0.0.1",
    target_port=26001,
)
result = engine.run("input text", server_name="autodraft")
print(result["generated_text"])
print(result["stats"])  # avg tree width/depth/nnodes, accept length, cost breakdown

MobiTree 는 항상 split-process 입니다 — draft 와 target 은 별도 프로세스로 실행되고 socket 으로 통신합니다. 같은 머신에서는 target_host="127.0.0.1" (기본값), 다른 머신에서는 target 머신의 IP 를 적으면 됩니다.

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

주의:

• pip install 없이 import 하려면 Python 의 sys.path 에 repo root 가 포함되어야 합니다. examples/draft.py 와 examples/target.py 는 자기 자신의 부모 디렉토리를 sys.path 에 끼워넣는 4줄짜리 bootstrap 이 들어 있어 python examples/draft.py 로 그대로 실행됩니다.
• 본인 코드에서 pip install 없이 쓰려면 다음 중 하나를 사용하세요:

  # (a) repo root 에서 -m 으로 실행
  python -m examples.draft
  
  # (b) PYTHONPATH 에 repo path 를 추가
  export PYTHONPATH=/path/to/MobiTree:$PYTHONPATH
  python my_script.py
  
  # (c) 본인 스크립트 첫 줄에 sys.path bootstrap 을 추가 (examples 와 동일한 패턴)
  

• target server 가 떠 있어야 draft 가 동작합니다. 먼저 python examples/target.py (또는 5-3 의 ./run_target.sh) 로 target 을 실행해 두세요.
• 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 디렉토리는 기본적으로 스크립트를 실행한 디렉토리의 ./data/ 에 생성됩니다 (즉 python my_script.py 했다면 그 명령이 실행된 디렉토리의 data/profile/, data/reference/). source-checkout 사용자가 repo root 에서 실행하던 동작과 자연스럽게 일치합니다.

처음 실행하면 (cache 가 없으면) target latency profile 과 draft latency profile 을 자동으로 굽습니다 (수십 초~수 분). 두 번째 실행부터는 cache hit 하여 빠르게 시작됩니다. cache 파일명은 GPU 이름을 포함하므로 다른 GPU 로 옮기면 새로 굽습니다. GPU 이름을 직접 지정하려면:

export MOBITREE_DEVICE_NAME="rtx-5080"   # 자유 라벨

지정하지 않으면 torch.cuda.get_device_name(0) 으로 자동 감지하고, GPU 가 없으면 "unknown-gpu" 라벨을 사용합니다.

다른 cache 위치 (예: 여러 프로젝트가 cache 를 공유하게 하고 싶을 때) 를 쓰려면:

export MOBITREE_DATA_DIR=/path/to/shared/cache

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

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

pip install "autodraft-sd[quant]"

trade-off 그래프를 PNG 로도 저장하려면 matplotlib extra:

pip install "autodraft-sd[plot]"

(matplotlib 가 없으면 JSON 만 저장됩니다.)

7-3-1. engine.run() 파라미터와 결과

result = engine.run(
    input_text="...",
    proactive=False,
    cs="balanced",            # "tps" / "balanced" (default) / "cost", 또는 0~100 숫자
    cost="total_cost",        # "total_cost" (default) / "api_cost" / "energy_total"
                              # / "draft_energy" / "target_energy"
    save_tradeoff=True,       # cs 0~100 reference trade-off curve 를 저장
    tradeoff_dir=None,        # 지정 안 하면 $MOBITREE_DATA_DIR/tradeoff
    server_name="autodraft",  # target 의 server_name 과 일치해야 cache 재사용됨
    # 그 외 run_draft 의 ~70 개 옵션은 **kwargs 로 그대로 forward 됩니다.
)

반환값 구조:

{
    "generated_text": "...",        # 모델 출력 텍스트
    "input_text": "...",
    "proactive": False,
    "cs": "balanced",
    "cost": "total_cost",
    "algorithm": "MobiTree",
    "stats": {                      # 한 줄 요약 통계
        "total_steps": 15,
        "total_new_tokens": 121,
        "total_time_seconds": 1.98,
        "tokens_per_second": 61.10,
        "tokens_per_step": 8.07,
        "avg_tree_width": 6.4,
        "avg_tree_depth": 7.2,
        "avg_nnodes": 50.5,         # 한 트리당 평균 노드 수
        "avg_accept_length": 7.07,
        "acceptance_ratio_avg": 0.98,
        "total_cost": 0.000125,     # 선택한 cost objective 의 총합 ($ 또는 kWh)
        "api_cost": 0.0,
        "draft_cost": 0.000031,
        "target_cost": 0.000094,
    },
    "tradeoff_files": {
        "json": "/path/to/data/tradeoff/tradeoff_<ts>.json",
        "png":  "/path/to/data/tradeoff/tradeoff_<ts>.png",  # matplotlib 있을 때만
    },
    "answer_row": { ... },          # 통합 결과의 answers[-1]
    "raw_result": { ... },          # 통합 결과 JSON 전체 (latency_statistics, accept_stats 등 포함)
}

7-4. 다른 머신에 떠 있는 target 사용

target 측은 python examples/target.py 로 띄우거나 (server 머신에서) 기존 shell 방식 그대로:

./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="...",
    target_host="TARGET_SERVER_IP",
    target_port=26001,
)

result = engine.run(
    "input text",
    proactive=True,
    cs="balanced",        # "tps" / "balanced" / "cost", or 0~100 number
    cost="total_cost",    # also: api_cost / energy_total / draft_energy / target_energy
    server_name="my-server",
)

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.toml 의 version 이 올라가 있다.
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.toml 의 name 값)
   • 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 할 때만 배포하세요.