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