pi05-base-p300x2 / tt-model.yaml
changh95's picture
Add files using upload-large-folder tool
f03ef2c verified
Raw History Blame Contribute Delete
11.3 kB
# SPDX-License-Identifier: Apache-2.0
# tt-model-manager container manifest (schema 5.1) for pi-0.5 (lerobot/pi05_base) on 2x Blackhole p300 (4 chips).
#
# The model is the multi-config pi0.5 megakernel of changh95/pi05-base-p150 (code/models/experimental/pi0: VISION,
# PREFIX and EXPERT as persistent generic_op programs on 110 cores of a chip, presets built from cameras 1-4 x prompt
# buckets 32/64/128/224 (+ extra buckets c1 L192, c2 L160, c4 L96) x action-row buckets 32/64) plus, in this package:
# - the single-robot mesh (tt/ttnn_pi05_sp.py): ONE model on all 4 chips of a 1x4 FABRIC_1D_RING mesh --
# camera-parallel VISION, sequence-parallel PREFIX with an in-kernel K / V exchange over the fabric, EXPERT on one
# chip -- with a static thermal-aware chip placement (tt/sp_thermal.py); actions bit-identical to one chip;
# - multi-robot serving with an SLO (tt/mr_sched.py, tt/ttnn_pi05_multichip.py dp): every chip runs the whole model,
# robots are bound to chips, the robot limit R_max comes from a start-up calibration against the p50 / p99 SLO and
# a new robot beyond it gets HTTP 503;
# - the batched EXPERT program (tt/megakernel/program_b.py, kernels/mkb_*), bit-identical per request;
# - prefix / expert disaggregation (chips 0-2 VISION + PREFIX, chip 3 the batched expert; K / V through the host);
# - faster / lower-energy programs (fp32 residual in L1, K / V-only last QKV, mask skip, VLM down K groups), all
# bit-identical to the previous release;
# - stability on QuietBox 2 (this revision): the PREFIX K / V feed runs its 4 full-grid multicasts one at a time with
# an acknowledged handoff (PE_KV_SERIAL + PE_KV_HANDOFF, bit-identical; ~11x fewer early-run chip hangs), and the
# server has a hang watchdog (PI05_HANG_S / PI05_HANG_ACTION / PI05_RESET_ON_START; see SERVING.md). The hang
# marker lives on the /cache volume so that the next container start resets the chips after a hang.
#
# Three serve profiles of one image (camera count / H / N fixed at server start; every request the same config):
# multi-robot (default) 4 chips, each the whole model; robots bound to chips, R_max from PI05_SLO_MS / _P99_MS
# single-robot all 4 chips as ONE model (SP mesh): the lowest latency for one robot
# disaggregated chips 0-2 VISION + PREFIX, chip 3 the batched expert (an experiment; below multi-robot)
#
# Run every tt-model command FROM THIS DIRECTORY (`root: code` resolves against the CWD):
# tt-model package --container tt-model.yaml --out <build dir>
# tt-model serve <build dir>/pi05-base-p300x2/tt_kernel_manifest.json [--profile single-robot]
# tt-model push <build dir>/pi05-base-p300x2 --publish
schema: "5.1"
repo: changh95/pi05-base-p300x2
name: pi05-base-p300x2
weights:
repo: lerobot/pi05_base
revision: b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba
allow_patterns:
- model.safetensors
- config.json
- policy_preprocessor.json
- policy_postprocessor.json
kind: tt-dit-server
arch: blackhole
source:
# tenstorrent/tt-metal main @ f856a38a361 (v0.80.0-dev20261001-17): the tree the p150 megakernel was validated on and
# this package was measured on (2x p300c). The port needs no tt-metal change: its kernels compile at run time.
tt_metal:
repo: https://github.com/tenstorrent/tt-metal
ref: f856a38a361939888f92d88f9e69b2f8a83fb713
code:
- models/common/lightweightmodule.py
extra_code:
- root: code
paths:
- models
ubuntu: "22.04"
python: "3.12"
runtime:
app: models.experimental.pi0_5.server.app:app
mesh_shape_env: TT_MESH_SHAPE
packages:
- "numpy>=1.24.4,<2"
- safetensors
- huggingface_hub
- transformers==5.12.1
- sentencepiece
- protobuf
serve:
port: 20000
hardware: p300x2
mesh_device: P300x2
env:
TT_WEIGHTS_REVISION: "b211f3d44c36b6acfcf7ae94a64e8e96f75a64ba"
PI05_MEGAKERNEL: "mc"
# Fixed at server start: cameras 1..4, action horizon 1..64, flow-matching steps 1..10.
PI05_NUM_IMAGES: "2"
PI05_ACTION_HORIZON: "50"
PI05_NUM_STEPS: "10"
PI05_TOKEN_LEN: "224"
PI05_TOKENIZER_REQUIRED: "1"
PI05_BATCH_WINDOW_MS: "0"
PI05_MODEL_NAME: "pi05-base-p300x2"
# Hang watchdog (server defaults, spelled out): a chip stuck 10 s is quarantined (multi-robot keeps serving on the
# others) or the server exits 75 (single-robot / disaggregated); a start after a hang resets the chips (tt-smi -r).
# The marker must outlive the container, so it goes on the persistent /cache volume.
PI05_HANG_S: "10"
PI05_HANG_ACTION: "degrade"
PI05_RESET_ON_START: "auto"
PI05_HANG_MARKER: "/cache/pi05_hang_marker"
default_profile: multi-robot
serve_profiles:
- name: multi-robot
description: "All 4 chips, each the whole model; each robot (request field robot_id) is bound to one chip. At start the server times every chip's cycle and derives the robots per chip and R_max from the SLO (p50 100 ms, p99 110 ms, clock guard 1.69 for chips clamped to 800 MHz when hot); a new robot beyond R_max gets HTTP 503. Hang watchdog: a stuck chip is quarantined, its robots move to the healthy chips (PI05_HANG_ACTION=exit: exit 75 for a supervisor restart). Measured in-process, 4 robots at 2 cameras / H 50 / N 10: p50 53 ms cool, p50 83.0 / p99 84.1 ms at thermal steady state, 52 req/s, 6.65 J per request."
env:
PI05_PROFILE: multi-robot
TT_MESH_SHAPE: "1x4"
PI05_LAYOUT: dp
PI05_SLO_MS: "100"
PI05_SLO_P99_MS: "110"
PI05_SLO_CLOCK_GUARD: "1.69"
- name: single-robot
description: "All 4 chips as ONE model (1x4 FABRIC_1D_RING mesh): camera-parallel VISION, sequence-parallel PREFIX, EXPERT on a cool chip (thermal placement); actions bit-identical to one chip. Host-to-host p50 at 2 cameras / 224 tokens / H 50: 38.5 ms cool (one chip: 52.6 ms). Per call at thermal steady state (trace replay, one chip at 800 MHz): 42.1 ms vs 38.0 ms cool; the legacy placement: 61.4 ms. On a hang the server exits 75 for a supervisor restart (the next start resets the chips)."
env:
PI05_PROFILE: single-robot
TT_MESH_SHAPE: "1x4"
PI05_LAYOUT: sp
- name: disaggregated
description: "Prefix / expert disaggregation: chips 0-2 run VISION + PREFIX, chip 3 the batched action expert on up to 4 requests (K / V caches copied through the host). Requests are accepted at any time; same config for every request. No SLO or admission control; per-request heterogeneous configs and LLM-style join/leave serving are future work. Measured below multi-robot on this model (see the card). On a hang the server exits 75 for a supervisor restart."
env:
PI05_PROFILE: disaggregated
TT_MESH_SHAPE: "1x4"
PI05_LAYOUT: pd
PI05_MAX_BATCH: "4"
verify:
- "import models.experimental.pi0.tt.ttnn_pi05_model as m; assert m.PI05MegakernelTTNN and m.PI05_DEVICE_PARAMS['worker_l1_size'] == 1395712, m.PI05_DEVICE_PARAMS"
- "from models.experimental.pi0.tt.megakernel import presets as PS; assert len(PS.PRESETS) == 38 and PS.CAMERAS == (1, 2, 3, 4) and PS.PROMPT_BUCKETS == (32, 64, 128, 224) and PS.SUFFIX_BUCKETS == (32, 64) and PS.EXTRA_BUCKETS == {1: (192,), 2: (160,), 4: (96,)}, (len(PS.PRESETS), PS.EXTRA_BUCKETS)"
- "from models.experimental.pi0.tt.megakernel import program_b as pb; assert pb.NB_MAX == 8 and pb.BatchedExpertMegakernel and pb.KVSlab"
- "import models.experimental.pi0.tt.ttnn_pi05_batched as b, models.experimental.pi0.tt.ttnn_pi05_multichip as mc; assert b.PI05BatchedTTNN and mc.PI05MultiChip"
- "import models.experimental.pi0.tt.ttnn_pi05_sp as sp, models.experimental.pi0.tt.sp_thermal as T, models.experimental.pi0.tt.mr_sched as R; assert sp.PI05SingleRobotMesh and T.ENV_PLACEMENT == 'PI05_SP_PLACEMENT' and T.ENV_HOT == 'PI05_SP_HOT' and T.ENV_MR_GUARD == 'PI05_MR_GUARD' and R.AdmissionRefused"
- "from models.experimental.pi0_5.server.mcx_backend import max_batch_for as f; assert (f(2, 50), f(4, 50), f(2, 10)) == (6, 4, 8), (f(2, 50), f(4, 50), f(2, 10))"
- "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp', PI05_MAX_BATCH='4'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.chips, c.max_batch) == ('mcx', 'dp', 4, 4), c"
- "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp', PI05_SLO_MS='100', PI05_SLO_P99_MS='110', PI05_SLO_CLOCK_GUARD='1.69'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.routing, c.slo_ms, c.slo_p99_ms, c.slo_clock_guard, c.max_batch) == ('mcx', 'dp', 'robot', 100.0, 110.0, 1.69, 6), c"
- "import os; [os.environ.pop(k, None) for k in ('PI05_NUM_IMAGES', 'PI05_ACTION_HORIZON', 'PI05_NUM_STEPS', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='sp'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert (c.backend, c.layout, c.chips, c.profile) == ('mc', 'sp', 4, 'single-robot'), c"
- "import os; [os.environ.pop(k, None) for k in ('PI05_LAYOUT', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x1'); from models.experimental.pi0_5.server import app as a; c = a.load_config(); assert c.backend == 'mc' and c.layout == 'mesh', c"
- "import re; from pathlib import Path; T = Path('/opt/tt-metal'); R = T / 'models/experimental/pi0'; inc = T / 'tt_metal/hw/inc'; ks = sorted(k for k in R.rglob('*') if k.suffix in ('.cpp', '.hpp', '.h')); assert len(ks) == 21, ks; pairs = [(k, i) for k in ks for i in re.findall(r'#include\\s+\"([^\"]+)\"', k.read_text())]; bad = [(str(k), i) for k, i in pairs if not ((k.parent / i).is_file() or (inc / i).is_file() or (T / i).is_file())]; assert not bad, bad"
- "import os; [os.environ.pop(k, None) for k in ('PI05_HANG_S', 'PI05_HANG_ACTION', 'PI05_RESET_ON_START', 'PI05_HANG_MARKER', 'PI05_LAYOUT', 'PI05_SLO_MS', 'PI05_SLO_P99_MS', 'PI05_MAX_BATCH', 'PI05_PROFILE')]; os.environ.update(PI05_MEGAKERNEL='mc', TT_MESH_SHAPE='1x4', PI05_LAYOUT='dp'); from models.experimental.pi0_5.server import app as a; from models.experimental.pi0.tt import mr_sched as R; c = a.load_config(); assert (c.hang_s, c.hang_action, c.hang_exit_code, c.reset_on_start) == (10.0, 'degrade', 75, 'auto') and R.ChipHung, c"
- "import ttnn; assert hasattr(ttnn, 'generic_op') and hasattr(ttnn, 'CreateDevices') and hasattr(ttnn, 'set_fabric_config') and hasattr(ttnn.FabricConfig, 'FABRIC_1D_RING') and hasattr(ttnn, 'open_mesh_device')"
- "import safetensors, huggingface_hub, sentencepiece, google.protobuf"
- "import transformers; assert transformers.__version__ == '5.12.1', transformers.__version__"
- "import numpy; assert numpy.__version__.startswith('1.'), numpy.__version__"
- "from pathlib import Path; assert Path('/opt/tt-metal/models/experimental/pi0_5/server/smoke_test.py').is_file()"