Load the pi0.5 policy at the size it runs at

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.

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

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(

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",