Unify Edge exporter run scripts

docsrc/tutorials/huggingface/index.rst

@@ -2,14 +2,15 @@ HuggingFace Models
 ==================
 
 Compile and accelerate HuggingFace models with Torch-TensorRT: large language models
-and visual language models via the ``tools/llm`` toolkit, Stable Diffusion via
-``torch.compile``, Flux via ``torch.export``, and LoRA weight-swapping via
-the Mutable Torch-TensorRT Module.
+and visual language models via the ``tools/llm`` toolkit, Edge-LLM VLAs via
+``EdgeExporter``, Stable Diffusion via ``torch.compile``, Flux via ``torch.export``,
+and LoRA weight-swapping via the Mutable Torch-TensorRT Module.
 
 .. toctree::
    :maxdepth: 1
 
    compile_hf_models
+   Edge Exporter <../../user_guide/edge_exporter>
    Example: Compiling Stable Diffusion with torch.compile <../_rendered_examples/dynamo/torch_compile_stable_diffusion>
    Example: Compiling FLUX.1-dev with the dynamo backend <../_rendered_examples/dynamo/torch_export_flux_dev>
    Example: Mutable Torch TensorRT Module <../_rendered_examples/dynamo/mutable_torchtrt_module_example>

docsrc/user_guide/index.rst

@@ -8,6 +8,7 @@ Conceptual guides and how-tos for Torch-TensorRT.
 
    torch_tensorrt_explained
    compilation/index
+   edge_exporter
    shapes_precision/index
    runtime_performance/index
    performance_tuning

tools/hf/__init__.py

@@ -0,0 +1,4 @@
+"""HuggingFace-facing export helpers for Torch-TensorRT.
+
+Use ``from exporters import EdgeExporter, EdgeConfig``.
+"""

tools/hf/exporters/__init__.py

@@ -0,0 +1,12 @@
+from .config import EdgeConfig
+from .exporter import EdgeExporter
+from .models.groot.spec import GrootSpec as _GrootSpec  # noqa: F401
+from .models.nanbeige.spec import NanbeigeSpec as _NanbeigeSpec  # noqa: F401
+from .models.nemotron.spec import NemotronSpec as _NemotronSpec  # noqa: F401
+from .models.pi05.spec import Pi05Spec as _Pi05Spec  # noqa: F401
+from .spec import (
+    ComponentBundle,
+    EdgeSpec,
+    get_edge_spec,
+    register_edge_spec,
+)

tools/hf/exporters/compile.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import json
+from inspect import Parameter, signature
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch_tensorrt
+
+from .measure import cuda_ms
+from .ops import _as_tuple, record_engine
+from .spec import ComponentBundle
+
+DEFAULT_TRT_SETTINGS: dict[str, Any] = {
+    "min_block_size": 1,
+    "require_full_compilation": True,
+    "immutable_weights": True,
+    "disable_tf32": True,
+    "truncate_double": True,
+}
+
+_TRT_COMPILE_KEYS = frozenset(DEFAULT_TRT_SETTINGS) | {
+    "use_fp32_acc",
+    "truncate_double",
+    "decompose_attention",
+    "offload_module_to_cpu",
+    "assume_dynamic_shape_support",
+    "use_explicit_typing",
+}
+
+
+def compile_component(
+    bundle: ComponentBundle,
+    *,
+    name: str,
+    engine_dir: Path,
+    trt_settings: dict[str, Any] | None = None,
+) -> tuple[str, tuple[torch.Tensor, ...], float]:
+    """Export one component, compile it, write ``engine_dir/<name>/``.
+
+    Family setattr is owned by ``EdgeSpec.apply_patches`` around this call.
+    ``execute_engine`` records the TensorRT module, not eager.
+    """
+    from .plugin.attn_patches import (
+        set_language_mask_type,
+    )
+
+    module = bundle.module.eval()
+    trace_args = tuple(bundle.trace_args)
+    save_args = tuple(bundle.save_args)
+    execute_args = tuple(bundle.execute_args or save_args)
+    out_dir = Path(engine_dir) / name
+    out_dir.mkdir(parents=True, exist_ok=True)
+    engine_path = str(out_dir)
+
+    if bundle.context_attention_mask_type is not None:
+        set_language_mask_type(bundle.context_attention_mask_type)
+
+    export_kwargs: dict[str, Any] = {"strict": False}
+    if bundle.input_specs is not None:
+        from torch_tensorrt.dynamo._tracer import build_dim_registry, get_dynamic_shapes
+
+        specs = tuple(bundle.input_specs)
+        leading = 0
+        for input_name in bundle.input_names:
+            if input_name.startswith("past_key_values"):
+                break
+            leading += 1
+        dim_registry = build_dim_registry(specs[:leading], {})
+        dynamic_shapes: dict[str, Any] = {}
+        positional_names: list[str] = []
+        var_pos_name: str | None = None
+        for param in signature(module.forward).parameters.values():
+            if param.kind == Parameter.VAR_POSITIONAL:
+                var_pos_name = param.name
+                break
+            if param.kind in (
+                Parameter.POSITIONAL_ONLY,
+                Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                positional_names.append(param.name)
+        for spec, param_name in zip(specs[:leading], positional_names[:leading]):
+            if param_name in ("inputs_embeds", "ds_stack"):
+                dynamic_shapes[param_name] = get_dynamic_shapes(spec, dim_registry)
+            else:
+                dynamic_shapes[param_name] = {}
+        if var_pos_name is not None:
+            dynamic_shapes[var_pos_name] = tuple(
+                get_dynamic_shapes(spec, dim_registry) for spec in specs[leading:]
+            )
+        export_kwargs["dynamic_shapes"] = dynamic_shapes
+
+    exported = torch.export.export(module, args=trace_args, **export_kwargs)
+    settings = {
+        k: v
+        for k, v in {
+            **DEFAULT_TRT_SETTINGS,
+            **(trt_settings or {}),
+            **bundle.trt_settings,
+        }.items()
+        if k in _TRT_COMPILE_KEYS
+    }
+
+    arg_inputs = (
+        tuple(bundle.input_specs) if bundle.input_specs is not None else trace_args
+    )
+    compiled = torch_tensorrt.dynamo.compile(
+        exported,
+        arg_inputs=arg_inputs,
+        **settings,
+    )
+
+    with torch.no_grad():
+        trt_out = _as_tuple(compiled(*execute_args))
+    trt_ms = cuda_ms(lambda: compiled(*execute_args))
+
+    record_engine(
+        engine_path,
+        component=name,
+        input_names=bundle.input_names,
+        outputs=trt_out,
+        module=compiled,
+    )
+    engine_file = bundle.engine_file
+
+    serialized = (
+        torch_tensorrt.dynamo.convert_exported_program_to_serialized_trt_engine(
+            exported,
+            arg_inputs=arg_inputs,
+            **settings,
+        )
+    )
+    (out_dir / engine_file).write_bytes(serialized)
+
+    _write_sidecar(out_dir, bundle, name, trt_out, engine_file=engine_file)
+    return engine_path, trt_out, trt_ms
+
+
+def _write_sidecar(
+    out_dir: Path,
+    bundle: ComponentBundle,
+    name: str,
+    outputs: tuple[torch.Tensor, ...],
+    *,
+    engine_file: str | None = None,
+) -> None:
+    config = {
+        "model_type": bundle.model_type,
+        "component": name,
+        "engine_file": engine_file or bundle.engine_file,
+        "input_names": list(bundle.input_names),
+        "output_names": list(bundle.output_names),
+        "outputs": [{"shape": list(t.shape), "dtype": str(t.dtype)} for t in outputs],
+    }
+    config.update(bundle.extra_config)
+    (out_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")

tools/hf/exporters/config.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class EdgeConfig:
+    """Knobs for :class:`~exporters.EdgeExporter`.
+
+    ``strict`` / ``dynamic`` / ``dynamic_shapes`` match HuggingFace
+    ``DynamoConfig`` so this can subclass it later without an API break.
+    """
+
+    strict: bool = False
+    dynamic: bool = False
+    dynamic_shapes: dict[str, Any] | None = None
+    prefer_deferred_runtime_asserts_over_guards: bool = False
+
+    engine_dir: Path | str | None = None
+    max_seq_len: int = 968
+    generation_reserve: int = 0
+    trt_settings: dict[str, Any] = field(default_factory=dict)
+    model_type: str | None = None

tools/hf/exporters/data.py

@@ -0,0 +1,131 @@
+"""LeRobot frame loading used by PI05 / GR00T ``prepare_sample_inputs``."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.constants import HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE, OBS_STATE
+from PIL import Image
+
+IMAGE_KEYS = ("observation.images.image", "observation.images.image2")
+DEFAULT_DATASET_ID = "lerobot/libero"
+DEFAULT_DATASET_REVISION = "v3.0"
+
+
+def _lerobot_dataset_has_metadata(root: Path) -> bool:
+    return (root / "meta" / "info.json").is_file()
+
+
+def _resolve_lerobot_dataset_root(
+    dataset_id: str,
+    revision: str = DEFAULT_DATASET_REVISION,
+) -> Path | None:
+    materialized = HF_LEROBOT_HOME / dataset_id
+    if _lerobot_dataset_has_metadata(materialized):
+        return materialized
+
+    hub_dir = HF_LEROBOT_HUB_CACHE / f"datasets--{dataset_id.replace('/', '--')}"
+    ref_file = hub_dir / "refs" / revision
+    if ref_file.is_file():
+        snapshot = hub_dir / "snapshots" / ref_file.read_text().strip()
+        if _lerobot_dataset_has_metadata(snapshot):
+            return snapshot
+
+    snapshots_dir = hub_dir / "snapshots"
+    if snapshots_dir.is_dir():
+        for snapshot in sorted(snapshots_dir.iterdir(), reverse=True):
+            if snapshot.is_dir() and _lerobot_dataset_has_metadata(snapshot):
+                return snapshot
+    return None
+
+
+def frame_from_test_data(
+    data: dict[str, Any],
+    policy: Any,
+    *,
+    fill_missing: bool = False,
+) -> dict[str, Any]:
+    frame = dict(data["images"])
+    frame[OBS_STATE] = data["state"]
+    frame["task"] = data.get("task", "")
+    if fill_missing:
+        for key, feature in policy.config.input_features.items():
+            if key.startswith("observation.images.") and key not in frame:
+                frame[key] = torch.zeros(feature.shape, dtype=torch.float32)
+    return frame
+
+
+def load_test_data(
+    dataset_id: str = DEFAULT_DATASET_ID,
+    *,
+    episode_index: int = 0,
+    frame_index: int = 0,
+) -> dict[str, Any]:
+    local_root = _resolve_lerobot_dataset_root(dataset_id)
+    dataset_kwargs: dict[str, Any] = {
+        "episodes": [episode_index],
+        "video_backend": "pyav",
+        "revision": DEFAULT_DATASET_REVISION,
+    }
+    if local_root is not None:
+        dataset_kwargs["root"] = local_root
+    dataset = LeRobotDataset(dataset_id, **dataset_kwargs)
+    frame = dataset[frame_index]
+    images = {key: frame[key] for key in IMAGE_KEYS if key in frame}
+    return {
+        "images": images,
+        "state": frame[OBS_STATE],
+        "task": frame.get("task", "") or "Perform the task.",
+    }
+
+
+def create_pil_messages(data: dict[str, Any]) -> list[dict[str, Any]]:
+    images = data["images"]
+    task = str(data.get("task", "") or "Perform the task.")
+    image_content = [
+        {"type": "image", "image": _tensor_image_to_pil(img)}
+        for _, img in sorted(images.items())
+    ]
+    return [
+        {
+            "role": "user",
+            "content": image_content + [{"type": "text", "text": str([task])}],
+        }
+    ]
+
+
+def pack_state(
+    state: torch.Tensor,
+    max_state_dim: int,
+    device: str | torch.device,
+) -> torch.Tensor:
+    state = torch.as_tensor(state, dtype=torch.float32, device=device)
+    if state.ndim == 1:
+        state = state.unsqueeze(0)
+    if state.ndim == 2:
+        state = state.unsqueeze(1)
+    bsz, _, state_dim = state.shape
+    if state_dim > max_state_dim:
+        state = state[:, :, :max_state_dim]
+    elif state_dim < max_state_dim:
+        pad = torch.zeros(
+            bsz,
+            1,
+            max_state_dim - state_dim,
+            dtype=state.dtype,
+            device=device,
+        )
+        state = torch.cat([state, pad], dim=-1)
+    return state
+
+
+def _tensor_image_to_pil(img: torch.Tensor) -> Image.Image:
+    img = img.detach().cpu()
+    if img.dtype.is_floating_point:
+        img = (img.clamp(0, 1) * 255).to(torch.uint8)
+    if img.ndim == 3 and img.shape[0] in (1, 3):
+        img = img.permute(1, 2, 0)
+    return Image.fromarray(img.numpy())