From c99e14061030ec5b46f4bf3c0f6517151f28faa7 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 8 Sep 2026 00:06:46 -0700 Subject: [PATCH] Always compile Edge engines and dual-profile PI05 language I/O. --- docsrc/user_guide/edge_exporter.rst | 91 ++-- tools/hf/exporters/compile.py | 161 +++--- tools/hf/exporters/config.py | 5 - tools/hf/exporters/exporter.py | 110 +--- tools/hf/exporters/models/common/helpers.py | 8 +- tools/hf/exporters/models/common/patches.py | 50 +- tools/hf/exporters/models/groot/spec.py | 340 ++++++++----- tools/hf/exporters/models/nemotron/patches.py | 11 +- tools/hf/exporters/models/nemotron/spec.py | 59 ++- tools/hf/exporters/models/pi05/helpers.py | 10 - tools/hf/exporters/models/pi05/spec.py | 476 +++++++++++++----- tools/hf/exporters/ops.py | 11 +- tools/hf/exporters/plugin/attn_patches.py | 35 +- tools/hf/exporters/spec.py | 69 ++- .../hf/exporters/tests/test_edge_exporter.py | 174 ++----- tools/hf/run_groot_export.py | 20 +- tools/hf/run_nemotron_export.py | 15 +- tools/hf/run_pi05_export.py | 23 +- 18 files changed, 926 insertions(+), 742 deletions(-) diff --git a/docsrc/user_guide/edge_exporter.rst b/docsrc/user_guide/edge_exporter.rst index 7263ebeb39..74009efe52 100644 --- a/docsrc/user_guide/edge_exporter.rst +++ b/docsrc/user_guide/edge_exporter.rst @@ -11,7 +11,7 @@ one ``torch.export`` graph that calls those engines. from hf.exporters import EdgeExporter, EdgeConfig exporter = EdgeExporter() - config = EdgeConfig(dryrun=True, engine_dir="/tmp/pi05_edge") + config = EdgeConfig(engine_dir="/tmp/pi05_edge") exported = exporter.export(policy, {"device": device, "dtype": torch.float16}, config) ``EdgeExporter`` is a HuggingFace ``DynamoExporter``. The public call is the same @@ -95,7 +95,6 @@ plus sample inputs, and call ``export()``. model_type="pi05", # optional when the spec can infer it engine_dir="/tmp/pi05_edge", max_seq_len=968, - dryrun=True, # skip TensorRT; still writes config.json + the outer graph ) exported = exporter.export(policy, {"device": device, "dtype": torch.float16}, config) @@ -108,13 +107,9 @@ plus sample inputs, and call ``export()``. Pass the **policy** for PI05 and GR00T (the spec needs the preprocessor), not an inner submodule. Pass the HuggingFace causal LM for Nemotron. -``dryrun=True`` walks the same export path without building TensorRT engines. Each -component directory still gets a ``config.json``. Use that to debug packing and -patches, then set ``dryrun=False`` (or pass ``--compile`` on the example scripts) -to emit ``.engine`` files. - -On a real compile, ``engine_dir//`` contains ``config.json`` and the -serialized engine (for example ``visual.engine``, ``language.engine``). +``engine_dir//`` contains ``config.json`` and the serialized engine +(for example ``visual.engine``, ``language.engine``). The smoke scripts always +compile. Export Program -------------- @@ -128,7 +123,7 @@ text embeddings. ``print(program.graph)`` prints that FX graph: each ``execute_engine`` node is one component, and the path in its args is the engine directory. -Here is a GR00T dryrun (``vision`` → ``scatter_image_tokens`` → +Here is a GR00T outer graph (``vision`` → ``scatter_image_tokens`` → ``language`` → ``context_projection`` → ``action``): .. code-block:: text @@ -182,14 +177,15 @@ Patches ------- Edge does not wrap the policy in a new module. It temporarily replaces -``Class.forward`` on the original HuggingFace / LeRobot class, compiles that -submodule, then restores the method (dryrun leaves the replacement in place). +``Class.forward`` on the original HuggingFace / LeRobot class so +``torch.export`` / TensorRT see plugin I/O, then restores the method. +Eager inference is the unpatched model. ``@register_patch`` does not install anything. It records a factory and a dotted class path on a backend (``"pi05"``, ``"groot"``, ``"nemotron"``). ``apply_patches(backend)`` imports that class and does -``setattr(Cls, "forward", factory(original))`` for the duration of -``export()``. +``setattr(Cls, "forward", factory(original))`` while each component is +traced and compiled. HuggingFace ``DynamoExporter`` uses the same two steps. The purpose is different. HF patches make the original modeling ``forward`` traceable. Edge @@ -288,15 +284,15 @@ Add a new model --------------- ``EdgeExporter.export`` never branches on PI05 vs GR00T. It loads an ``EdgeSpec`` -and loops ``spec.components``. A new architecture is a new spec plus a patch -backend. +and compiles whatever ``prepare`` returns. A new architecture is a new spec +plus a patch backend. Create ``tools/hf/exporters/models//``: .. code-block:: text / - spec.py # EdgeSpec: components, sample inputs, prepare, run + spec.py # EdgeSpec: sample inputs, prepare, run patches.py # @register_patch factories on this family's backend helpers.py # optional packing / submodule lookup @@ -316,8 +312,6 @@ exporter package loads: @register_edge_spec("my_vla") class MyVlaSpec(EdgeSpec): - components = ("vision", "language", "action") - def apply_patches(self, model=None): from hf.exporters.plugin.attn_patches import apply_patches from .patches import MY_VLA @@ -340,17 +334,43 @@ Reuse the shared plugin attention factories when the layout matches 3. Select submodules and flatten I/O ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -``prepare(name, model, sample, upstream, config)`` returns a ``ComponentBundle``: +``prepare(model, sample, config)`` returns a dict of ``ComponentBundle`` values +(one per engine). The exporter compiles that dict; it does not branch on +component names. * ``module`` — the original submodule to compile (vision tower, decoder, action head) * ``trace_args`` / ``save_args`` — positional tensors for ``torch.export`` / the engine * ``input_names`` / ``output_names`` — written into ``config.json`` * ``context_attention_mask_type`` — padding vs causal for the language plugin -``capture_upstream`` maps this engine's outputs into keys the next ``prepare`` -needs (image tokens, prefix KV, context embeddings). +Pack later-stage example tensors inside ``prepare`` (unpatched eager, or zeros +of the trace shape). ``run()`` still chains **engine** outputs at runtime. + +4. Capture unpatched eager +^^^^^^^^^^^^^^^^^^^^^^^^^^ -4. Call engines in ``run()`` +``capture_eager_outputs(model, sample, config)`` runs the original HuggingFace / +LeRobot forwards **before** ``apply_patches``. Return one tensor per component +(the value e2e passes to ``parity``). The exporter compares that to TensorRT. + +.. code-block:: python + + def capture_eager_outputs(self, model, sample, config, bench=None): + paligemma = ... + language = paligemma.language_model + with torch.no_grad(): + visual_embeds = paligemma.multi_modal_projector( + paligemma.vision_tower(sample["pixel_values"]).last_hidden_state + ) + lm = language( + inputs_embeds=sample["prefix_embs"], + attention_mask=sample["prefix_attention_mask"], + position_ids=sample["prefix_position_ids"], + return_dict=True, + ) + return {"vision": visual_embeds, "language": lm.last_hidden_state, ...} + +5. Call engines in ``run()`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ``run(engines, sample)`` is the outer graph. Call ``call_engine`` for each @@ -367,7 +387,7 @@ or your own custom op). engines["action"], "action", sample["step_actions"], ..., lm[2], lm[3] ) -5. Collate sample inputs +6. Collate sample inputs ^^^^^^^^^^^^^^^^^^^^^^^^ ``prepare_sample_inputs`` turns the caller payload into the stem dict ``prepare`` @@ -378,21 +398,16 @@ so the example scripts can pass only ``device`` and ``dtype``. Example scripts --------------- -The smoke scripts live next to the other Dynamo examples. Default is **dryrun** -(no TensorRT). Pass ``--compile`` to build engines. +The smoke scripts live in ``tools/hf``. They always build TensorRT engines. .. code-block:: bash - cd TensorRT/examples/dynamo + cd TensorRT/tools/hf python run_pi05_export.py - python run_pi05_export.py --compile --engine-dir /tmp/pi05_edge - python run_groot_export.py - python run_groot_export.py --compile --engine-dir /tmp/groot_edge - python run_nemotron_export.py --prompt "Hello." - python run_nemotron_export.py --compile --checkpoint nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + python run_nemotron_export.py --checkpoint nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 .. list-table:: :header-rows: 1 @@ -414,9 +429,6 @@ The smoke scripts live next to the other Dynamo examples. Default is **dryrun** Each script loads plugins, forces ``eager`` attention, calls ``EdgeExporter.export``, prints ``exporter.engines``, and runs the returned program once. -To export only one component while you debug a family, set -``EdgeConfig(components=("vision",))`` (or pass a subset of ``spec.components``). - Configuration ------------- @@ -432,18 +444,9 @@ Configuration * - ``engine_dir`` - ``"edge_engines"`` - Output directory; one subdirectory per component - * - ``dryrun`` - - ``False`` - - Skip TensorRT; keep patched Python modules for ``execute_engine`` - * - ``skip_runtime_export`` - - ``False`` - - Return the runtime module without ``torch.export`` of the outer graph * - ``model_type`` - inferred - ``"pi05"``, ``"groot"``, ``"nemotron_h"`` - * - ``components`` - - spec default - - Subset of engines to compile * - ``max_seq_len`` - ``968`` - KV / RoPE capacity for language diff --git a/tools/hf/exporters/compile.py b/tools/hf/exporters/compile.py index 81777e5c66..8030f3319a 100644 --- a/tools/hf/exporters/compile.py +++ b/tools/hf/exporters/compile.py @@ -1,13 +1,14 @@ 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, parity +from .measure import cuda_ms from .ops import _as_tuple, record_engine from .spec import ComponentBundle @@ -34,17 +35,12 @@ def compile_component( *, name: str, engine_dir: Path, - dryrun: bool = False, trt_settings: dict[str, Any] | None = None, - bench: dict[str, tuple[float, float]] | None = None, -) -> tuple[str, tuple[torch.Tensor, ...]]: +) -> tuple[str, tuple[torch.Tensor, ...], float]: """Export one component, compile it, write ``engine_dir//``. - Returns ``(engine_dir, example_outputs)`` from a patched eager run so the - exporter can chain components without a second unpatched forward. - - Family setattr is owned by ``EdgeSpec.apply_patches``, not this helper. - ``dryrun`` records the patched eager module for ``execute_engine``. + 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, @@ -61,81 +57,84 @@ def compile_component( if bundle.context_attention_mask_type is not None: set_language_mask_type(bundle.context_attention_mask_type) - patched = bundle.patch_fn(module) if bundle.patch_fn is not None else None - try: - with torch.no_grad(): - example = module(*execute_args) - outputs = _as_tuple(example) - record_engine( - engine_path, - component=name, - input_names=bundle.input_names, - outputs=outputs, - module=module, - ) - if dryrun: - _write_sidecar(out_dir, bundle, name, outputs, dryrun=True) - return engine_path, outputs - - exported = torch.export.export(module, args=trace_args, strict=False) - settings = { - k: v - for k, v in { - **DEFAULT_TRT_SETTINGS, - **(trt_settings or {}), - **bundle.trt_settings, - }.items() - if k in _TRT_COMPILE_KEYS - } - - compiled = torch_tensorrt.dynamo.compile( - exported, - arg_inputs=trace_args, - **settings, - ) - - with torch.no_grad(): - trt_out = _as_tuple(compiled(*execute_args)) - for i, (eager_t, trt_t) in enumerate(zip(outputs, trt_out)): - if not isinstance(eager_t, torch.Tensor) or not isinstance( - trt_t, torch.Tensor + 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, ): - continue - label = name if i == 0 else f"{name}[{i}]" - parity(f"{label} A vs C (TRT)", eager_t, trt_t) - - eager_ms = cuda_ms(lambda: module(*execute_args)) - trt_ms = cuda_ms(lambda: compiled(*execute_args)) - if bench is not None: - bench[name] = (eager_ms, trt_ms) - - record_engine( - engine_path, - component=name, - input_names=bundle.input_names, - outputs=outputs, - module=compiled, - ) - engine_file = bundle.engine_file - - serialized = ( - torch_tensorrt.dynamo.convert_exported_program_to_serialized_trt_engine( - exported, - arg_inputs=trace_args, - **settings, + 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, outputs, engine_file=engine_file) - return engine_path, outputs - finally: - if not dryrun and patched is not None: - from .plugin.plugin_utils import ( - restore_attention, - ) + ) + (out_dir / engine_file).write_bytes(serialized) - restore_attention(patched) # type: ignore[no-untyped-call] + _write_sidecar(out_dir, bundle, name, trt_out, engine_file=engine_file) + return engine_path, trt_out, trt_ms def _write_sidecar( @@ -145,7 +144,6 @@ def _write_sidecar( outputs: tuple[torch.Tensor, ...], *, engine_file: str | None = None, - dryrun: bool = False, ) -> None: config = { "model_type": bundle.model_type, @@ -153,7 +151,6 @@ def _write_sidecar( "engine_file": engine_file or bundle.engine_file, "input_names": list(bundle.input_names), "output_names": list(bundle.output_names), - "dryrun": dryrun, "outputs": [{"shape": list(t.shape), "dtype": str(t.dtype)} for t in outputs], } config.update(bundle.extra_config) diff --git a/tools/hf/exporters/config.py b/tools/hf/exporters/config.py index d81055589b..3a822da16b 100644 --- a/tools/hf/exporters/config.py +++ b/tools/hf/exporters/config.py @@ -11,8 +11,6 @@ class EdgeConfig: ``strict`` / ``dynamic`` / ``dynamic_shapes`` match HuggingFace ``DynamoConfig`` so this can subclass it later without an API break. - ``components`` is ``None`` to use the spec default (1 engine for an LLM, - 3–4 for a VLA). """ strict: bool = False @@ -23,8 +21,5 @@ class EdgeConfig: engine_dir: Path | str | None = None max_seq_len: int = 968 generation_reserve: int = 0 - components: tuple[str, ...] | None = None trt_settings: dict[str, Any] = field(default_factory=dict) - dryrun: bool = False - skip_runtime_export: bool = False model_type: str | None = None diff --git a/tools/hf/exporters/exporter.py b/tools/hf/exporters/exporter.py index f93ca57bc0..a46b5ad964 100644 --- a/tools/hf/exporters/exporter.py +++ b/tools/hf/exporters/exporter.py @@ -1,9 +1,6 @@ from __future__ import annotations -import contextlib import copy -import inspect -import logging from collections.abc import MutableMapping from pathlib import Path from typing import Any @@ -11,16 +8,19 @@ import torch import torch.nn as nn from torch.export import ExportedProgram -from transformers.exporters.exporter_dynamo import DynamoExporter +from transformers.exporters.exporter_dynamo import ( + DynamoExporter, + get_auto_dynamic_shapes, + patch_forward_signature, +) from . import ops as _ops # noqa: F401 from .compile import compile_component from .config import EdgeConfig +from .measure import parity from .runtime import EdgeRuntimeModule from .spec import get_edge_spec -logger = logging.getLogger(__name__) - def _clone_export_kwargs(sample_inputs: MutableMapping[str, Any]) -> dict[str, Any]: """Copy example kwargs into graph leaves. @@ -41,122 +41,60 @@ class EdgeExporter(DynamoExporter): # type: ignore[misc] def __init__(self) -> None: super().__init__() self.engines: dict[str, str] = {} - self.runtime: EdgeRuntimeModule | None = None self.sample: dict[str, Any] = {} self.bench: dict[str, tuple[float, float]] = {} - self._dryrun_patches: contextlib.ExitStack | None = None def export( self, model: nn.Module, sample_inputs: MutableMapping[str, Any], config: EdgeConfig | dict[str, Any], - ) -> ExportedProgram | EdgeRuntimeModule: + ) -> ExportedProgram: if isinstance(config, dict): config = EdgeConfig(**config) elif not isinstance(config, EdgeConfig): raise TypeError(f"Expected EdgeConfig or dict, got {type(config)}") - # Family spec owns flatten / stitch. The exporter only loops - # over component names (vision, language, action, ...). spec = get_edge_spec(model, config.model_type) - names = config.components or spec.components - if not names: - raise ValueError(f"{type(spec).__name__} has empty components") - - # Caller payload (policy batch, tokenizer ids, ...) -> shared sample - # dict used by prepare and the stitched runtime. sample = spec.prepare_sample_inputs(model, sample_inputs, config) - + bundles = spec.prepare(model, sample, config) engine_dir = Path(config.engine_dir or "edge_engines") engine_dir.mkdir(parents=True, exist_ok=True) + eager_ms: dict[str, float] = {} + eager = spec.capture_eager_outputs(model, sample, config, bench=eager_ms) + engines: dict[str, str] = {} - upstream: dict[str, Any] = {} self.bench = {} - - def _compile_components() -> None: - for name in names: - bundle = spec.prepare(name, model, sample, upstream, config) - engines[name], outs = compile_component( + with spec.apply_patches(model): + for name, bundle in bundles.items(): + engines[name], trt_out, trt_ms = compile_component( bundle, name=name, engine_dir=engine_dir, - dryrun=config.dryrun, trt_settings=config.trt_settings, - bench=self.bench, ) - upstream.update(spec.capture_upstream(name, outs, sample, bundle)) + self.bench[name] = (eager_ms.get(name, 0.0), trt_ms) + out_name = bundle.parity_output or bundle.output_names[0] + trt = trt_out[bundle.output_names.index(out_name)] + ref = eager.get(name) + if isinstance(ref, torch.Tensor) and isinstance(trt, torch.Tensor): + parity(f"{name} eager vs TRT", ref, trt) - # Family setattr. Dryrun leaves them installed so execute_engine still - # hits the patched original module after export() returns. - if config.dryrun: - if self._dryrun_patches is not None: - self._dryrun_patches.close() - self._dryrun_patches = contextlib.ExitStack() - self._dryrun_patches.enter_context(spec.apply_patches(model)) - _compile_components() - else: - with spec.apply_patches(model): - _compile_components() - - # One module whose forward is spec.run() over execute_engine calls. runtime = EdgeRuntimeModule(spec, engines) runtime_kwargs = _clone_export_kwargs(spec.runtime_kwargs(sample)) - self.engines = engines - self.runtime = runtime self.sample = dict(runtime_kwargs) - if config.skip_runtime_export: - return runtime - # torch.export the stitched graph so the product is one ExportedProgram. - return self._export_runtime(runtime, runtime_kwargs, config) - - def _export_runtime( - self, - model: nn.Module, - sample_inputs: MutableMapping[str, Any], - config: EdgeConfig, - ) -> ExportedProgram: - try: - from transformers.exporters.exporter_dynamo import ( - get_auto_dynamic_shapes, - patch_forward_signature, - register_cache_pytrees_for_model, - reset_model_state, - ) - from transformers.exporters.utils import prepare_for_export - except ImportError: - return torch.export.export( - model, - args=(), - kwargs=_clone_export_kwargs(sample_inputs), - strict=config.strict, - dynamic_shapes=config.dynamic_shapes, - ) - - sample_inputs = _clone_export_kwargs(sample_inputs) - model, sample_inputs, _output_flags = prepare_for_export(model, sample_inputs) dynamic_shapes = config.dynamic_shapes - if config.dynamic and dynamic_shapes is None: - dynamic_shapes = get_auto_dynamic_shapes(sample_inputs) - - if inspect.getmodule(model) is not None: - try: - register_cache_pytrees_for_model(model) - except Exception: - logger.debug("register_cache_pytrees_for_model skipped", exc_info=True) + dynamic_shapes = get_auto_dynamic_shapes(runtime_kwargs) - with ( - reset_model_state(model), - patch_forward_signature(model, sample_inputs), - ): + with patch_forward_signature(runtime, runtime_kwargs): return torch.export.export( - model, + runtime, args=(), - kwargs=_clone_export_kwargs(sample_inputs), + kwargs=_clone_export_kwargs(runtime_kwargs), strict=config.strict, dynamic_shapes=dynamic_shapes, prefer_deferred_runtime_asserts_over_guards=( diff --git a/tools/hf/exporters/models/common/helpers.py b/tools/hf/exporters/models/common/helpers.py index 5f4951e07c..2eb153e7bf 100644 --- a/tools/hf/exporters/models/common/helpers.py +++ b/tools/hf/exporters/models/common/helpers.py @@ -16,7 +16,11 @@ def causal_lm_flat( dtype: torch.dtype, seq_len: int | None = None, ) -> tuple[tuple[torch.Tensor, ...], dict[str, Any]]: - """inputs_embeds, rope, ctx, kv_start, last_token_ids, ds_stack, *kvs.""" + """inputs_embeds, rope, ctx, kv_start, last_token_ids, ds_stack, *kvs. + + ``ds_stack`` is ``[num_layers, B, S, H]`` for every family. PI05 fills it + with zeros so the per-layer add is a no-op; GR00T writes residuals. + """ decoder = getattr(language, "model", language) cfg = language.config bsz, prompt_len, hidden = inputs_embeds.shape @@ -36,7 +40,7 @@ def causal_lm_flat( ctx_len = torch.full((bsz,), seq_len, device=device, dtype=torch.int32) last_token_ids = torch.full((bsz, 1), seq_len - 1, device=device, dtype=torch.int64) kv_start = torch.empty(0, dtype=torch.int32, device=device) - ds_stack = torch.zeros(0, bsz, seq_len, hidden, device=device, dtype=dtype) + ds_stack = torch.zeros(num_layers, bsz, seq_len, hidden, device=device, dtype=dtype) kvs = [ torch.zeros( bsz, 2, num_kv, int(max_seq_len), head_dim, device=device, dtype=dtype diff --git a/tools/hf/exporters/models/common/patches.py b/tools/hf/exporters/models/common/patches.py index c9434f0a27..c2d916325c 100644 --- a/tools/hf/exporters/models/common/patches.py +++ b/tools/hf/exporters/models/common/patches.py @@ -27,34 +27,6 @@ def language_decoder(language: nn.Module) -> nn.Module: raise AttributeError(f"{type(language).__name__} has no decoder .layers") -def gather_last_token_hidden( - hidden_states: torch.Tensor, - last_token_ids: torch.Tensor, -) -> torch.Tensor: - """Gather [B, S, H] at last_token_ids [B] or [B, 1] -> [B, H] for lm_head.""" - if last_token_ids.ndim == 1: - indices = last_token_ids - else: - indices = last_token_ids.squeeze(-1) - batch_idx = torch.arange( - hidden_states.shape[0], - device=hidden_states.device, - dtype=torch.long, - ) - return hidden_states[batch_idx, indices] - - -def _lm_head_logits( - lm: nn.Module, lm_head: nn.Module | None, last_hidden: torch.Tensor -) -> torch.Tensor: - if lm_head is not None: - return lm_head(last_hidden).float() - embed = getattr(lm, "embed_tokens", None) - if embed is None: - raise AttributeError(f"{type(lm).__name__} has no lm_head or embed_tokens") - return F.linear(last_hidden, embed.weight).float() - - def causal_lm_plugin_forward( lm: nn.Module, inputs_embeds: torch.Tensor, @@ -67,11 +39,14 @@ def causal_lm_plugin_forward( lm_head: nn.Module | None = None, select_layer: int = -1, ): - """Prefill loop used by Edge language engines (plugin attention + prefix KV).""" + """Prefill loop used by Edge language engines (plugin attention + prefix KV). + + ``ds_stack`` is ``[num_layers, B, S, H]``. Each layer adds its slice; PI05 + passes zeros so the add does not change hidden. + """ lm_dtype = next(lm.parameters()).dtype hidden = _as_tensor(inputs_embeds).to(dtype=lm_dtype) seq_len = inputs_embeds.shape[1] - num_ds = int(ds_stack.shape[0]) context_hidden = hidden if select_layer == 0 else None new_kvs = [] @@ -94,8 +69,7 @@ def causal_lm_plugin_forward( hidden = residual + hidden new_kvs.append(kv) - if i < num_ds: - hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) + hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) if select_layer > 0 and (i + 1) == select_layer: context_hidden = hidden @@ -104,8 +78,16 @@ def causal_lm_plugin_forward( if context_hidden is None: context_hidden = hidden - last_hidden = gather_last_token_hidden(hidden, last_token_ids) - logits = _lm_head_logits(lm, lm_head, last_hidden) + indices = last_token_ids if last_token_ids.ndim == 1 else last_token_ids.squeeze(-1) + last_hidden = hidden[ + torch.arange(hidden.shape[0], device=hidden.device, dtype=torch.long), + indices, + ] + if lm_head is not None: + logits = lm_head(last_hidden).float() + else: + embed = getattr(lm, "embed_tokens", None) + logits = F.linear(last_hidden, embed.weight).float() prefix_k = torch.stack([kv[:, 0, :, :seq_len, :] for kv in new_kvs], dim=0) prefix_v = torch.stack([kv[:, 1, :, :seq_len, :] for kv in new_kvs], dim=0) return logits, context_hidden, prefix_k, prefix_v diff --git a/tools/hf/exporters/models/groot/spec.py b/tools/hf/exporters/models/groot/spec.py index 82a542b57b..c019e5851b 100644 --- a/tools/hf/exporters/models/groot/spec.py +++ b/tools/hf/exporters/models/groot/spec.py @@ -42,8 +42,6 @@ def _causal_lm(language: nn.Module) -> nn.Module: @register_edge_spec("groot", "gr00t") class GrootSpec(EdgeSpec): # type: ignore[misc] - components = ("vision", "language", "context_projection", "action") - def apply_patches(self, model=None): return apply_groot_patches(model) @@ -118,141 +116,241 @@ def prepare_sample_inputs( "embodiment_id": make_embodiment_id(policy, state, device, torch.long), } + def capture_eager_outputs( + self, model, sample, config, bench=None + ) -> dict[str, torch.Tensor]: + del config + from ...measure import cuda_ms + + found = _groot(model) + eagle = found.backbone.eagle_model + language = _causal_lm(eagle.language_model) + px = sample["pixel_values"] + lm_hidden = sample["lm_hidden"] + action_head = found.action_head + + with torch.no_grad(): + visual_embeds = eagle.extract_feature(px) + lm = language( + inputs_embeds=sample["inputs_embeds"], + attention_mask=sample.get("attention_mask"), + return_dict=True, + ) + context_embs = found.backbone.eagle_linear(lm_hidden) + vlln = found.action_head.vlln + weight = getattr(vlln, "weight", None) + if weight is not None: + context_embs = context_embs.to(dtype=weight.dtype) + context_embs = vlln(context_embs) + context_embs = found.action_head.vl_self_attention(context_embs) + state_features = action_head.state_encoder( + sample["state"], sample["embodiment_id"] + ) + action_features = action_head.action_encoder( + sample["step_actions"], + sample["step_timestep"], + sample["embodiment_id"], + ) + if action_head.config.add_pos_embed: + pos_ids = torch.arange( + action_features.shape[1], + dtype=torch.long, + device=action_features.device, + ) + action_features = action_features + action_head.position_embedding( + pos_ids + ).unsqueeze(0) + future_tokens = action_head.future_tokens.weight.unsqueeze(0).expand( + sample["context_embs"].shape[0], + -1, + -1, + ) + sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) + expert_out = action_head.model( + hidden_states=sa_embs, + encoder_hidden_states=sample["context_embs"], + timestep=sample["step_timestep"], + ) + action_hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out + ) + if isinstance(action_hidden, (tuple, list)): + action_hidden = action_hidden[0] + velocity = action_head.action_decoder( + action_hidden[:, -int(action_head.config.action_horizon) :], + sample["embodiment_id"], + ) + + if bench is not None: + bench["vision"] = cuda_ms(lambda: eagle.extract_feature(px)) + bench["language"] = cuda_ms( + lambda: language( + inputs_embeds=sample["inputs_embeds"], + attention_mask=sample.get("attention_mask"), + return_dict=True, + ) + ) + return { + "vision": visual_embeds, + "language": ( + lm.last_hidden_state if hasattr(lm, "last_hidden_state") else lm[0] + ), + "context_projection": context_embs, + "action": velocity, + } + def prepare( self, - name: str, model: nn.Module, sample: MutableMapping[str, Any], - upstream: Mapping[str, Any], config: Any, - ) -> ComponentBundle: + ) -> dict[str, ComponentBundle]: from ...plugin.attention import ( ContextAttentionMaskType, ) found = _groot(model) eagle = found.backbone.eagle_model - device = sample["pixel_values"].device - dtype = sample["pixel_values"].dtype + px = sample["pixel_values"] + device = px.device + dtype = px.dtype - if name == "vision": - px = sample["pixel_values"] - return ComponentBundle( - module=_export_module(eagle, sample), - trace_args=(px,), - save_args=(px,), - input_names=["pixel_values"], - output_names=["visual_embeds"], - model_type="vit", - engine_file="visual.engine", - ) - - if name == "language": - language = _causal_lm(eagle.language_model) - input_ids = sample["input_ids"] - input_embs = language.get_input_embeddings()(input_ids) - image_token_index = getattr( - eagle, "image_token_index", eagle.config.image_token_index - ) - mask = input_ids == image_token_index - sample["image_token_mask"] = mask - vis = upstream["visual_embeds"] - hidden = input_embs.shape[-1] - flat = input_embs.clone().reshape(-1, hidden) - vis_flat = vis.reshape(-1, hidden).to(device=flat.device, dtype=flat.dtype) - n = int(mask.reshape(-1).sum().item()) - flat[mask.reshape(-1)] = vis_flat[:n] - inputs_embeds = ( - flat.reshape_as(input_embs).to(device=device, dtype=dtype).contiguous() - ) - sample["lang_embeds"] = input_embs.to(device=device, dtype=dtype) - max_seq_len = max(int(config.max_seq_len), int(inputs_embeds.shape[1])) - packed, meta = causal_lm_flat( - language, - inputs_embeds, - max_seq_len=max_seq_len, - device=device, - dtype=dtype, - ) - sample.update(split_flat_to_kwargs(packed, meta["input_names"])) + vision = ComponentBundle( + module=_export_module(eagle, sample), + trace_args=(px,), + save_args=(px,), + input_names=["pixel_values"], + output_names=["visual_embeds"], + model_type="vit", + engine_file="visual.engine", + trt_settings={ + "disable_tf32": False, + "use_fp32_acc": False, + "use_explicit_typing": False, + "decompose_attention": True, + }, + ) - return ComponentBundle( - module=_export_module(language, sample), - trace_args=packed, - save_args=packed, - input_names=meta["input_names"], - output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], - context_attention_mask_type=int(ContextAttentionMaskType.CAUSAL), - model_type="language", - engine_file="language.engine", - ) + language = _causal_lm(eagle.language_model) + input_ids = sample["input_ids"] + input_embs = language.get_input_embeddings()(input_ids) + image_token_index = getattr( + eagle, "image_token_index", eagle.config.image_token_index + ) + mask = input_ids == image_token_index + sample["image_token_mask"] = mask + with torch.no_grad(): + vis = eagle.extract_feature(px) + hidden = input_embs.shape[-1] + flat_emb = input_embs.clone().reshape(-1, hidden) + vis_flat = vis.reshape(-1, hidden).to( + device=flat_emb.device, dtype=flat_emb.dtype + ) + n = int(mask.reshape(-1).sum().item()) + flat_emb[mask.reshape(-1)] = vis_flat[:n] + inputs_embeds = ( + flat_emb.reshape_as(input_embs).to(device=device, dtype=dtype).contiguous() + ) + sample["lang_embeds"] = input_embs.to(device=device, dtype=dtype) + max_seq_len = max(int(config.max_seq_len), int(inputs_embeds.shape[1])) + packed, meta = causal_lm_flat( + language, + inputs_embeds, + max_seq_len=max_seq_len, + device=device, + dtype=dtype, + ) + sample.update(split_flat_to_kwargs(packed, meta["input_names"])) - if name == "context_projection": - hidden = upstream["lm_hidden"].to(dtype=dtype) - return ComponentBundle( - module=_export_module(found, sample), - trace_args=(hidden,), - save_args=(hidden,), - input_names=["lm_hidden_states"], - output_names=["vl_embs"], - model_type="context_projection", - engine_file="context_projection.engine", - ) + language_bundle = ComponentBundle( + module=_export_module(language, sample), + trace_args=packed, + save_args=packed, + input_names=meta["input_names"], + output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], + parity_output="lm_hidden_states", + context_attention_mask_type=int(ContextAttentionMaskType.CAUSAL), + model_type="language", + engine_file="language.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "assume_dynamic_shape_support": True, + }, + ) - if name == "action": - bsz = int(upstream["context_embs"].shape[0]) - horizon = int(found.action_head.config.action_horizon) - action_dim = int(found.action_head.config.action_dim) - step_actions = sample.get( - "step_actions", - torch.randn(bsz, horizon, action_dim, device=device, dtype=dtype), - ) - step_timestep = sample.get( - "step_timestep", - torch.zeros(bsz, device=device, dtype=dtype), - ) - sample["step_actions"] = step_actions - sample["step_timestep"] = step_timestep - args = ( - step_actions, - step_timestep, - upstream["context_embs"].to(device=device, dtype=dtype), - sample["state"], - sample["embodiment_id"], - ) - return ComponentBundle( - module=_export_module(found.action_head, sample), - trace_args=args, - save_args=args, - input_names=[ - "actions", - "timestep", - "context_embs", - "state", - "embodiment_id", - ], - output_names=["velocity"], - model_type="action", - engine_file="action.engine", - ) - raise KeyError(name) + bsz, seq_len, hidden_size = inputs_embeds.shape + lm_hidden = torch.zeros(bsz, seq_len, hidden_size, device=device, dtype=dtype) + sample["lm_hidden"] = lm_hidden + context_projection = ComponentBundle( + module=_export_module(found, sample), + trace_args=(lm_hidden,), + save_args=(lm_hidden,), + input_names=["lm_hidden_states"], + output_names=["vl_embs"], + model_type="context_projection", + engine_file="context_projection.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + }, + ) - def capture_upstream( - self, - name: str, - outputs: Any, - sample: Mapping[str, Any], - bundle: ComponentBundle, - ) -> dict[str, Any]: - if name == "vision": - vis = outputs[0] if isinstance(outputs, tuple) else outputs - return {"visual_embeds": vis} - if name == "language": - return {"lm_hidden": outputs[1]} - if name == "context_projection": - ctx = outputs[0] if isinstance(outputs, tuple) else outputs - return {"context_embs": ctx} - return {} + out_dim = int(found.backbone.eagle_linear.out_features) + context_embs = torch.zeros(bsz, seq_len, out_dim, device=device, dtype=dtype) + horizon = int(found.action_head.config.action_horizon) + action_dim = int(found.action_head.config.action_dim) + step_actions = sample.get( + "step_actions", + torch.randn(bsz, horizon, action_dim, device=device, dtype=dtype), + ) + step_timestep = sample.get( + "step_timestep", + torch.zeros(bsz, device=device, dtype=dtype), + ) + sample["step_actions"] = step_actions + sample["step_timestep"] = step_timestep + sample["context_embs"] = context_embs + args = ( + step_actions, + step_timestep, + context_embs, + sample["state"], + sample["embodiment_id"], + ) + action = ComponentBundle( + module=_export_module(found.action_head, sample), + trace_args=args, + save_args=args, + input_names=[ + "actions", + "timestep", + "context_embs", + "state", + "embodiment_id", + ], + output_names=["velocity"], + model_type="action", + engine_file="action.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + }, + ) + return { + "vision": vision, + "language": language_bundle, + "context_projection": context_projection, + "action": action, + } def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] diff --git a/tools/hf/exporters/models/nemotron/patches.py b/tools/hf/exporters/models/nemotron/patches.py index 7b772500d9..568a5f436c 100644 --- a/tools/hf/exporters/models/nemotron/patches.py +++ b/tools/hf/exporters/models/nemotron/patches.py @@ -5,11 +5,12 @@ from contextlib import contextmanager from typing import Any, Callable, Iterator +import torch + from ...plugin.attn_patches import ( apply_patches, register_patch, ) -from ..common.patches import gather_last_token_hidden from .helpers import _decoder, _kind NEMOTRON = "nemotron" @@ -68,7 +69,13 @@ def forward( hidden = mixer(hidden) hidden = residual + hidden hidden = decoder.norm_f(hidden) - last = gather_last_token_hidden(hidden, last_token_ids) + indices = ( + last_token_ids if last_token_ids.ndim == 1 else last_token_ids.squeeze(-1) + ) + last = hidden[ + torch.arange(hidden.shape[0], device=hidden.device, dtype=torch.long), + indices, + ] logits = self.lm_head(last).float() return (logits, *present_kv, *present_conv, *present_ssm) diff --git a/tools/hf/exporters/models/nemotron/spec.py b/tools/hf/exporters/models/nemotron/spec.py index b4a283f410..2f41a99dc7 100644 --- a/tools/hf/exporters/models/nemotron/spec.py +++ b/tools/hf/exporters/models/nemotron/spec.py @@ -28,8 +28,6 @@ @register_edge_spec("nemotron_h", "nemotron") class NemotronSpec(EdgeSpec): # type: ignore[misc] - components = ("language",) - def apply_patches(self, model=None): return apply_nemotron_patches(model) @@ -65,17 +63,33 @@ def prepare_sample_inputs( "bsz": embeddings.shape[0], } + def capture_eager_outputs( + self, model, sample, config, bench=None + ) -> dict[str, torch.Tensor]: + del config + from ...measure import cuda_ms + + kwargs = { + "inputs_embeds": sample["inputs_embeds"], + "return_dict": True, + } + if sample.get("attention_mask") is not None: + kwargs["attention_mask"] = sample["attention_mask"] + with torch.no_grad(): + out = model(**kwargs) + logits = out.logits if hasattr(out, "logits") else out[0] + if bench is not None: + bench["language"] = cuda_ms(lambda: model(**kwargs)) + return {"language": logits} + def prepare( self, - name: str, model: nn.Module, sample: MutableMapping[str, Any], - upstream: Mapping[str, Any], config: Any, - ) -> ComponentBundle: + ) -> dict[str, ComponentBundle]: from ...rope import make_rope_rotary_cos_sin - del name, upstream embeds = sample["inputs_embeds"] device, dtype = embeds.device, embeds.dtype bsz, seq_len, _ = embeds.shape @@ -107,18 +121,27 @@ def prepare( *[f"ssm_state_{i}" for i in range(nm)], ] sample.update(split_flat_to_kwargs(flat, names)) - return ComponentBundle( - module=model.eval(), - trace_args=flat, - save_args=flat, - input_names=names, - output_names=["logits"] - + [f"present_kv_{i}" for i in range(na)] - + [f"present_conv_{i}" for i in range(nm)] - + [f"present_ssm_{i}" for i in range(nm)], - model_type="nemotron", - engine_file="language.engine", - ) + return { + "language": ComponentBundle( + module=model.eval(), + trace_args=flat, + save_args=flat, + input_names=names, + output_names=["logits"] + + [f"present_kv_{i}" for i in range(na)] + + [f"present_conv_{i}" for i in range(nm)] + + [f"present_ssm_{i}" for i in range(nm)], + model_type="nemotron", + engine_file="language.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "assume_dynamic_shape_support": True, + }, + ) + } def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: leading = [ diff --git a/tools/hf/exporters/models/pi05/helpers.py b/tools/hf/exporters/models/pi05/helpers.py index 5a1608837c..0808afac77 100644 --- a/tools/hf/exporters/models/pi05/helpers.py +++ b/tools/hf/exporters/models/pi05/helpers.py @@ -1,7 +1,6 @@ from __future__ import annotations import torch -import torch.nn as nn from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks @@ -225,15 +224,6 @@ def make_pi05_suffix_position_and_mask(core, prefix_pad_masks, x_t, device): return position_ids, attention_mask -def _core(model: nn.Module) -> nn.Module: - if hasattr(model, "paligemma_with_expert"): - return model - inner = getattr(model, "model", None) - if isinstance(inner, nn.Module) and hasattr(inner, "paligemma_with_expert"): - return inner - raise RuntimeError("PI05 spec expected a policy or paligemma_with_expert module") - - def _nchw_to_hwc(pixel_values): if pixel_values.ndim != 4: return pixel_values diff --git a/tools/hf/exporters/models/pi05/spec.py b/tools/hf/exporters/models/pi05/spec.py index d53dd05dcc..c5f0bb74b6 100644 --- a/tools/hf/exporters/models/pi05/spec.py +++ b/tools/hf/exporters/models/pi05/spec.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +import torch_tensorrt from ...ops import call_engine, fuse_prefix from ...spec import ( @@ -19,7 +20,6 @@ ) from ..common.patches import language_decoder from .helpers import ( - _core, build_pi05_prefix_embs, make_pi05_suffix_position_and_mask, pi05_compact_index, @@ -29,8 +29,6 @@ @register_edge_spec("pi05") class Pi05Spec(EdgeSpec): # type: ignore[misc] - components = ("vision", "language", "action") - def apply_patches(self, model=None): """Install vision, language, and action setattr replacements.""" del model @@ -38,6 +36,95 @@ def apply_patches(self, model=None): return apply_patches(PI05) + def create_dynamic_shapes( + self, + input_names: list[str], + trace_args: tuple[Any, ...], + *, + max_seq_len: int, + ) -> tuple[Any, ...]: + """Prefill/decode ``torch_tensorrt.Input`` specs (e2e language dual-profile).""" + named = dict(zip(input_names, trace_args)) + embs = named["inputs_embeds"] + ds = named["ds_stack"] + kv = next( + tensor + for name, tensor in zip(input_names, trace_args) + if name.startswith("past_key_values_") + ) + bsz = int(embs.shape[0]) + hidden = int(embs.shape[-1]) + opt_prefill = max(int(max_seq_len) // 2, 1) + num_ds = int(ds.shape[0]) + num_kv = int(kv.shape[2]) + head_dim = int(kv.shape[-1]) + prefill_profile = { + "min_shape": (1, 1, hidden), + "opt_shape": (bsz, opt_prefill, hidden), + "max_shape": (bsz, max_seq_len, hidden), + } + decode_profile = { + "min_shape": (1, 1, hidden), + "opt_shape": (bsz, 1, hidden), + "max_shape": (bsz, 1, hidden), + } + kv_profile = { + "min_shape": (1, 2, num_kv, 1, head_dim), + "opt_shape": (bsz, 2, num_kv, max_seq_len, head_dim), + "max_shape": (bsz, 2, num_kv, max_seq_len, head_dim), + } + ds_prefill = { + "min_shape": (num_ds, 1, 1, hidden), + "opt_shape": (num_ds, bsz, opt_prefill, hidden), + "max_shape": (num_ds, bsz, max_seq_len, hidden), + } + ds_decode = { + "min_shape": (num_ds, 1, 1, hidden), + "opt_shape": (num_ds, bsz, 1, hidden), + "max_shape": (num_ds, bsz, 1, hidden), + } + input_specs = [] + for name, tensor in zip(input_names, trace_args): + if name == "inputs_embeds": + input_specs.append( + torch_tensorrt.Input( + profiles=[prefill_profile, decode_profile], + shared_dims={1: "seq_len"}, + dtype=tensor.dtype, + format=torch.contiguous_format, + name=name, + ) + ) + elif name == "ds_stack": + input_specs.append( + torch_tensorrt.Input( + profiles=[ds_prefill, ds_decode], + shared_dims={2: "seq_len"}, + dtype=tensor.dtype, + format=torch.contiguous_format, + name=name, + ) + ) + elif name.startswith("past_key_values_"): + input_specs.append( + torch_tensorrt.Input( + profiles=[kv_profile, kv_profile], + dtype=tensor.dtype, + format=torch.contiguous_format, + name=name, + ) + ) + else: + input_specs.append( + torch_tensorrt.Input( + shape=tuple(tensor.shape), + dtype=tensor.dtype, + format=torch.contiguous_format, + name=name, + ) + ) + return tuple(input_specs) + def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any ) -> MutableMapping[str, Any]: @@ -79,7 +166,7 @@ def prepare_sample_inputs( ).contiguous() tokens = batch[OBS_LANGUAGE_TOKENS].to(device=device, dtype=torch.long) masks = batch[OBS_LANGUAGE_ATTENTION_MASK].to(device=device, dtype=torch.bool) - core = _core(policy) + core = policy if hasattr(policy, "paligemma_with_expert") else policy.model lang_embeds = core.paligemma_with_expert.embed_language_tokens(tokens) return { "pixel_values": pixel_values, @@ -90,145 +177,278 @@ def prepare_sample_inputs( "lang_embeds": lang_embeds.to(device=device, dtype=dtype).contiguous(), } - def prepare( - self, - name: str, - model: nn.Module, - sample: MutableMapping[str, Any], - upstream: Mapping[str, Any], - config: Any, - ) -> ComponentBundle: - from ...plugin.attention import ( - ContextAttentionMaskType, - ) + def capture_eager_outputs( + self, model, sample, config, bench=None + ) -> dict[str, torch.Tensor]: + del config + from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding - core = _core(model) - paligemma = core.paligemma_with_expert.paligemma.model - device = sample["pixel_values"].device - dtype = sample["pixel_values"].dtype + from ...measure import cuda_ms + from ...prefix_cache import PrefixKVCache - if name == "vision": - px = sample["pixel_values"] - return ComponentBundle( - module=paligemma.eval(), - trace_args=(px,), - save_args=(px,), - input_names=["pixel_values"], - output_names=["visual_embeds"], - model_type="vit", - engine_file="visual.engine", - ) + core = model if hasattr(model, "paligemma_with_expert") else model.model + paligemma = core.paligemma_with_expert.paligemma.model + language = paligemma.language_model + px = sample["pixel_values"] - if name == "language": - paligemma = core.paligemma_with_expert.paligemma.model - language = paligemma.language_model - embs, pad, _attn, _pos = build_pi05_prefix_embs( - core, - sample["img_masks"], - sample["tokens"], - sample["masks"], - upstream["visual_embeds"], - sample["images"], + with torch.no_grad(): + tower_out = paligemma.vision_tower(px) + hidden = getattr(tower_out, "last_hidden_state", tower_out) + visual_embeds = paligemma.multi_modal_projector(hidden) + lm_dtype = next(language.parameters()).dtype + prefix_embs = sample["prefix_embs"].to(dtype=lm_dtype) + lm = language( + inputs_embeds=prefix_embs, + attention_mask=sample["prefix_attention_mask"], + position_ids=sample["prefix_position_ids"], + return_dict=True, ) - compact_len = int(embs.shape[1]) - vis = upstream["visual_embeds"] - per_cam = int(sample["images"][0].shape[0]) - seq_per_image = int( - vis.reshape(len(sample["images"]), per_cam, -1, vis.shape[-1]).shape[2] + suffix_embs = core.action_in_proj(sample["step_actions"]) + time_emb = create_sinusoidal_pos_embedding( + sample["step_timestep"], + core.action_in_proj.out_features, + min_period=core.config.min_period, + max_period=core.config.max_period, + device=sample["step_timestep"].device, + ).to(dtype=suffix_embs.dtype) + adarms_cond = torch.nn.functional.silu( + core.time_mlp_out(torch.nn.functional.silu(core.time_mlp_in(time_emb))) ) - sample["compact_index"] = pi05_compact_index( - sample["img_masks"], - sample["images"], - seq_per_image, - sample["masks"], - device, + expert_out = core.paligemma_with_expert.gemma_expert.model( + inputs_embeds=suffix_embs, + attention_mask=sample["suffix_attention_mask"], + position_ids=sample["suffix_position_ids"], + past_key_values=PrefixKVCache(sample["prefix_k"], sample["prefix_v"]), + use_cache=False, + adarms_cond=adarms_cond, ) - sample["prefix_pad_mask"] = pad - max_seq_len = max( - int(config.max_seq_len), compact_len + int(config.generation_reserve) + action_hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out ) - flat, meta = causal_lm_flat( - language, - embs.to(device=device, dtype=dtype), - max_seq_len=max_seq_len, - device=device, - dtype=dtype, - seq_len=compact_len, - ) - sample.update(split_flat_to_kwargs(flat, meta["input_names"])) - - return ComponentBundle( - module=language_decoder(language).eval(), - trace_args=flat, - save_args=flat, - input_names=meta["input_names"], - output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], - context_attention_mask_type=int(ContextAttentionMaskType.PADDING), - extra_config={"prefix_pad_mask_len": compact_len}, - model_type="language", - engine_file="language.engine", + if isinstance(action_hidden, (tuple, list)): + action_hidden = action_hidden[0] + velocity = core.action_out_proj( + action_hidden[:, -int(core.config.chunk_size) :] ) - if name == "action": - bsz = int(sample["lang_embeds"].shape[0]) - core_mod = _core(model) - step_actions = sample.get("step_actions") - if step_actions is None: - step_actions = torch.randn( - bsz, - int(core_mod.config.chunk_size), - int(core_mod.config.max_action_dim), - device=device, - dtype=dtype, + if bench is not None: + bench["vision"] = cuda_ms( + lambda: paligemma.multi_modal_projector( + paligemma.vision_tower(px).last_hidden_state ) - sample["step_actions"] = step_actions - step_timestep = sample.get( - "step_timestep", - torch.full((bsz,), 1.0, device=device, dtype=torch.float32), ) - sample["step_timestep"] = step_timestep - prefix_k = upstream["prefix_k"].to(device=device, dtype=dtype) - prefix_v = upstream["prefix_v"].to(device=device, dtype=dtype) - pos, mask = make_pi05_suffix_position_and_mask( # type: ignore[no-untyped-call] - core_mod, sample["prefix_pad_mask"], step_actions, device + bench["language"] = cuda_ms( + lambda: language( + inputs_embeds=prefix_embs, + attention_mask=sample["prefix_attention_mask"], + position_ids=sample["prefix_position_ids"], + return_dict=True, + ) ) - sample["suffix_position_ids"] = pos - sample["suffix_attention_mask"] = mask - args = (step_actions, step_timestep, prefix_k, prefix_v, pos, mask) - return ComponentBundle( - module=core.eval(), - trace_args=args, - save_args=args, - input_names=[ - "x_t", - "timestep", - "prefix_k", - "prefix_v", - "position_ids", - "attention_mask", - ], - output_names=["velocity"], - model_type="action", - engine_file="action.engine", + bench["action"] = cuda_ms( + lambda: core.action_out_proj( + core.paligemma_with_expert.gemma_expert.model( + inputs_embeds=suffix_embs, + attention_mask=sample["suffix_attention_mask"], + position_ids=sample["suffix_position_ids"], + past_key_values=PrefixKVCache( + sample["prefix_k"], sample["prefix_v"] + ), + use_cache=False, + adarms_cond=adarms_cond, + ).last_hidden_state[:, -int(core.config.chunk_size) :] + ) ) - raise KeyError(name) + return { + "vision": visual_embeds, + "language": lm.last_hidden_state, + "action": velocity, + } - def capture_upstream( + def prepare( self, - name: str, - outputs: Any, - sample: Mapping[str, Any], - bundle: ComponentBundle, - ) -> dict[str, Any]: - if name == "vision": - vis = outputs[0] if isinstance(outputs, tuple) else outputs - # Engine output is [B, S, H]. Language packing still uses [N, H]. - if vis.ndim == 3: - vis = vis.reshape(-1, vis.shape[-1]) - return {"visual_embeds": vis} - if name == "language": - return {"prefix_k": outputs[2], "prefix_v": outputs[3]} - return {} + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + from ...plugin.attention import ( + ContextAttentionMaskType, + ) + + core = model if hasattr(model, "paligemma_with_expert") else model.model + paligemma = core.paligemma_with_expert.paligemma.model + language = paligemma.language_model + px = sample["pixel_values"] + device = px.device + dtype = px.dtype + + vision = ComponentBundle( + module=paligemma.eval(), + trace_args=(px,), + save_args=(px,), + input_names=["pixel_values"], + output_names=["visual_embeds"], + model_type="vit", + engine_file="visual.engine", + trt_settings={ + "disable_tf32": False, + "use_fp32_acc": False, + "use_explicit_typing": False, + "decompose_attention": True, + }, + ) + + with torch.no_grad(): + tower_out = paligemma.vision_tower(px) + hidden = getattr(tower_out, "last_hidden_state", tower_out) + visual_embeds = paligemma.multi_modal_projector(hidden) + if visual_embeds.ndim == 3: + visual_embeds = visual_embeds.reshape(-1, visual_embeds.shape[-1]) + + embs, pad, attn, pos = build_pi05_prefix_embs( + core, + sample["img_masks"], + sample["tokens"], + sample["masks"], + visual_embeds, + sample["images"], + ) + compact_len = int(embs.shape[1]) + per_cam = int(sample["images"][0].shape[0]) + seq_per_image = int( + visual_embeds.reshape( + len(sample["images"]), per_cam, -1, visual_embeds.shape[-1] + ).shape[2] + ) + sample["compact_index"] = pi05_compact_index( + sample["img_masks"], + sample["images"], + seq_per_image, + sample["masks"], + device, + ) + sample["prefix_embs"] = embs + sample["prefix_pad_mask"] = pad + sample["prefix_attention_mask"] = attn + sample["prefix_position_ids"] = pos + if int(config.generation_reserve) < 0: + raise ValueError("generation_reserve must be non-negative") + max_seq_len = max( + int(config.max_seq_len), + compact_len + int(config.generation_reserve), + ) + flat, meta = causal_lm_flat( + language, + embs.to(device=device, dtype=dtype), + max_seq_len=max_seq_len, + device=device, + dtype=dtype, + seq_len=compact_len, + ) + sample.update(split_flat_to_kwargs(flat, meta["input_names"])) + + embs_t, rope, ctx, kv_start, last, ds, *kvs = flat + opt_prefill = max(max_seq_len // 2, 1) + trace_len = min(int(embs_t.shape[1]), opt_prefill) + trace_args = ( + embs_t[:, :trace_len].contiguous(), + rope, + torch.full_like(ctx, trace_len), + kv_start, + torch.full_like(last, trace_len - 1), + ds[:, :, :trace_len].contiguous(), + *kvs, + ) + + decoder = language_decoder(language) + language_bundle = ComponentBundle( + module=decoder.eval(), + trace_args=trace_args, + save_args=flat, + execute_args=flat, + input_specs=self.create_dynamic_shapes( + meta["input_names"], trace_args, max_seq_len=max_seq_len + ), + input_names=meta["input_names"], + output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], + parity_output="lm_hidden_states", + context_attention_mask_type=int(ContextAttentionMaskType.PADDING), + extra_config={"prefix_pad_mask_len": compact_len}, + model_type="language", + engine_file="language.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "assume_dynamic_shape_support": True, + }, + ) + + bsz = int(sample["lang_embeds"].shape[0]) + step_actions = sample.get("step_actions") + if step_actions is None: + step_actions = torch.randn( + bsz, + int(core.config.chunk_size), + int(core.config.max_action_dim), + device=device, + dtype=dtype, + ) + sample["step_actions"] = step_actions + step_timestep = sample.get( + "step_timestep", + torch.full((bsz,), 1.0, device=device, dtype=torch.float32), + ) + sample["step_timestep"] = step_timestep + cfg = language.config + num_kv = int(cfg.num_key_value_heads) + head_dim = int( + getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) + ) + prefix_k = torch.zeros( + len(decoder.layers), + bsz, + num_kv, + compact_len, + head_dim, + device=device, + dtype=dtype, + ) + prefix_v = torch.zeros_like(prefix_k) + sample["prefix_k"] = prefix_k + sample["prefix_v"] = prefix_v + pos, mask = make_pi05_suffix_position_and_mask( # type: ignore[no-untyped-call] + core, sample["prefix_pad_mask"], step_actions, device + ) + sample["suffix_position_ids"] = pos + sample["suffix_attention_mask"] = mask + args = (step_actions, step_timestep, prefix_k, prefix_v, pos, mask) + action = ComponentBundle( + module=core.eval(), + trace_args=args, + save_args=args, + input_names=[ + "x_t", + "timestep", + "prefix_k", + "prefix_v", + "position_ids", + "attention_mask", + ], + output_names=["velocity"], + model_type="action", + engine_file="action.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + }, + ) + return {"vision": vision, "language": language_bundle, "action": action} def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] diff --git a/tools/hf/exporters/ops.py b/tools/hf/exporters/ops.py index 9755d451af..a71f791827 100644 --- a/tools/hf/exporters/ops.py +++ b/tools/hf/exporters/ops.py @@ -2,12 +2,19 @@ from __future__ import annotations +import sys from typing import Any import torch -_ENGINE_META: dict[str, dict[str, Any]] = {} -_COMPILED_MODULES: dict[str, torch.nn.Module] = {} +# One process-wide table so pytest dual-imports of this module still share +# execute_engine state with record_engine. +_REGISTRY: dict[str, Any] = sys.modules.setdefault( + "_edge_llm_engine_registry", + {"meta": {}, "modules": {}}, +) +_ENGINE_META: dict[str, dict[str, Any]] = _REGISTRY["meta"] +_COMPILED_MODULES: dict[str, torch.nn.Module] = _REGISTRY["modules"] def record_engine( diff --git a/tools/hf/exporters/plugin/attn_patches.py b/tools/hf/exporters/plugin/attn_patches.py index 4b30576132..2693dbb853 100644 --- a/tools/hf/exporters/plugin/attn_patches.py +++ b/tools/hf/exporters/plugin/attn_patches.py @@ -2,8 +2,8 @@ Same contract as ``transformers.exporters.utils.register_patch``: one factory per backend, listed against every attention class that shares that layout. Patches are -installed only while ``apply_patches`` is active (or left installed on dryrun so -``execute_engine`` still hits the plugin). +installed only while ``apply_patches`` is active so ``torch.export`` sees +plugin I/O. Eager inference uses the original HuggingFace forward. Language dispatch: Edge prefill calls ``self_attn(..., rope_rotary_cos_sin=...)``. The PI05 action expert is often the same class (GemmaAttention / PiGemmaModel) @@ -224,14 +224,33 @@ def forward(self, hidden_states, attention_mask=None, **kwargs): "transformers.models.qwen3.modeling_qwen3.Qwen3Attention.forward", ) def _patch_language_attention(original: Callable) -> Callable: - def forward(self, hidden_states, *args, **kwargs): - rope_rotary_cos_sin = kwargs.get("rope_rotary_cos_sin") + """Same I/O as ``PluginAttention.forward``. + + ``rope_rotary_cos_sin`` is a real parameter (not ``kwargs.get``) so + ``torch.export`` specializes the plugin branch. The HF expert path is + ``rope_rotary_cos_sin is None``. + """ + + def forward( + self, + hidden_states, + rope_rotary_cos_sin=None, + attention_mask=None, + position_ids=None, + past_key_value=None, + ctx_len=None, + kvcache_start_index=None, + **kwargs, + ): if rope_rotary_cos_sin is None: - return original(self, hidden_states, *args, **kwargs) + return original( + self, + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) - past_key_value = kwargs.get("past_key_value") - ctx_len = kwargs.get("ctx_len") - kvcache_start_index = kwargs.get("kvcache_start_index") if rope_rotary_cos_sin.dtype != torch.float32: raise ValueError("rope_rotary_cos_sin must be FP32") if past_key_value is None: diff --git a/tools/hf/exporters/spec.py b/tools/hf/exporters/spec.py index 0186cdf260..3dbea351df 100644 --- a/tools/hf/exporters/spec.py +++ b/tools/hf/exporters/spec.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from typing import Any +import torch import torch.nn as nn _SPECS: dict[str, type[EdgeSpec]] = {} @@ -35,11 +36,12 @@ class ComponentBundle: save_args: tuple[Any, ...] input_names: list[str] output_names: list[str] + parity_output: str | None = None extra_config: dict[str, Any] = field(default_factory=dict) trt_settings: dict[str, Any] = field(default_factory=dict) - patch_fn: Callable[[nn.Module], Any] | None = None context_attention_mask_type: int | None = None execute_args: tuple[Any, ...] | None = None + input_specs: tuple[Any, ...] | None = None model_type: str = "edge" engine_file: str = "engine.engine" @@ -47,20 +49,20 @@ class ComponentBundle: class EdgeSpec(ABC): """Per-family flatten / runtime wiring. - ``EdgeExporter.export`` never branches on PI05 vs Nemotron. It only loops - ``spec.components``. + ``EdgeExporter.export`` never branches on PI05 vs Nemotron. It compiles + whatever ``prepare`` returns. """ - components: tuple[str, ...] = () - def apply_patches( self, model: nn.Module | None = None ) -> AbstractContextManager[None]: - """Install this family's setattr replacements for the whole ``export()``. + """Install this family's setattr replacements for TensorRT tracing. - Default is a no-op. Families register factories on their own backend - and return ``apply_patches(backend)``. ``model`` is the export root; - Nemotron uses it to wrap hybrid mixers. + Installed only around ``compile_component``. Eager inference is the + original HuggingFace / LeRobot forward. Default is a no-op. Families + register factories on their own backend and return + ``apply_patches(backend)``. ``model`` is the export root; Nemotron uses + it to wrap hybrid mixers. """ del model return nullcontext() @@ -75,32 +77,53 @@ def prepare_sample_inputs( """Caller payload → stem dict used by prepare/run.""" @abstractmethod - def prepare( + def capture_eager_outputs( self, - name: str, model: nn.Module, sample: MutableMapping[str, Any], - upstream: Mapping[str, Any], config: Any, - ) -> ComponentBundle: - """Select the original submodule and build its trace/save tuple.""" + bench: dict[str, float] | None = None, + ) -> dict[str, torch.Tensor]: + """Unpatched HF / LeRobot tensors, keyed like ``prepare``. + + Called before ``apply_patches``. One tensor per component: the value + e2e passes to ``parity`` (vision embeds, language ``last_hidden_state``, + action velocity). Optional ``bench`` records unpatched CUDA-event ms. + """ + + def create_dynamic_shapes( + self, + input_names: list[str], + trace_args: tuple[Any, ...], + *, + max_seq_len: int, + ) -> tuple[Any, ...] | None: + """``torch_tensorrt.Input`` specs for dual-profile language compile. + + Default is no specs (static ``trace_args``). PI05 overrides this. + """ + del input_names, trace_args, max_seq_len + return None - def capture_upstream( + @abstractmethod + def prepare( self, - name: str, - outputs: Any, - sample: Mapping[str, Any], - bundle: ComponentBundle, - ) -> dict[str, Any]: - """Map this engine's outputs into keys the next ``prepare`` needs.""" - return {} + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + """Build every component bundle in one call. + + Example tensors for later stages come from unpatched eager packing + (or zeros of the trace shape), not from the previous engine. + """ @abstractmethod def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: """Packing + ``execute_engine`` calls. This is the dumped graph.""" def runtime_kwargs(self, sample: Mapping[str, Any]) -> dict[str, Any]: - """Tensor kwargs for ``torch.export`` of :class:`EdgeRuntimeModule`.""" + """Tensor kwargs for ``torch.export`` of the outer VLA graph.""" return { key: value for key, value in sample.items() diff --git a/tools/hf/exporters/tests/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py index 23437408fd..590ae602b2 100644 --- a/tools/hf/exporters/tests/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -6,29 +6,55 @@ import pytest import torch import torch.nn as nn +import torch_tensorrt from exporters import EdgeConfig, EdgeExporter, register_edge_spec from exporters.ops import call_engine from exporters.spec import ComponentBundle, EdgeSpec, registered_specs +from torch.export import ExportedProgram + + +def _install_fake_trt(monkeypatch) -> None: + monkeypatch.setattr(torch_tensorrt.dynamo, "compile", _fake_trt_compile) + monkeypatch.setattr( + torch_tensorrt.dynamo, + "convert_exported_program_to_serialized_trt_engine", + lambda *args, **kwargs: b"fake-engine", + ) + + +def _fake_trt_compile(exported, arg_inputs=None, **kwargs): + del arg_inputs, kwargs + return exported.module() @register_edge_spec("dummy_edge") class DummySpec(EdgeSpec): - components = ("language",) - def prepare_sample_inputs(self, model, raw, config): return {"x": raw["x"]} - def prepare(self, name, model, sample, upstream, config) -> ComponentBundle: + def capture_eager_outputs(self, model, sample, config, bench=None): + del config + from exporters.measure import cuda_ms + + with torch.no_grad(): + y = model(sample["x"]) + if bench is not None: + bench["language"] = cuda_ms(lambda: model(sample["x"])) + return {"language": y} + + def prepare(self, model, sample, config) -> dict[str, ComponentBundle]: x = sample["x"] - return ComponentBundle( - module=model.eval(), - trace_args=(x,), - save_args=(x,), - input_names=["x"], - output_names=["y"], - model_type="dummy", - engine_file="language.engine", - ) + return { + "language": ComponentBundle( + module=model.eval(), + trace_args=(x,), + save_args=(x,), + input_names=["x"], + output_names=["y"], + model_type="dummy", + engine_file="language.engine", + ) + } def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]): return call_engine(engines["language"], "language", sample["x"])[0] @@ -44,137 +70,29 @@ def test_builtin_specs_are_registered(): @pytest.mark.unit -def test_edge_exporter_dryrun_runtime(tmp_path): - torch.manual_seed(0) - model = nn.Linear(4, 4) - sample = {"x": torch.randn(2, 4)} - exporter = EdgeExporter() - runtime = exporter.export( - model, - sample, - EdgeConfig( - dryrun=True, - skip_runtime_export=True, - model_type="dummy_edge", - engine_dir=tmp_path, - ), - ) - assert "language" in exporter.engines - assert (tmp_path / "language" / "config.json").is_file() - with torch.no_grad(): - got = runtime(x=sample["x"]) - expected = model(sample["x"]) - torch.testing.assert_close(got, expected) - - -@pytest.mark.unit -def test_edge_exporter_dryrun_exported_program(tmp_path): +def test_edge_exporter_exported_program(tmp_path, monkeypatch): + _install_fake_trt(monkeypatch) torch.manual_seed(0) model = nn.Linear(4, 4) - # Packing tensors are intermediates, not graph leaves. sample = {"x": torch.randn(2, 4) + 1} exporter = EdgeExporter() program = exporter.export( model, sample, EdgeConfig( - dryrun=True, model_type="dummy_edge", engine_dir=tmp_path, ), ) - assert program is not None + assert isinstance(program, ExportedProgram) + assert "language" in exporter.engines + assert (tmp_path / "language" / "config.json").is_file() with torch.no_grad(): out = program.module()(x=sample["x"]) expected = model(sample["x"]) torch.testing.assert_close(out, expected) -class _NativeAttn(nn.Module): - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4) - - def forward(self, hidden_states, **kwargs): - raise TypeError("cannot unpack non-iterable NoneType object") - - -class _PluginAttn(nn.Module): - def __init__(self, inner: nn.Module): - super().__init__() - self.linear = inner.linear - - def forward(self, hidden_states, **kwargs): - return self.linear(hidden_states) - - -class _Layer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = _NativeAttn() - - -class _PatchedWrapper(nn.Module): - def __init__(self): - super().__init__() - self.layer = _Layer() - - def forward(self, x): - return self.layer.self_attn(x, rope_rotary_cos_sin=x) - - -@register_edge_spec("patch_edge") -class _PatchSpec(EdgeSpec): - components = ("language",) - - def prepare_sample_inputs(self, model, raw, config): - return {"x": raw["x"]} - - def prepare(self, name, model, sample, upstream, config) -> ComponentBundle: - x = sample["x"] - - def _patch(mod): - orig = mod.layer.self_attn - mod.layer.self_attn = _PluginAttn(orig).eval() - return [(mod.layer, orig)] - - return ComponentBundle( - module=model.eval(), - trace_args=(x,), - save_args=(x,), - input_names=["x"], - output_names=["y"], - patch_fn=_patch, - model_type="dummy", - engine_file="language.engine", - ) - - def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]): - return call_engine(engines["language"], "language", sample["x"])[0] - - -@pytest.mark.unit -def test_edge_exporter_dryrun_keeps_attention_patch(tmp_path): - """Language wrappers pass plugin kwargs; native HF attention cannot run them.""" - torch.manual_seed(0) - model = _PatchedWrapper() - sample = {"x": torch.randn(2, 4)} - exporter = EdgeExporter() - runtime = exporter.export( - model, - sample, - EdgeConfig( - dryrun=True, - skip_runtime_export=True, - model_type="patch_edge", - engine_dir=tmp_path, - ), - ) - with torch.no_grad(): - got = runtime(x=sample["x"]) - assert got.shape == (2, 4) - - @pytest.mark.unit def test_attn_patch_attribute_restores(): from exporters.plugin.attn_patches import patch_attribute @@ -364,7 +282,7 @@ def test_eagle_vision_patch_extracts_features(): _patch_eagle_image_features, ) - class Dummy: + class Dummy(nn.Module): def extract_feature(self, pixel_values): return pixel_values + 1 @@ -380,7 +298,7 @@ def forward(self, *args, **kwargs): def test_groot_patches_live_eagle_class(): from exporters.models.groot.patches import apply_groot_patches - class Eagle: + class Eagle(nn.Module): def extract_feature(self, pixel_values): return pixel_values + 1 @@ -409,7 +327,7 @@ def test_eagle_vision_keeps_vlm_forward_with_input_ids(): _patch_eagle_image_features, ) - class Dummy: + class Dummy(nn.Module): def extract_feature(self, pixel_values): raise AssertionError("extract_feature should not run") diff --git a/tools/hf/run_groot_export.py b/tools/hf/run_groot_export.py index bf918c9add..d21007e825 100644 --- a/tools/hf/run_groot_export.py +++ b/tools/hf/run_groot_export.py @@ -7,10 +7,6 @@ from __future__ import annotations -import argparse -import sys -from pathlib import Path - import torch # noqa: E402 import torch_tensorrt # noqa: E402 from exporters import EdgeConfig, EdgeExporter @@ -49,13 +45,6 @@ def load_groot(device: torch.device) -> GrootPolicy: def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--compile", action="store_true", help="Build TRT engines (default: dryrun)" - ) - parser.add_argument("--engine-dir", default="/tmp/groot_edge_exporter") - args = parser.parse_args() - load_plugins_for_trt() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -68,7 +57,9 @@ def main() -> None: force_hf_attention(eagle.language_model, "eager") exporter = EdgeExporter() - config = EdgeConfig(model_type="groot", engine_dir=args.engine_dir, max_seq_len=968) + config = EdgeConfig( + model_type="groot", engine_dir="/tmp/groot_edge_exporter", max_seq_len=968 + ) # Spec tokenizes libero via Eagle chat template because we pass the policy. sample_inputs = {"device": device, "dtype": dtype} @@ -78,10 +69,7 @@ def main() -> None: print("runtime keys:", sorted(exporter.sample)) with torch.no_grad(): - if hasattr(program, "module"): - velocity = program.module()(**exporter.sample) - else: - velocity = program(**exporter.sample) + velocity = program.module()(**exporter.sample) out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity print("velocity", tuple(out.shape), "mean", float(out.float().mean())) diff --git a/tools/hf/run_nemotron_export.py b/tools/hf/run_nemotron_export.py index 977cfe0056..14da09de3e 100644 --- a/tools/hf/run_nemotron_export.py +++ b/tools/hf/run_nemotron_export.py @@ -8,8 +8,6 @@ from __future__ import annotations import argparse -import sys -from pathlib import Path import torch import torch_tensorrt @@ -39,10 +37,6 @@ def load_nemotron(checkpoint: str, device: torch.device, dtype: torch.dtype): def main() -> None: parser = argparse.ArgumentParser() - parser.add_argument( - "--compile", action="store_true", help="Build TRT engines (default: dryrun)" - ) - parser.add_argument("--engine-dir", default="/tmp/nemotron_edge_exporter") parser.add_argument( "--checkpoint", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", @@ -66,10 +60,8 @@ def main() -> None: exporter = EdgeExporter() config = EdgeConfig( model_type="nemotron_h", - engine_dir=args.engine_dir, + engine_dir="/tmp/nemotron_edge_exporter", max_seq_len=args.max_seq_len, - dryrun=not args.compile, - skip_runtime_export=False, ) program = exporter.export(model, sample_inputs, config=config) @@ -77,10 +69,7 @@ def main() -> None: print("runtime keys:", sorted(exporter.sample)) with torch.no_grad(): - if hasattr(program, "module"): - out = program.module()(**exporter.sample) - else: - out = program(**exporter.sample) + out = program.module()(**exporter.sample) logits = out[0] if isinstance(out, (tuple, list)) else out print("logits", tuple(logits.shape), "mean", float(logits.float().mean())) diff --git a/tools/hf/run_pi05_export.py b/tools/hf/run_pi05_export.py index 2c2017c0e0..f7abd318bc 100644 --- a/tools/hf/run_pi05_export.py +++ b/tools/hf/run_pi05_export.py @@ -7,10 +7,6 @@ from __future__ import annotations -import argparse -import sys -from pathlib import Path - import torch import torch_tensorrt from exporters import EdgeConfig, EdgeExporter @@ -52,20 +48,13 @@ def load_pi05(device: torch.device) -> PI05Policy: def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--compile", action="store_true", help="Build TRT engines (default: dryrun)" - ) - parser.add_argument("--engine-dir", default="/tmp/pi05_edge_exporter") - args = parser.parse_args() - load_plugins_for_trt() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float16 policy = load_pi05(device) - # Weights on GPU; spec still needs the policy object for the preprocessor. + policy.model.to(device=device, dtype=dtype).eval() paligemma = policy.model.paligemma_with_expert.paligemma.model force_hf_attention(paligemma.vision_tower, "eager") @@ -75,15 +64,12 @@ def main() -> None: exporter = EdgeExporter() config = EdgeConfig( model_type="pi05", # optional; inferred from paligemma_with_expert - engine_dir=args.engine_dir, + engine_dir="/tmp/pi05", max_seq_len=968, ) - # Spec loads libero + preprocessor because we pass the policy, not a tensor dict. sample_inputs = {"device": device, "dtype": dtype} - program = exporter.export(policy, sample_inputs, config=config) - print("engines:", exporter.engines) # Runtime kwargs are tensors only (pixel_values, lang_embeds, rope, KVs, …). @@ -91,10 +77,7 @@ def main() -> None: print("runtime keys:", sorted(runtime_kwargs)) with torch.no_grad(): - if hasattr(program, "module"): - velocity = program.module()(**runtime_kwargs) - else: - velocity = program(**runtime_kwargs) + velocity = program.module()(**runtime_kwargs) out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity print("velocity", tuple(out.shape), "mean", float(out.float().mean()))