From be03783127fb968da4fba439c082afb7b73d81d9 Mon Sep 17 00:00:00 2001 From: Dan Wahl Date: Thu, 20 Aug 2026 23:35:30 -0500 Subject: [PATCH] Load the pi0.5 policy at the size it runs at The two vocabulary heads no action chunk reads are dropped, and a checkpoint is built in host memory so what reaches the GPU is what runs there. On the shipped Kinova checkpoint the server holds 8593 MiB where it held 10103. An int8 knob in vla_serving.yaml holds the language backbone and the vision tower at eight bits, taking that to 5375 MiB for 0.14s per 50-step chunk against 0.13s, both inside the real-time budget at 10 fps. /health reports which of the two is running. What eight bits costs in success rate is unmeasured on this checkpoint. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01JFNEZkKa2NN7LMWbzM37Nq --- src/vla_sim/config/vla_serving.yaml | 7 ++ .../docker/Dockerfile.vla_inference_server | 8 ++- .../docker/test_vla_inference_server.py | 37 ++++++++++ src/vla_sim/docker/vla_inference_server.py | 70 ++++++++++++++++++- 4 files changed, 118 insertions(+), 4 deletions(-) diff --git a/src/vla_sim/config/vla_serving.yaml b/src/vla_sim/config/vla_serving.yaml index e4b750ab7..c4082ddf0 100644 --- a/src/vla_sim/config/vla_serving.yaml +++ b/src/vla_sim/config/vla_serving.yaml @@ -26,6 +26,13 @@ fps: 10.0 # order of magnitude slower on cpu. device: auto +# Hold the pi0.5 language backbone and vision tower at eight bits, most of the +# weights: roughly 3 GiB less GPU memory for a little more time per chunk. false +# loads every weight at full width, and a policy family other than pi0.5 ignores +# this without saying so. What it costs in success rate is unmeasured on this +# checkpoint. +int8: false + # Trained observation.state width: 7 arm joints plus 1 gripper. 0 trusts # config.json, which for this checkpoint reports its padded 32-dim # architecture width instead of the real one. diff --git a/src/vla_sim/docker/Dockerfile.vla_inference_server b/src/vla_sim/docker/Dockerfile.vla_inference_server index f1ef72769..f04b4228c 100644 --- a/src/vla_sim/docker/Dockerfile.vla_inference_server +++ b/src/vla_sim/docker/Dockerfile.vla_inference_server @@ -15,16 +15,22 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # torch/torchvision are pinned (within lerobot 0.6.0's >=2.7,<2.12 range) so # builds are reproducible and the CPU-only pre-install cannot be re-resolved # to a CUDA build by the lerobot install. +# torchao supplies the int8 weights vla_serving.yaml can ask for. Its wheel +# declares no torch dependency, so pip cannot catch a mismatch; torchao skips +# loading its compiled kernels with a warning when torch is older than the +# release expects, which leaves the pairing below this file's to keep. ARG TORCH_INDEX= ARG TORCH_VERSION=2.11.0 ARG TORCHVISION_VERSION=0.26.0 +ARG TORCHAO_VERSION=0.17.0 RUN if [ -n "$TORCH_INDEX" ]; then \ pip install --no-cache-dir --index-url "$TORCH_INDEX" \ --extra-index-url https://pypi.org/simple \ "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION"; \ fi \ && pip install --no-cache-dir "lerobot[pi,smolvla]==0.6.0" \ - "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION" + "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION" \ + "torchao==$TORCHAO_VERSION" COPY vla_inference_server.py /app/vla_inference_server.py WORKDIR /app diff --git a/src/vla_sim/docker/test_vla_inference_server.py b/src/vla_sim/docker/test_vla_inference_server.py index 0a46d5948..a8a657b65 100644 --- a/src/vla_sim/docker/test_vla_inference_server.py +++ b/src/vla_sim/docker/test_vla_inference_server.py @@ -210,6 +210,42 @@ def test_non_numeric_yaml_values_park_in_config_error(self) -> None: self.assertEqual(args.fps, 0.0) self.assertEqual(args.state_dim, 0) + def test_non_boolean_int8_parks_in_config_error(self) -> None: + """A quoted or misspelled int8 lands in config_error with the flag off. + Coercing it with bool() would read every non-empty string as true, so the + server would quantize the policy the operator asked to leave alone.""" + # GIVEN a serving config whose int8 is a string rather than a boolean + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f: + f.write('int8: "false"\n') + path = f.name + try: + # WHEN parsing arguments against that config + with patch("sys.argv", ["vla_inference_server.py", "--config", path]): + args = parse_args() + finally: + os.unlink(path) + + # THEN the value is reported and the flag stays off + self.assertIn("int8", args.config_error) + self.assertFalse(args.int8) + + def test_boolean_int8_is_honored(self) -> None: + """A real YAML boolean reaches the flag, so the knob works as documented.""" + # GIVEN a serving config asking for int8 + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f: + f.write("int8: true\n") + path = f.name + try: + # WHEN parsing arguments against that config + with patch("sys.argv", ["vla_inference_server.py", "--config", path]): + args = parse_args() + finally: + os.unlink(path) + + # THEN the flag is on and nothing is reported + self.assertTrue(args.int8) + self.assertEqual(args.config_error, "") + class TestLoadPolicyMissingCheckpoint(unittest.TestCase): """load_policy: an unset checkpoint parks the error state, never exits.""" @@ -424,6 +460,7 @@ def __init__( native_map: dict | None = None, ) -> None: self.device = "cpu" + self.int8 = False self._infer_error = infer_error # Like PolicyRunner, derived once at construction. self.request_names = request_camera_names( diff --git a/src/vla_sim/docker/vla_inference_server.py b/src/vla_sim/docker/vla_inference_server.py index b3ca3c292..677e1a861 100644 --- a/src/vla_sim/docker/vla_inference_server.py +++ b/src/vla_sim/docker/vla_inference_server.py @@ -77,10 +77,13 @@ from huggingface_hub import hf_hub_download from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError +from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import RTCAttentionSchedule from lerobot.policies.factory import get_policy_class, make_pre_post_processors from lerobot.policies.rtc.configuration_rtc import RTCConfig +from torchao.quantization import Int8WeightOnlyConfig, quantize_ + # pi0.5 checkpoints save a processor pipeline that references # 'relative_actions_processor', an alias lerobot does not always auto-register; # without it make_pre_post_processors raises @@ -111,6 +114,14 @@ # src/vla_sim/config/. Overridable with --config for a standalone `docker run`. DEFAULT_CONFIG_PATH = "/vla_config/vla_serving.yaml" +# The two modules int8 quantizes, named by prefix. Both run once per chunk and +# hand the action expert a cache, so they sit the far side of the model from the +# commands. +INT8_MODULE_PREFIXES = ( + "paligemma.model.language_model", + "paligemma.model.vision_tower", +) + def load_serving_config(path: str) -> dict: """Read the per-config model-serving YAML into a dict of knob values. @@ -319,7 +330,8 @@ class PolicyRunner: Loading passes policy_cfg by keyword and overrides the device on both processors, which merged pi0.5 checkpoints need: their config declares a padded 32-dim state while the saved normalizer stats carry the trained - width. + width. It loads on the cpu, trims pi0.5's unused vocabulary heads and + quantizes when asked, before moving to the serving device. """ def __init__( @@ -330,6 +342,7 @@ def __init__( guidance_horizon: int, rtc_schedule: str, state_dim: int, + int8: bool = False, ): self.device = device self.state_dim = state_dim @@ -339,9 +352,40 @@ def __init__( # Resolve before the slow checkpoint load so a schedule typo fails fast. schedule = resolve_rtc_schedule(rtc_schedule) - self.policy = get_policy_class(policy_type).from_pretrained(checkpoint) + # Loaded on the cpu and moved once the trimming below has run: building on + # the gpu first would peak at the untrimmed size, which is the size this is + # here to stay under. + policy_config = PreTrainedConfig.from_pretrained(checkpoint) + policy_config.device = "cpu" + self.policy = get_policy_class(policy_type).from_pretrained( + checkpoint, config=policy_config + ) + if policy_type == "pi05": + model = self.policy.model + # The two vocabulary heads are ~1.5 GiB of weights no chunk ever reads, + # since actions leave through action_out_proj. Dropping them is why this + # load is smaller than an unmodified one even with int8 off. + model.paligemma_with_expert.paligemma.lm_head = None + model.paligemma_with_expert.gemma_expert.lm_head = None + + if int8: + # version=2 gives each output channel its own scale rather than one + # for the whole tensor. Inductor's config is declined because it turns + # on TF32 for every float32 matmul in the process. + quantize_( + model.paligemma_with_expert, + Int8WeightOnlyConfig(version=2, set_inductor_config=False), + filter_fn=lambda module, name: ( + isinstance(module, torch.nn.Linear) + and name.startswith(INT8_MODULE_PREFIXES) + ), + ) self.policy.to(device) + # The config named the load device, and something downstream reading it + # would otherwise be told the weights are still on the host. + policy_config.device = device self.policy.eval() + self.int8 = int8 and policy_type == "pi05" # infer() passes the horizon per call on every RTC request, so the # config's own execution_horizon never applies; only the enable and @@ -525,7 +569,8 @@ def load_policy(state: ServerState, args: argparse.Namespace) -> None: ) log( f"loading {policy_type} checkpoint '{args.checkpoint}' on '{device}' " - f"(torch {torch.__version__}) ..." + f"(torch {torch.__version__}" + f"{', int8' if args.int8 and policy_type == 'pi05' else ''}) ..." ) runner = PolicyRunner( args.checkpoint, @@ -534,6 +579,7 @@ def load_policy(state: ServerState, args: argparse.Namespace) -> None: args.guidance_horizon, args.rtc_schedule, args.state_dim, + args.int8, ) image_features = [ @@ -719,6 +765,7 @@ def do_GET(self): health["detail"] = state.detail elif state.status == "ready": health["device"] = state.runner.device + health["int8"] = state.runner.int8 self._send(200, health) def _authorized(self) -> bool: @@ -826,6 +873,16 @@ def numeric_default(name: str, cast, builtin): config_errors.append(f"{name}: '{value}' is not a number") return builtin + def flag_default(name: str, builtin: bool): + # Parks the same way a bad number does. bool() would take any non-empty + # string as True, so a quoted `int8: "false"` would serve the opposite of + # what the operator wrote. + value = resolve_default(config.get(name), builtin) + if isinstance(value, bool): + return value + config_errors.append(f"{name}: '{value}' is not true or false") + return builtin + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--config", @@ -856,6 +913,13 @@ def numeric_default(name: str, cast, builtin): default=str(resolve_default(config.get("device"), "auto")), help="torch device: auto | cpu | cuda", ) + parser.add_argument( + "--int8", + action=argparse.BooleanOptionalAction, + default=flag_default("int8", False), + help="hold the pi0.5 backbone and vision tower at eight bits; " + "other policy families are unaffected", + ) parser.add_argument("--port", type=int, default=8973) parser.add_argument( "--state-dim",