From 56173196357f534aa4b4e823cd16996df46e9084 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 12:43:55 -0700 Subject: [PATCH 01/11] nemotron, gr00t and pi05 running. --- examples/dynamo/run_groot_export.py | 112 ++++++ examples/dynamo/run_nemotron_export.py | 108 ++++++ examples/dynamo/run_pi05_export.py | 117 ++++++ py/torch_tensorrt/hf/__init__.py | 4 + py/torch_tensorrt/hf/exporters/__init__.py | 20 + py/torch_tensorrt/hf/exporters/compile.py | 133 +++++++ py/torch_tensorrt/hf/exporters/config.py | 30 ++ py/torch_tensorrt/hf/exporters/exporter.py | 161 ++++++++ py/torch_tensorrt/hf/exporters/ops.py | 132 +++++++ py/torch_tensorrt/hf/exporters/runtime.py | 19 + py/torch_tensorrt/hf/exporters/spec.py | 139 +++++++ .../hf/exporters/specs/__init__.py | 5 + .../hf/exporters/specs/_language.py | 79 ++++ py/torch_tensorrt/hf/exporters/specs/groot.py | 347 ++++++++++++++++++ .../hf/exporters/specs/nemotron.py | 255 +++++++++++++ py/torch_tensorrt/hf/exporters/specs/pi05.py | 316 ++++++++++++++++ setup.py | 6 + tests/py/dynamo/hf/test_edge_exporter.py | 179 +++++++++ 18 files changed, 2162 insertions(+) create mode 100644 examples/dynamo/run_groot_export.py create mode 100644 examples/dynamo/run_nemotron_export.py create mode 100644 examples/dynamo/run_pi05_export.py create mode 100644 py/torch_tensorrt/hf/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/compile.py create mode 100644 py/torch_tensorrt/hf/exporters/config.py create mode 100644 py/torch_tensorrt/hf/exporters/exporter.py create mode 100644 py/torch_tensorrt/hf/exporters/ops.py create mode 100644 py/torch_tensorrt/hf/exporters/runtime.py create mode 100644 py/torch_tensorrt/hf/exporters/spec.py create mode 100644 py/torch_tensorrt/hf/exporters/specs/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/specs/_language.py create mode 100644 py/torch_tensorrt/hf/exporters/specs/groot.py create mode 100644 py/torch_tensorrt/hf/exporters/specs/nemotron.py create mode 100644 py/torch_tensorrt/hf/exporters/specs/pi05.py create mode 100644 tests/py/dynamo/hf/test_edge_exporter.py diff --git a/examples/dynamo/run_groot_export.py b/examples/dynamo/run_groot_export.py new file mode 100644 index 00000000000..91253911047 --- /dev/null +++ b/examples/dynamo/run_groot_export.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Smoke EdgeExporter on GR00T (4 engines: vision, language, action_context, action). + +Pass the LeRobot GrootPolicy, not policy._groot_model — prepare_sample_inputs +needs GrootEagleEncodeStep / embodiment_id from the policy wrapper. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ +_TRT_PY = _REPO_ROOT / "py" +_TEST = Path("/home/micwilliams/workspace/Test") + +_test = str(_TEST) +while _test in sys.path: + sys.path.remove(_test) +sys.path.insert(0, _test) + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 + +_src_pkg = str(_TRT_PY / "torch_tensorrt") +if _src_pkg not in list(torch_tensorrt.__path__): + torch_tensorrt.__path__.append(_src_pkg) + +from lerobot.configs import FeatureType, PolicyFeature +from lerobot.policies.groot import GrootPolicy +from lerobot.policies.groot.configuration_groot import GrootConfig +from lerobot.utils.constants import ACTION, OBS_STATE +from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter +from trt.plugin.plugin_utils import load_plugins_for_trt +from trt.utils import configure_thor_pytorch, force_hf_attention + + +def load_groot(device: torch.device) -> GrootPolicy: + config = GrootConfig( + base_model_path="nvidia/GR00T-N1.5-3B", + device=str(device), + embodiment_tag="new_embodiment", + chunk_size=50, + n_action_steps=50, + max_state_dim=64, + max_action_dim=32, + image_size=(224, 224), + tokenizer_assets_repo="lerobot/eagle2hg-processor-groot-n1p5", + input_features={ + "observation.images.image": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + "observation.images.image2": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,))}, + ) + return GrootPolicy(config).to(device).eval() + + +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() + + configure_thor_pytorch() + load_plugins_for_trt() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float16 + + policy = load_groot(device) + model = policy._groot_model.to(device=device, dtype=dtype).eval() + eagle = model.backbone.eagle_model + force_hf_attention(eagle.vision_model, "eager") + force_hf_attention(eagle.language_model, "eager") + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="groot", + engine_dir=args.engine_dir, + max_seq_len=968, + dryrun=not args.compile, + skip_runtime_export=False, + ) + + # Spec tokenizes libero via Eagle chat template because we pass the policy. + sample_inputs = {"device": device, "dtype": dtype} + program = exporter.export(policy, sample_inputs, config=config) + + print("engines:", exporter.engines) + print("saved:", exporter.save_engines()) + print("runtime keys:", sorted(exporter.sample)) + + with torch.no_grad(): + if hasattr(program, "module"): + velocity = program.module()(**exporter.sample) + else: + velocity = program(**exporter.sample) + + out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity + print("velocity", tuple(out.shape), "mean", float(out.float().mean())) + + +if __name__ == "__main__": + main() diff --git a/examples/dynamo/run_nemotron_export.py b/examples/dynamo/run_nemotron_export.py new file mode 100644 index 00000000000..c5b59578571 --- /dev/null +++ b/examples/dynamo/run_nemotron_export.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Smoke EdgeExporter on Nemotron-H (one language engine: attn + mamba + MoE). + +Pass the HF causal LM. Collation is tokenizer → input_ids; the spec embeds and +pads to max_seq_len. apply_mamba_stub() must run before from_pretrained. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ +_TRT_PY = _REPO_ROOT / "py" +_TEST = Path("/home/micwilliams/workspace/Test") +_NEMOTRON = _TEST / "nemotron" + +for p in (_NEMOTRON, _TEST): + s = str(p) + while s in sys.path: + sys.path.remove(s) + sys.path.insert(0, s) + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 + +_src_pkg = str(_TRT_PY / "torch_tensorrt") +if _src_pkg not in list(torch_tensorrt.__path__): + torch_tensorrt.__path__.append(_src_pkg) + +from mamba_stub import apply as apply_mamba_stub +from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter +from transformers import AutoModelForCausalLM, AutoTokenizer +from trt.plugin.plugin_utils import load_plugins_for_trt +from trt.utils import configure_thor_pytorch + + +def load_nemotron(checkpoint: str, device: torch.device, dtype: torch.dtype): + apply_mamba_stub() + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + torch_dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + return model, tokenizer + + +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", + ) + parser.add_argument("--prompt", default="Hello.") + parser.add_argument("--max-seq-len", type=int, default=128) + args = parser.parse_args() + + configure_thor_pytorch() + load_plugins_for_trt() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float16 + + model, tokenizer = load_nemotron(args.checkpoint, device, dtype) + encoded = tokenizer(args.prompt, return_tensors="pt") + sample_inputs = { + "input_ids": encoded["input_ids"].to(device), + "attention_mask": encoded["attention_mask"].to(device), + } + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="nemotron_h", + engine_dir=args.engine_dir, + max_seq_len=args.max_seq_len, + dryrun=not args.compile, + skip_runtime_export=False, + ) + program = exporter.export(model, sample_inputs, config=config) + + print("engines:", exporter.engines) + print("saved:", exporter.save_engines()) + print("runtime keys:", sorted(exporter.sample)) + + with torch.no_grad(): + if hasattr(program, "module"): + out = program.module()(**exporter.sample) + else: + out = program(**exporter.sample) + + logits = out[0] if isinstance(out, (tuple, list)) else out + print("logits", tuple(logits.shape), "mean", float(logits.float().mean())) + + +if __name__ == "__main__": + main() diff --git a/examples/dynamo/run_pi05_export.py b/examples/dynamo/run_pi05_export.py new file mode 100644 index 00000000000..f7a25d4c53c --- /dev/null +++ b/examples/dynamo/run_pi05_export.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ +_TRT_PY = _REPO_ROOT / "py" +_TEST = Path("/home/micwilliams/workspace/Test") + +# ``trt.*`` lives in the Test tree. Always move it to the front. +_test = str(_TEST) +while _test in sys.path: + sys.path.remove(_test) +sys.path.insert(0, _test) + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 + +_src_pkg = str(_TRT_PY / "torch_tensorrt") +if _src_pkg not in list(torch_tensorrt.__path__): + torch_tensorrt.__path__.append(_src_pkg) + +from lerobot.configs import FeatureType, PolicyFeature +from lerobot.policies.pi05 import PI05Policy +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE +from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter +from trt.plugin.plugin_utils import load_plugins_for_trt +from trt.utils import configure_thor_pytorch, force_hf_attention + + +def load_pi05(device: torch.device) -> PI05Policy: + policy = PI05Policy.from_pretrained("lerobot/pi05_libero_base").eval() + cfg = policy.config + cfg.device = str(device) + cfg.chunk_size = 50 + cfg.n_action_steps = 50 + cfg.max_state_dim = 32 + cfg.max_action_dim = 32 + cfg.input_features = { + f"{OBS_IMAGES}.image": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image2": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image3": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image4": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(32,)), + } + cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,))} + cfg.empty_cameras = 0 + cfg.validate_features() + return policy + + +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() + + configure_thor_pytorch() + 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") + force_hf_attention(paligemma.language_model, "eager") + force_hf_attention(policy.model.paligemma_with_expert.gemma_expert.model, "eager") + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="pi05", # optional; inferred from paligemma_with_expert + engine_dir=args.engine_dir, + max_seq_len=968, + dryrun=not args.compile, # True = no TRT, still writes config.json + runtime graph + skip_runtime_export=False, # False = also torch.export the stitched execute_engine graph + # components=("vision",), # uncomment to export only vision + ) + + # 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) + print("saved:", exporter.save_engines()) + + # Runtime kwargs are tensors only (pixel_values, lang_embeds, rope, KVs, …). + runtime_kwargs = exporter.sample + print("runtime keys:", sorted(runtime_kwargs)) + + with torch.no_grad(): + if hasattr(program, "module"): + velocity = program.module()(**runtime_kwargs) + else: + velocity = program(**runtime_kwargs) + + out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity + print("velocity", tuple(out.shape), "mean", float(out.float().mean())) + + +if __name__ == "__main__": + main() diff --git a/py/torch_tensorrt/hf/__init__.py b/py/torch_tensorrt/hf/__init__.py new file mode 100644 index 00000000000..1a52009e770 --- /dev/null +++ b/py/torch_tensorrt/hf/__init__.py @@ -0,0 +1,4 @@ +"""HuggingFace-facing export helpers for Torch-TensorRT. + +Use ``from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig``. +""" diff --git a/py/torch_tensorrt/hf/exporters/__init__.py b/py/torch_tensorrt/hf/exporters/__init__.py new file mode 100644 index 00000000000..5d763b03cb8 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/__init__.py @@ -0,0 +1,20 @@ +from torch_tensorrt.hf.exporters.config import EdgeConfig +from torch_tensorrt.hf.exporters.exporter import EdgeExporter +from torch_tensorrt.hf.exporters.spec import ( + ComponentBundle, + EdgeSpec, + get_edge_spec, + register_edge_spec, +) +from torch_tensorrt.hf.exporters.specs import groot as _groot # noqa: F401 +from torch_tensorrt.hf.exporters.specs import nemotron as _nemotron # noqa: F401 +from torch_tensorrt.hf.exporters.specs import pi05 as _pi05 # noqa: F401 + +__all__ = [ + "ComponentBundle", + "EdgeConfig", + "EdgeExporter", + "EdgeSpec", + "get_edge_spec", + "register_edge_spec", +] diff --git a/py/torch_tensorrt/hf/exporters/compile.py b/py/torch_tensorrt/hf/exporters/compile.py new file mode 100644 index 00000000000..8430729dec2 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/compile.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.ops import _as_tuple, record_engine +from torch_tensorrt.hf.exporters.spec import ComponentBundle +from trt.plugin.plugin_utils import restore_attention + +DEFAULT_TRT_SETTINGS: dict[str, Any] = { + "min_block_size": 1, + "require_full_compilation": True, + "immutable_weights": True, + "disable_tf32": 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( + module: nn.Module, + bundle: ComponentBundle, + *, + name: str, + engine_dir: Path, + dryrun: bool = False, + trt_settings: dict[str, Any] | None = None, +) -> tuple[str, tuple[torch.Tensor, ...]]: + """Export one wrapper, 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. + + ``dryrun`` records the patched eager module for ``execute_engine`` and + leaves the patch in place. A real compile restores HF attention after the + TensorRT module is recorded. + """ + module = 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) + + 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 + } + import torch_tensorrt + + compiled = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=trace_args, + **settings, + ) + 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, + ) + ) + (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: + restore_attention(patched) + + +def _write_sidecar( + out_dir: Path, + bundle: ComponentBundle, + name: str, + outputs: tuple[torch.Tensor, ...], + *, + engine_file: str | None = None, + dryrun: bool = False, +) -> 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), + "dryrun": dryrun, + "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") diff --git a/py/torch_tensorrt/hf/exporters/config.py b/py/torch_tensorrt/hf/exporters/config.py new file mode 100644 index 00000000000..ae9f140366e --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/config.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass +class EdgeConfig: + """Knobs for :class:`~torch_tensorrt.hf.exporters.EdgeExporter`. + + ``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 + 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 + 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/py/torch_tensorrt/hf/exporters/exporter.py b/py/torch_tensorrt/hf/exporters/exporter.py new file mode 100644 index 00000000000..8acad2d5074 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/exporter.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import copy +import inspect +import logging +from collections.abc import MutableMapping +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from torch.export import ExportedProgram +from torch_tensorrt.hf.exporters import ops as _ops # noqa: F401 +from torch_tensorrt.hf.exporters.compile import compile_component +from torch_tensorrt.hf.exporters.config import EdgeConfig +from torch_tensorrt.hf.exporters.runtime import EdgeRuntimeModule +from torch_tensorrt.hf.exporters.spec import get_edge_spec +from transformers.exporters.exporter_dynamo import DynamoExporter + +logger = logging.getLogger(__name__) + + +def _clone_export_kwargs(sample_inputs: MutableMapping[str, Any]) -> dict[str, Any]: + """Copy example kwargs into graph leaves. + + Vision/language packing produces intermediate tensors. ``copy.deepcopy`` + refuses those (``Only Tensors created explicitly by the user...``). + """ + cloned: dict[str, Any] = {} + for key, value in dict(sample_inputs).items(): + if isinstance(value, torch.Tensor): + cloned[key] = value.detach().contiguous().clone() + else: + cloned[key] = copy.deepcopy(value) + return cloned + + +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] = {} + + def export( + self, + model: nn.Module, + sample_inputs: MutableMapping[str, Any], + config: EdgeConfig | dict[str, Any], + ) -> ExportedProgram | EdgeRuntimeModule: + if isinstance(config, dict): + config = EdgeConfig(**config) + elif not isinstance(config, EdgeConfig): + raise TypeError(f"Expected EdgeConfig or dict, got {type(config)}") + + 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") + + sample = spec.prepare_sample_inputs(model, sample_inputs, config) + + engine_dir = Path(config.engine_dir or "edge_engines") + engine_dir.mkdir(parents=True, exist_ok=True) + + engines: dict[str, str] = {} + upstream: dict[str, Any] = {} + + for name in names: + module = spec.wrap(name, model, sample, config) + bundle = spec.prepare(name, model, sample, upstream, config, module) + engines[name], outs = compile_component( + module, + bundle, + name=name, + engine_dir=engine_dir, + dryrun=config.dryrun, + trt_settings=config.trt_settings, + ) + upstream.update(spec.capture_upstream(name, outs, sample, bundle)) + + 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 + 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) + + with ( + reset_model_state(model), + patch_forward_signature(model, sample_inputs), + ): + return torch.export.export( + model, + args=(), + kwargs=_clone_export_kwargs(sample_inputs), + strict=config.strict, + dynamic_shapes=dynamic_shapes, + prefer_deferred_runtime_asserts_over_guards=( + config.prefer_deferred_runtime_asserts_over_guards + ), + ) + + def save_engines(self, out_dir: str | Path | None = None) -> dict[str, Path]: + if not self.engines: + raise RuntimeError("save_engines() requires export() first") + if out_dir is None: + return {name: Path(path) for name, path in self.engines.items()} + import shutil + + dest = Path(out_dir) + dest.mkdir(parents=True, exist_ok=True) + written: dict[str, Path] = {} + for name, path in self.engines.items(): + target = dest / name + if Path(path).resolve() != target.resolve(): + if target.exists(): + shutil.rmtree(target) + shutil.copytree(path, target) + written[name] = target + return written diff --git a/py/torch_tensorrt/hf/exporters/ops.py b/py/torch_tensorrt/hf/exporters/ops.py new file mode 100644 index 00000000000..7623c8b43ee --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/ops.py @@ -0,0 +1,132 @@ +"""Named Edge-LLM runtime ops. Engines are values; behavior is the operator.""" + +from __future__ import annotations + +from typing import Any + +import torch + +_ENGINE_META: dict[str, dict[str, Any]] = {} +_COMPILED_MODULES: dict[str, torch.nn.Module] = {} + + +def record_engine( + path: str, + *, + component: str, + input_names: list[str], + outputs: tuple[torch.Tensor, ...], + module: torch.nn.Module | None = None, +) -> None: + _ENGINE_META[path] = { + "component": component, + "input_names": list(input_names), + "output_shapes": [tuple(t.shape) for t in outputs], + "output_dtypes": [t.dtype for t in outputs], + } + if module is not None: + _COMPILED_MODULES[path] = module + + +def _as_tuple(value: Any) -> tuple[torch.Tensor, ...]: + if isinstance(value, tuple): + return tuple(value) + if isinstance(value, list): + return tuple(value) + return (value,) + + +@torch.library.custom_op("edge_llm::execute_engine", mutates_args=()) # type: ignore[misc] +def execute_engine( + engine_path: str, component: str, tensors: list[torch.Tensor] +) -> list[torch.Tensor]: + compiled = _COMPILED_MODULES.get(engine_path) + if compiled is not None: + out = compiled(*tensors) + return list(_as_tuple(out)) + raise RuntimeError( + f"No in-process module for engine {engine_path!r} ({component}). " + "Compile first, or load a serialized engine at this path." + ) + + +@execute_engine.register_fake # type: ignore[misc] +def _( + engine_path: str, component: str, tensors: list[torch.Tensor] +) -> list[torch.Tensor]: + meta = _ENGINE_META.get(engine_path) + device = tensors[0].device if tensors else torch.device("cpu") + if meta is None: + return [torch.empty_like(t) for t in tensors] + return [ + torch.empty(shape, dtype=dtype, device=device) + for shape, dtype in zip(meta["output_shapes"], meta["output_dtypes"]) + ] + + +def call_engine( + engine_path: str, component: str, *tensors: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """Python helper so specs can pass ``*tensors`` instead of a list.""" + out = torch.ops.edge_llm.execute_engine.default( + engine_path, component, list(tensors) + ) + return tuple(out) + + +@torch.library.custom_op("edge_llm::fuse_prefix", mutates_args=()) # type: ignore[misc] +def fuse_prefix( + vision_tokens: torch.Tensor, + lang_embeds: torch.Tensor, + compact_index: torch.Tensor, +) -> torch.Tensor: + hidden = lang_embeds.shape[-1] + batch = lang_embeds.shape[0] + vis = vision_tokens + if vis.ndim == 2: + vis = vis.reshape(batch, -1, hidden) + embs = torch.cat([vis, lang_embeds], dim=1) + index = compact_index.to(dtype=torch.long) + return torch.gather(embs, 1, index.unsqueeze(-1).expand(-1, -1, hidden)) + + +@fuse_prefix.register_fake # type: ignore[misc] +def _( + vision_tokens: torch.Tensor, + lang_embeds: torch.Tensor, + compact_index: torch.Tensor, +) -> torch.Tensor: + batch, compact_len = compact_index.shape + return torch.empty( + batch, + compact_len, + lang_embeds.shape[-1], + dtype=lang_embeds.dtype, + device=lang_embeds.device, + ) + + +@torch.library.custom_op("edge_llm::scatter_image_tokens", mutates_args=()) # type: ignore[misc] +def scatter_image_tokens( + vision_tokens: torch.Tensor, + lang_embeds: torch.Tensor, + image_token_mask: torch.Tensor, +) -> torch.Tensor: + hidden = lang_embeds.shape[-1] + vis = vision_tokens.reshape(-1, hidden).to(dtype=lang_embeds.dtype) + out = lang_embeds.clone() + flat = out.reshape(-1, hidden) + mask = image_token_mask.reshape(-1).to(dtype=torch.bool) + n = int(mask.sum().item()) + if n: + flat[mask] = vis[:n] + return flat.reshape_as(lang_embeds) + + +@scatter_image_tokens.register_fake # type: ignore[misc] +def _( + vision_tokens: torch.Tensor, + lang_embeds: torch.Tensor, + image_token_mask: torch.Tensor, +) -> torch.Tensor: + return torch.empty_like(lang_embeds) diff --git a/py/torch_tensorrt/hf/exporters/runtime.py b/py/torch_tensorrt/hf/exporters/runtime.py new file mode 100644 index 00000000000..8e98947140a --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/runtime.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch.nn as nn +from torch_tensorrt.hf.exporters.spec import EdgeSpec + + +class EdgeRuntimeModule(nn.Module): # type: ignore[misc] + """Graph dumped by ``EdgeExporter``: packing in Python, compute as execute_engine.""" + + def __init__(self, spec: EdgeSpec, engines: Mapping[str, str]) -> None: + super().__init__() + self.spec = spec + self.engines = dict(engines) + + def forward(self, **sample: Any) -> Any: + return self.spec.run(self.engines, sample) diff --git a/py/torch_tensorrt/hf/exporters/spec.py b/py/torch_tensorrt/hf/exporters/spec.py new file mode 100644 index 00000000000..98f94a9e406 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/spec.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping, MutableMapping +from dataclasses import dataclass, field +from typing import Any + +import torch.nn as nn + +_SPECS: dict[str, type[EdgeSpec]] = {} + + +def register_edge_spec(*model_types: str) -> Callable[[type[EdgeSpec]], type[EdgeSpec]]: + """Register an :class:`EdgeSpec` for one or more ``config.model_type`` keys.""" + + def decorator(cls: type[EdgeSpec]) -> type[EdgeSpec]: + for model_type in model_types: + _SPECS[model_type] = cls + return cls + + return decorator + + +def registered_specs() -> dict[str, type[EdgeSpec]]: + return dict(_SPECS) + + +@dataclass +class ComponentBundle: + """Everything :func:`compile_component` needs to build one engine.""" + + trace_args: tuple[Any, ...] + save_args: tuple[Any, ...] + input_names: list[str] + output_names: list[str] + input_specs: Any = 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 + execute_args: tuple[Any, ...] | None = None + model_type: str = "edge" + engine_file: str = "engine.engine" + + +class EdgeSpec(ABC): + """Per-family wrap / flatten / runtime wiring. + + ``EdgeExporter.export`` never branches on PI05 vs Nemotron. It only loops + ``spec.components``. + """ + + components: tuple[str, ...] = () + + @abstractmethod + def prepare_sample_inputs( + self, + model: nn.Module, + raw: Mapping[str, Any], + config: Any, + ) -> MutableMapping[str, Any]: + """Caller payload → stem dict used by wrap/prepare/run.""" + + @abstractmethod + def wrap( + self, + name: str, + model: nn.Module, + sample: Mapping[str, Any], + config: Any, + ) -> nn.Module: + """Replace a submodule with an export wrapper (not an HF forward patch).""" + + @abstractmethod + def prepare( + self, + name: str, + model: nn.Module, + sample: MutableMapping[str, Any], + upstream: Mapping[str, Any], + config: Any, + module: nn.Module, + ) -> ComponentBundle: + """Build the trace/save tuple for ``module``.""" + + def capture_upstream( + 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 {} + + @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`.""" + return { + key: value + for key, value in sample.items() + if hasattr(value, "dtype") and hasattr(value, "device") + } + + +def infer_model_type(model: nn.Module, explicit: str | None = None) -> str: + if explicit: + return explicit + config = getattr(model, "config", None) + model_type = getattr(config, "model_type", None) + if isinstance(model_type, str) and model_type in _SPECS: + return model_type + if hasattr(model, "paligemma_with_expert") or hasattr( + getattr(model, "model", None), "paligemma_with_expert" + ): + return "pi05" + if hasattr(model, "_groot_model") or ( + getattr(getattr(model, "backbone", None), "eagle_model", None) is not None + ): + return "groot" + name = type(model).__name__.lower() + if "nemotron" in name: + return "nemotron_h" + if isinstance(model_type, str): + return model_type + raise KeyError( + f"No EdgeSpec for {type(model).__name__}. " + f"Pass EdgeConfig(model_type=...) or register one. " + f"Known: {sorted(_SPECS)}" + ) + + +def get_edge_spec(model: nn.Module, model_type: str | None = None) -> EdgeSpec: + key = infer_model_type(model, model_type) + if key not in _SPECS: + raise KeyError(f"No EdgeSpec registered for {key!r}. Known: {sorted(_SPECS)}") + return _SPECS[key]() diff --git a/py/torch_tensorrt/hf/exporters/specs/__init__.py b/py/torch_tensorrt/hf/exporters/specs/__init__.py new file mode 100644 index 00000000000..ed8ebec6682 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/specs/__init__.py @@ -0,0 +1,5 @@ +"""Family specs. Import is a side effect: ``@register_edge_spec`` attaches.""" + +from torch_tensorrt.hf.exporters.specs import groot, nemotron, pi05 + +__all__ = ["groot", "nemotron", "pi05"] diff --git a/py/torch_tensorrt/hf/exporters/specs/_language.py b/py/torch_tensorrt/hf/exporters/specs/_language.py new file mode 100644 index 00000000000..25a8d86154b --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/specs/_language.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +import torch.nn as nn + + +def causal_lm_flat( + language: nn.Module, + inputs_embeds: torch.Tensor, + *, + max_seq_len: int, + device: torch.device, + 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.""" + decoder = getattr(language, "model", language) + cfg = language.config + bsz, prompt_len, hidden = inputs_embeds.shape + seq_len = int(seq_len or prompt_len) + num_layers = len(decoder.layers) + num_kv = int(cfg.num_key_value_heads) + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + try: + from trt.rope import make_rope_rotary_cos_sin + + rope = make_rope_rotary_cos_sin( + cfg, int(max_seq_len), device, language_model=language + ) + except ImportError: + rope = torch.zeros(int(max_seq_len), 2, 1, head_dim, device=device, dtype=dtype) + + 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) + kvs = [ + torch.zeros( + bsz, 2, num_kv, int(max_seq_len), head_dim, device=device, dtype=dtype + ) + for _ in range(num_layers) + ] + flat = (inputs_embeds, rope, ctx_len, kv_start, last_token_ids, ds_stack, *kvs) + names = [ + "inputs_embeds", + "rope_rotary_cos_sin", + "context_lengths", + "kvcache_start_index", + "last_token_ids", + "ds_stack", + *[f"past_key_values_{i}" for i in range(num_layers)], + ] + return flat, { + "input_names": names, + "num_layers": num_layers, + "hidden_size": hidden, + "head_dim": head_dim, + "num_key_value_heads": num_kv, + } + + +def kv_kwargs( + sample: Mapping[str, Any], prefix: str = "past_key_values_" +) -> list[torch.Tensor]: + tensors = [] + idx = 0 + while f"{prefix}{idx}" in sample: + tensors.append(sample[f"{prefix}{idx}"]) + idx += 1 + return tensors + + +def split_flat_to_kwargs( + flat: tuple[torch.Tensor, ...], names: list[str] +) -> dict[str, torch.Tensor]: + return dict(zip(names, flat)) diff --git a/py/torch_tensorrt/hf/exporters/specs/groot.py b/py/torch_tensorrt/hf/exporters/specs/groot.py new file mode 100644 index 00000000000..a9c7d00d7ae --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/specs/groot.py @@ -0,0 +1,347 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.ops import call_engine, scatter_image_tokens +from torch_tensorrt.hf.exporters.spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from torch_tensorrt.hf.exporters.specs._language import ( + causal_lm_flat, + kv_kwargs, + split_flat_to_kwargs, +) + + +def _groot(model: nn.Module) -> nn.Module: + if hasattr(model, "_groot_model"): + return model._groot_model + backbone = getattr(model, "backbone", None) + if backbone is not None and hasattr(backbone, "eagle_model"): + return model + raise RuntimeError( + "GR00T spec expected GrootPolicy or a module with backbone.eagle_model" + ) + + +@register_edge_spec("groot", "gr00t") +class GrootSpec(EdgeSpec): + components = ("vision", "language", "action_context", "action") + + def prepare_sample_inputs( + self, model: nn.Module, raw: Mapping[str, Any], config: Any + ) -> MutableMapping[str, Any]: + if "pixel_values" in raw and "input_ids" in raw: + return dict(raw) + + from lerobot.policies.factory import make_pre_post_processors + from lerobot.policies.groot.processor_groot import GrootEagleEncodeStep + from trt.data import create_pil_messages, load_test_data, pack_state + from trt.executor.models.groot.helpers import make_embodiment_id + + policy = model + device = raw.get( + "device", torch.device("cuda" if torch.cuda.is_available() else "cpu") + ) + dtype = raw.get("dtype", torch.float16) + cfg = getattr(policy, "config", None) + pre_processor, _ = make_pre_post_processors( + cfg, + None, + preprocessor_overrides={"device_processor": {"device": str(device)}}, + ) + eagle_step = next( + s for s in pre_processor.steps if isinstance(s, GrootEagleEncodeStep) + ) + proc = eagle_step.proc + data = raw.get("data") or load_test_data( + raw.get("dataset_id", "lerobot/libero"), episode_index=0, frame_index=0 + ) + messages = create_pil_messages(data) + text = proc.apply_chat_template( + messages, tokenize=False, **{"add_generation_prompt": True} + ) + image_inputs, video_inputs = proc.process_vision_info(messages) + tokenized = proc( + text=[text], + images=image_inputs, + videos=video_inputs, + return_tensors="pt", + padding=True, + **{ + "images_kwargs": { + "min_dynamic_tiles": 1, + "max_dynamic_tiles": 1, + "use_thumbnail": False, + } + }, + ) + state = ( + pack_state( + data["state"], + max_state_dim=int(getattr(cfg, "max_state_dim", 64)), + device=device, + ) + .to(device=device, dtype=dtype) + .contiguous() + ) + return { + "pixel_values": tokenized["pixel_values"].to(device=device, dtype=dtype), + "input_ids": tokenized["input_ids"].to(device=device, dtype=torch.long), + "attention_mask": tokenized["attention_mask"].to( + device=device, dtype=torch.long + ), + "state": state, + "embodiment_id": make_embodiment_id(policy, state, device, torch.long), + } + + def wrap( + self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any + ) -> nn.Module: + from trt.modules.export.diffusion import ( + GrootDiTStepEncoderExportModule, + StaticActionVelocityStepExportModule, + TRTDynamicCategorySpecificMLPExportModule, + ) + from trt.modules.export.language import ( + CausalLMExportModule, + ContextProjectionExportModule, + language_decoder, + ) + from trt.modules.export.vision import GridVisionExportModule + + found = _groot(model) + eagle = found.backbone.eagle_model + device = sample["pixel_values"].device + dtype = sample["pixel_values"].dtype + # GR00T loads DiT / category MLPs in bf16. Export tensors are fp16. + # The e2e path casts the whole wrapper; without that, sa_embs is + # float32 against bf16 attn.to_q weights. + if name == "vision": + module = GridVisionExportModule( + vision_model=eagle.vision_model, + projector=eagle.mlp1, + sample_pixel_values=sample["pixel_values"], + select_layer=eagle.select_layer, + pixel_shuffle=eagle.use_pixel_shuffle, + downsample_ratio=getattr(eagle, "downsample_ratio", 0.5), + vision_kwargs={}, + ) + elif name == "language": + language = eagle.language_model + module = CausalLMExportModule( + language_decoder(language), language.lm_head, select_layer=-1 + ) + elif name == "action_context": + module = ContextProjectionExportModule( + found.backbone.eagle_linear, + found.action_head.vlln, + found.action_head.vl_self_attention, + ) + elif name == "action": + module = StaticActionVelocityStepExportModule( + step_encoder=GrootDiTStepEncoderExportModule( + found.action_head, sample.get("embodiment_id") + ), + action_expert=found.action_head.model, + velocity_decoder=TRTDynamicCategorySpecificMLPExportModule( + found.action_head.action_decoder + ), + output_tokens=int(found.action_head.config.action_horizon), + cast_hidden_fp32=False, + ) + else: + raise KeyError(name) + return module.eval().to(device=device, dtype=dtype) + + def prepare( + self, + name: str, + model: nn.Module, + sample: MutableMapping[str, Any], + upstream: Mapping[str, Any], + config: Any, + module: nn.Module, + ) -> ComponentBundle: + from trt.plugin.attention import ContextAttentionMaskType + from trt.plugin.plugin_utils import ( + patch_language_attention, + patch_vision_attention, + ) + + found = _groot(model) + eagle = found.backbone.eagle_model + device = sample["pixel_values"].device + dtype = sample["pixel_values"].dtype + + if name == "vision": + px = sample["pixel_values"] + seq_len = int(getattr(module, "seq_len", 1) or 1) + batch = int(getattr(module, "batch_size", 1) or 1) + + def _patch(mod: nn.Module) -> Any: + vision = getattr(mod.vision_model, "vision_model", mod.vision_model) + return patch_vision_attention( + vision, batch_size=batch, seq_len=seq_len, name="SigLIP" + ) + + return ComponentBundle( + trace_args=(px,), + save_args=(px,), + input_names=["pixel_values"], + output_names=["visual_embeds"], + patch_fn=_patch, + model_type="vit", + engine_file="visual.engine", + ) + + if name == "language": + language = 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"])) + cfg = language.config + head_dim = int( + getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) + ) + + def _patch(mod: nn.Module) -> Any: + decoder = getattr(mod, "lm", mod) + return patch_language_attention( + decoder, + hidden_size=int(cfg.hidden_size), + num_attention_heads=int(cfg.num_attention_heads), + num_key_value_heads=int(cfg.num_key_value_heads), + head_dim=head_dim, + context_attention_mask_type=ContextAttentionMaskType.CAUSAL, + ) + + return ComponentBundle( + trace_args=packed, + save_args=packed, + input_names=meta["input_names"], + output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], + patch_fn=_patch, + model_type="language", + engine_file="language.engine", + ) + + if name == "action_context": + hidden = upstream["lm_hidden"].to(dtype=dtype) + return ComponentBundle( + trace_args=(hidden,), + save_args=(hidden,), + input_names=["lm_hidden_states"], + output_names=["vl_embs"], + model_type="context_projection", + engine_file="context_projection.engine", + ) + + 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( + 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) + + 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 == "action_context": + ctx = outputs[0] if isinstance(outputs, tuple) else outputs + return {"context_embs": ctx} + return {} + + def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: + vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] + embeds = scatter_image_tokens( + vis, sample["lang_embeds"], sample["image_token_mask"] + ) + lm = call_engine( + engines["language"], + "language", + embeds, + sample["rope_rotary_cos_sin"], + sample["context_lengths"], + sample["kvcache_start_index"], + sample["last_token_ids"], + sample["ds_stack"], + *kv_kwargs(sample), + ) + ctx = call_engine(engines["action_context"], "action_context", lm[1])[0] + return call_engine( + engines["action"], + "action", + sample["step_actions"], + sample["step_timestep"], + ctx, + sample["state"], + sample["embodiment_id"], + ) diff --git a/py/torch_tensorrt/hf/exporters/specs/nemotron.py b/py/torch_tensorrt/hf/exporters/specs/nemotron.py new file mode 100644 index 00000000000..8fc8090b080 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/specs/nemotron.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.ops import call_engine +from torch_tensorrt.hf.exporters.spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from torch_tensorrt.hf.exporters.specs._language import kv_kwargs, split_flat_to_kwargs + + +def _decoder(model: nn.Module) -> nn.Module: + return getattr(model, "backbone", None) or model.model + + +def _kind(mixer: nn.Module) -> str: + name = type(mixer).__name__ + if "Mamba" in name: + return "mamba" + if "Attention" in name: + return "attention" + if "MoE" in name or "Moe" in name: + return "moe" + return "mlp" + + +class NemotronExportModule(nn.Module): # type: ignore[misc] + """Hybrid decoder: plugin attention / mamba / moe, native MLP.""" + + def __init__(self, model: nn.Module): + super().__init__() + decoder = _decoder(model) + self.layers = decoder.layers + self.norm = decoder.norm_f + self.lm_head = model.lm_head + self.kinds = [_kind(block.mixer) for block in self.layers] + self.num_attn = self.kinds.count("attention") + self.num_mamba = self.kinds.count("mamba") + + def forward( + self, + inputs_embeds: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + context_lengths: torch.Tensor, + kvcache_start_index: torch.Tensor, + last_token_ids: torch.Tensor, + *states: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + from trt.modules.export.language import gather_last_token_hidden + + na, nm = self.num_attn, self.num_mamba + kvs = list(states[:na]) + convs = list(states[na : na + nm]) + ssms = list(states[na + nm : na + 2 * nm]) + kv_i = conv_i = 0 + hidden = inputs_embeds + present_kv, present_conv, present_ssm = [], [], [] + for block, kind in zip(self.layers, self.kinds): + residual = hidden + hidden = block.norm(hidden) + mixer = block.mixer + if kind == "attention": + hidden, kv = mixer( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=kvs[kv_i], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, + ) + present_kv.append(kv) + kv_i += 1 + elif kind == "mamba": + hidden, conv_out, ssm_out = mixer( + hidden, convs[conv_i], ssms[conv_i], context_lengths + ) + present_conv.append(conv_out) + present_ssm.append(ssm_out) + conv_i += 1 + else: + hidden = mixer(hidden) + hidden = residual + hidden + hidden = self.norm(hidden) + last = gather_last_token_hidden(hidden, last_token_ids) + logits = self.lm_head(last).float() + return (logits, *present_kv, *present_conv, *present_ssm) + + +def allocate_plugin_states( + model: nn.Module, + config: Any, + batch: int, + max_seq_len: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + kinds = [_kind(block.mixer) for block in _decoder(model).layers] + head_dim = int( + getattr(config, "head_dim", 0) + or config.hidden_size // config.num_attention_heads + ) + conv_dim = int(config.mamba_num_heads) * int(config.mamba_head_dim) + 2 * int( + config.n_groups + ) * int(config.ssm_state_size) + conv_kernel = int(getattr(config, "conv_kernel", 4)) + kvs, convs, ssms = [], [], [] + for kind in kinds: + if kind == "attention": + kvs.append( + torch.zeros( + batch, + 2, + int(config.num_key_value_heads), + max_seq_len, + head_dim, + device=device, + dtype=dtype, + ) + ) + elif kind == "mamba": + convs.append( + torch.zeros(batch, conv_dim, conv_kernel, device=device, dtype=dtype) + ) + ssms.append( + torch.zeros( + batch, + int(config.mamba_num_heads), + int(config.mamba_head_dim), + int(config.ssm_state_size), + device=device, + dtype=dtype, + ) + ) + return kvs, convs, ssms + + +@register_edge_spec("nemotron_h", "nemotron") +class NemotronSpec(EdgeSpec): + components = ("language",) + + def prepare_sample_inputs( + self, model: nn.Module, raw: Mapping[str, Any], config: Any + ) -> MutableMapping[str, Any]: + if "inputs_embeds" in raw: + return dict(raw) + device = raw["input_ids"].device + embeddings = model.get_input_embeddings()(raw["input_ids"]) + mask = raw.get("attention_mask") + max_seq_len = int(config.max_seq_len) + bsz, prompt_len, hidden = embeddings.shape + if prompt_len < max_seq_len: + pad = max_seq_len - prompt_len + embeddings = torch.cat( + [ + embeddings, + torch.zeros( + bsz, pad, hidden, device=device, dtype=embeddings.dtype + ), + ], + dim=1, + ) + if mask is not None: + mask = torch.cat( + [mask, torch.ones(bsz, pad, device=device, dtype=mask.dtype)], + dim=1, + ) + return { + "inputs_embeds": embeddings, + "attention_mask": mask, + "bsz": embeddings.shape[0], + } + + def wrap( + self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any + ) -> nn.Module: + from trt.plugin.moe import PluginNemotronMoE + from trt.plugin.plugin_utils import patch_nemotron_mixers + + patch_nemotron_mixers(model, model.config) + for block in _decoder(model).layers: + if isinstance(block.mixer, PluginNemotronMoE): + block.mixer.prepare_for_export() + return NemotronExportModule(model).eval() + + def prepare( + self, + name: str, + model: nn.Module, + sample: MutableMapping[str, Any], + upstream: Mapping[str, Any], + config: Any, + module: nn.Module, + ) -> ComponentBundle: + from trt.rope import make_rope_rotary_cos_sin + + embeds = sample["inputs_embeds"] + device, dtype = embeds.device, embeds.dtype + bsz, seq_len, _ = embeds.shape + rope = make_rope_rotary_cos_sin( + model.config, + int(config.max_seq_len), + device, + language_model=_decoder(model), + ) + 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) + kvs, convs, ssms = allocate_plugin_states( + model, model.config, bsz, int(config.max_seq_len), device, dtype + ) + flat = (embeds, rope, ctx_len, kv_start, last_token_ids, *kvs, *convs, *ssms) + na, nm = module.num_attn, module.num_mamba + names = [ + "inputs_embeds", + "rope_rotary_cos_sin", + "context_lengths", + "kvcache_start_index", + "last_token_ids", + *[f"past_key_values_{i}" for i in range(na)], + *[f"conv_state_{i}" for i in range(nm)], + *[f"ssm_state_{i}" for i in range(nm)], + ] + sample.update(split_flat_to_kwargs(flat, names)) + return ComponentBundle( + 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", + ) + + def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: + leading = [ + sample["inputs_embeds"], + sample["rope_rotary_cos_sin"], + sample["context_lengths"], + sample["kvcache_start_index"], + sample["last_token_ids"], + ] + states = ( + kv_kwargs(sample) + + kv_kwargs(sample, "conv_state_") + + kv_kwargs(sample, "ssm_state_") + ) + return call_engine(engines["language"], "language", *leading, *states) diff --git a/py/torch_tensorrt/hf/exporters/specs/pi05.py b/py/torch_tensorrt/hf/exporters/specs/pi05.py new file mode 100644 index 00000000000..33f013c8c37 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/specs/pi05.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.ops import call_engine, fuse_prefix +from torch_tensorrt.hf.exporters.spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from torch_tensorrt.hf.exporters.specs._language import ( + causal_lm_flat, + kv_kwargs, + split_flat_to_kwargs, +) + + +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: torch.Tensor) -> torch.Tensor: + if pixel_values.ndim != 4: + return pixel_values + if pixel_values.shape[1] in (1, 3, 4) and pixel_values.shape[-1] not in (1, 3, 4): + return pixel_values.permute(0, 2, 3, 1).contiguous() + return pixel_values + + +@register_edge_spec("pi05") +class Pi05Spec(EdgeSpec): + components = ("vision", "language", "action") + + def prepare_sample_inputs( + self, model: nn.Module, raw: Mapping[str, Any], config: Any + ) -> MutableMapping[str, Any]: + if "pixel_values" in raw and "tokens" in raw: + return dict(raw) + + from lerobot.policies.factory import make_pre_post_processors + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + from trt.data import frame_from_test_data, load_test_data + + policy = model if hasattr(model, "_preprocess_images") else None + if policy is None: + raise ValueError( + "PI05 prepare_sample_inputs needs a LeRobot policy or a pre-collated " + "dict with pixel_values/tokens/masks" + ) + device = raw.get("device", next(policy.parameters()).device) + dtype = raw.get("dtype", torch.float16) + data = raw.get("data") or load_test_data( + raw.get("dataset_id", "lerobot/libero"), episode_index=0, frame_index=0 + ) + frame = frame_from_test_data(data, policy, fill_missing=True) + pre_processor, _ = make_pre_post_processors( + policy.config, + None, + preprocessor_overrides={"device_processor": {"device": str(device)}}, + ) + batch = pre_processor(frame) + images, img_masks = policy._preprocess_images(batch) + pixel_values = torch.cat( + [img.to(device=device, dtype=dtype) for img in images], dim=0 + ).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) + lang_embeds = core.paligemma_with_expert.embed_language_tokens(tokens) + return { + "pixel_values": pixel_values, + "images": images, + "img_masks": img_masks, + "tokens": tokens, + "masks": masks, + "lang_embeds": lang_embeds.to(device=device, dtype=dtype).contiguous(), + } + + def wrap( + self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any + ) -> nn.Module: + from trt.modules.export.diffusion import ( + PI05PrefixKVStepEncoderExportModule, + StaticActionVelocityStepExportModule, + ) + from trt.modules.export.language import CausalLMExportModule, language_decoder + from trt.modules.export.vision import GridVisionExportModule + + core = _core(model) + paligemma = core.paligemma_with_expert.paligemma.model + if name == "vision": + px = sample["pixel_values"] + sample_px = px.float() if px.dtype != torch.float32 else px + return GridVisionExportModule( + vision_model=paligemma.vision_tower.float(), + projector=paligemma.multi_modal_projector, + sample_pixel_values=sample_px, + select_layer=-1, + pixel_shuffle=False, + downsample_ratio=0.5, + force_float32_input=True, + ).eval() + if name == "language": + language = paligemma.language_model + lm_head = core.paligemma_with_expert.paligemma.lm_head + return CausalLMExportModule(language_decoder(language), lm_head).eval() + if name == "action": + return StaticActionVelocityStepExportModule( + step_encoder=PI05PrefixKVStepEncoderExportModule(core), + action_expert=core.paligemma_with_expert.gemma_expert.model, + velocity_decoder=core.action_out_proj, + output_tokens=int(core.config.chunk_size), + cast_hidden_fp32=False, + ).eval() + raise KeyError(name) + + def prepare( + self, + name: str, + model: nn.Module, + sample: MutableMapping[str, Any], + upstream: Mapping[str, Any], + config: Any, + module: nn.Module, + ) -> ComponentBundle: + from trt.executor.models.pi05.helpers import ( + build_pi05_prefix_embs, + make_pi05_suffix_position_and_mask, + pi05_compact_index, + ) + from trt.plugin.attention import ContextAttentionMaskType + from trt.plugin.plugin_utils import ( + patch_language_attention, + patch_vision_attention, + ) + + core = _core(model) + device = sample["pixel_values"].device + dtype = sample["pixel_values"].dtype + + if name == "vision": + px = sample["pixel_values"] + seq_len = int(getattr(module, "seq_len", 1) or 1) + batch = int(getattr(module, "batch_size", 1) or 1) + + def _patch(mod: nn.Module) -> Any: + return patch_vision_attention( + mod.vision_model, batch_size=batch, seq_len=seq_len, name="SigLIP" + ) + + return ComponentBundle( + trace_args=(px,), + save_args=(_nchw_to_hwc(px),), + input_names=["pixel_values"], + output_names=["visual_embeds"], + patch_fn=_patch, + model_type="vit", + engine_file="visual.engine", + ) + + 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"], + ) + 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] + ) + sample["compact_index"] = pi05_compact_index( + sample["img_masks"], + sample["images"], + seq_per_image, + sample["masks"], + device, + ) + sample["prefix_pad_mask"] = pad + 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"])) + cfg = language.config + head_dim = int( + getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) + ) + + def _patch(mod: nn.Module) -> Any: + decoder = getattr(mod, "lm", mod) + return patch_language_attention( + decoder, + hidden_size=int(cfg.hidden_size), + num_attention_heads=int(cfg.num_attention_heads), + num_key_value_heads=int(cfg.num_key_value_heads), + head_dim=head_dim, + context_attention_mask_type=ContextAttentionMaskType.PADDING, + ) + + return ComponentBundle( + trace_args=flat, + save_args=flat, + input_names=meta["input_names"], + output_names=["logits", "lm_hidden_states", "prefix_k", "prefix_v"], + patch_fn=_patch, + extra_config={"prefix_pad_mask_len": compact_len}, + model_type="language", + engine_file="language.engine", + ) + + 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, + ) + 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( + core_mod, 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) + return ComponentBundle( + 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", + ) + raise KeyError(name) + + 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 {"prefix_k": outputs[2], "prefix_v": outputs[3]} + return {} + + def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: + vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] + prefix = fuse_prefix(vis, sample["lang_embeds"], sample["compact_index"]) + lm = call_engine( + engines["language"], + "language", + prefix, + sample["rope_rotary_cos_sin"], + sample["context_lengths"], + sample["kvcache_start_index"], + sample["last_token_ids"], + sample["ds_stack"], + *kv_kwargs(sample), + ) + return call_engine( + engines["action"], + "action", + sample["step_actions"], + sample["step_timestep"], + lm[2], + lm[3], + sample["suffix_position_ids"], + sample["suffix_attention_mask"], + ) diff --git a/setup.py b/setup.py index 1c994f6f39b..10269cfe912 100644 --- a/setup.py +++ b/setup.py @@ -616,6 +616,9 @@ def run(self): "torch_tensorrt.dynamo.runtime", "torch_tensorrt.dynamo.tools", "torch_tensorrt.executorch", + "torch_tensorrt.hf", + "torch_tensorrt.hf.exporters", + "torch_tensorrt.hf.exporters.specs", "torch_tensorrt.runtime", ] @@ -655,6 +658,9 @@ def run(self): "torch_tensorrt.dynamo.runtime": "py/torch_tensorrt/dynamo/runtime", "torch_tensorrt.dynamo.tools": "py/torch_tensorrt/dynamo/tools", "torch_tensorrt.executorch": "py/torch_tensorrt/executorch", + "torch_tensorrt.hf": "py/torch_tensorrt/hf", + "torch_tensorrt.hf.exporters": "py/torch_tensorrt/hf/exporters", + "torch_tensorrt.hf.exporters.specs": "py/torch_tensorrt/hf/exporters/specs", "torch_tensorrt.runtime": "py/torch_tensorrt/runtime", } diff --git a/tests/py/dynamo/hf/test_edge_exporter.py b/tests/py/dynamo/hf/test_edge_exporter.py new file mode 100644 index 00000000000..bd9ae2aacbe --- /dev/null +++ b/tests/py/dynamo/hf/test_edge_exporter.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pytest +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter, register_edge_spec +from torch_tensorrt.hf.exporters.ops import call_engine +from torch_tensorrt.hf.exporters.spec import ComponentBundle, EdgeSpec, registered_specs + + +@register_edge_spec("dummy_edge") +class DummySpec(EdgeSpec): + components = ("language",) + + def prepare_sample_inputs(self, model, raw, config): + return {"x": raw["x"]} + + def wrap(self, name, model, sample, config) -> nn.Module: + return model.eval() + + def prepare(self, name, model, sample, upstream, config, module) -> ComponentBundle: + x = sample["x"] + return ComponentBundle( + 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] + + +@pytest.mark.unit +def test_builtin_specs_are_registered(): + keys = registered_specs() + assert "pi05" in keys + assert "groot" in keys + assert "nemotron_h" in keys + assert "dummy_edge" in keys + + +@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): + 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 + 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 wrap(self, name, model, sample, config) -> nn.Module: + return model.eval() + + def prepare(self, name, model, sample, upstream, config, module) -> 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( + 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) From b5fa004adc843cb23d02f53efd939f4a843142f1 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 13:09:34 -0700 Subject: [PATCH 02/11] Vendor Edge-LLM plugins and patches into hf.exporters. Drop the Test/trt import path so PI05, GR00T, and Nemotron export is self-contained, and rename GR00T action_context to context_projection. --- examples/dynamo/run_groot_export.py | 14 +- examples/dynamo/run_nemotron_export.py | 16 +- examples/dynamo/run_pi05_export.py | 13 +- py/torch_tensorrt/hf/exporters/compile.py | 11 +- py/torch_tensorrt/hf/exporters/data.py | 131 ++++ py/torch_tensorrt/hf/exporters/exporter.py | 19 - .../hf/exporters/helpers/__init__.py | 1 + .../hf/exporters/helpers/groot.py | 27 + .../hf/exporters/helpers/pi05.py | 224 +++++++ py/torch_tensorrt/hf/exporters/mamba_stub.py | 88 +++ .../hf/exporters/patches/__init__.py | 1 + .../hf/exporters/patches/diffusion.py | 338 +++++++++++ .../hf/exporters/patches/language.py | 148 +++++ .../hf/exporters/patches/vision.py | 158 +++++ .../hf/exporters/plugin/__init__.py | 1 + .../hf/exporters/plugin/attention.py | 557 ++++++++++++++++++ .../hf/exporters/plugin/mamba.py | 330 +++++++++++ py/torch_tensorrt/hf/exporters/plugin/moe.py | 423 +++++++++++++ .../hf/exporters/plugin/plugin_converter.py | 366 ++++++++++++ .../hf/exporters/plugin/plugin_utils.py | 519 ++++++++++++++++ .../hf/exporters/prefix_cache.py | 308 ++++++++++ py/torch_tensorrt/hf/exporters/rope.py | 307 ++++++++++ .../hf/exporters/specs/_language.py | 2 +- py/torch_tensorrt/hf/exporters/specs/groot.py | 54 +- .../hf/exporters/specs/nemotron.py | 20 +- py/torch_tensorrt/hf/exporters/specs/pi05.py | 38 +- py/torch_tensorrt/hf/exporters/utils.py | 48 ++ pyproject.toml | 12 + setup.py | 6 + 29 files changed, 4074 insertions(+), 106 deletions(-) create mode 100644 py/torch_tensorrt/hf/exporters/data.py create mode 100644 py/torch_tensorrt/hf/exporters/helpers/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/helpers/groot.py create mode 100644 py/torch_tensorrt/hf/exporters/helpers/pi05.py create mode 100644 py/torch_tensorrt/hf/exporters/mamba_stub.py create mode 100644 py/torch_tensorrt/hf/exporters/patches/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/patches/diffusion.py create mode 100644 py/torch_tensorrt/hf/exporters/patches/language.py create mode 100644 py/torch_tensorrt/hf/exporters/patches/vision.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/attention.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/mamba.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/moe.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py create mode 100644 py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py create mode 100644 py/torch_tensorrt/hf/exporters/prefix_cache.py create mode 100644 py/torch_tensorrt/hf/exporters/rope.py create mode 100644 py/torch_tensorrt/hf/exporters/utils.py diff --git a/examples/dynamo/run_groot_export.py b/examples/dynamo/run_groot_export.py index 91253911047..19baf8558f9 100644 --- a/examples/dynamo/run_groot_export.py +++ b/examples/dynamo/run_groot_export.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Smoke EdgeExporter on GR00T (4 engines: vision, language, action_context, action). +"""Smoke EdgeExporter on GR00T (4 engines: vision, language, context_projection, action). Pass the LeRobot GrootPolicy, not policy._groot_model — prepare_sample_inputs needs GrootEagleEncodeStep / embodiment_id from the policy wrapper. @@ -8,17 +8,10 @@ from __future__ import annotations import argparse -import sys from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ _TRT_PY = _REPO_ROOT / "py" -_TEST = Path("/home/micwilliams/workspace/Test") - -_test = str(_TEST) -while _test in sys.path: - sys.path.remove(_test) -sys.path.insert(0, _test) import torch # noqa: E402 import torch_tensorrt # noqa: E402 @@ -32,8 +25,8 @@ from lerobot.policies.groot.configuration_groot import GrootConfig from lerobot.utils.constants import ACTION, OBS_STATE from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from trt.plugin.plugin_utils import load_plugins_for_trt -from trt.utils import configure_thor_pytorch, force_hf_attention +from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt +from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch, force_hf_attention def load_groot(device: torch.device) -> GrootPolicy: @@ -95,7 +88,6 @@ def main() -> None: program = exporter.export(policy, sample_inputs, config=config) print("engines:", exporter.engines) - print("saved:", exporter.save_engines()) print("runtime keys:", sorted(exporter.sample)) with torch.no_grad(): diff --git a/examples/dynamo/run_nemotron_export.py b/examples/dynamo/run_nemotron_export.py index c5b59578571..7c854ae0917 100644 --- a/examples/dynamo/run_nemotron_export.py +++ b/examples/dynamo/run_nemotron_export.py @@ -8,19 +8,10 @@ from __future__ import annotations import argparse -import sys from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ _TRT_PY = _REPO_ROOT / "py" -_TEST = Path("/home/micwilliams/workspace/Test") -_NEMOTRON = _TEST / "nemotron" - -for p in (_NEMOTRON, _TEST): - s = str(p) - while s in sys.path: - sys.path.remove(s) - sys.path.insert(0, s) import torch # noqa: E402 import torch_tensorrt # noqa: E402 @@ -29,11 +20,11 @@ if _src_pkg not in list(torch_tensorrt.__path__): torch_tensorrt.__path__.append(_src_pkg) -from mamba_stub import apply as apply_mamba_stub from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter +from torch_tensorrt.hf.exporters.mamba_stub import apply as apply_mamba_stub +from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt +from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch from transformers import AutoModelForCausalLM, AutoTokenizer -from trt.plugin.plugin_utils import load_plugins_for_trt -from trt.utils import configure_thor_pytorch def load_nemotron(checkpoint: str, device: torch.device, dtype: torch.dtype): @@ -91,7 +82,6 @@ def main() -> None: program = exporter.export(model, sample_inputs, config=config) print("engines:", exporter.engines) - print("saved:", exporter.save_engines()) print("runtime keys:", sorted(exporter.sample)) with torch.no_grad(): diff --git a/examples/dynamo/run_pi05_export.py b/examples/dynamo/run_pi05_export.py index f7a25d4c53c..fc84c8ce07f 100644 --- a/examples/dynamo/run_pi05_export.py +++ b/examples/dynamo/run_pi05_export.py @@ -2,18 +2,10 @@ from __future__ import annotations import argparse -import sys from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ _TRT_PY = _REPO_ROOT / "py" -_TEST = Path("/home/micwilliams/workspace/Test") - -# ``trt.*`` lives in the Test tree. Always move it to the front. -_test = str(_TEST) -while _test in sys.path: - sys.path.remove(_test) -sys.path.insert(0, _test) import torch # noqa: E402 import torch_tensorrt # noqa: E402 @@ -26,8 +18,8 @@ from lerobot.policies.pi05 import PI05Policy from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from trt.plugin.plugin_utils import load_plugins_for_trt -from trt.utils import configure_thor_pytorch, force_hf_attention +from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt +from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch, force_hf_attention def load_pi05(device: torch.device) -> PI05Policy: @@ -97,7 +89,6 @@ def main() -> None: program = exporter.export(policy, sample_inputs, config=config) print("engines:", exporter.engines) - print("saved:", exporter.save_engines()) # Runtime kwargs are tensors only (pixel_values, lang_embeds, rope, KVs, …). runtime_kwargs = exporter.sample diff --git a/py/torch_tensorrt/hf/exporters/compile.py b/py/torch_tensorrt/hf/exporters/compile.py index 8430729dec2..df5c5c5f7a4 100644 --- a/py/torch_tensorrt/hf/exporters/compile.py +++ b/py/torch_tensorrt/hf/exporters/compile.py @@ -6,9 +6,9 @@ import torch import torch.nn as nn +import torch_tensorrt from torch_tensorrt.hf.exporters.ops import _as_tuple, record_engine from torch_tensorrt.hf.exporters.spec import ComponentBundle -from trt.plugin.plugin_utils import restore_attention DEFAULT_TRT_SETTINGS: dict[str, Any] = { "min_block_size": 1, @@ -79,7 +79,6 @@ def compile_component( }.items() if k in _TRT_COMPILE_KEYS } - import torch_tensorrt compiled = torch_tensorrt.dynamo.compile( exported, @@ -107,8 +106,12 @@ def compile_component( _write_sidecar(out_dir, bundle, name, outputs, engine_file=engine_file) return engine_path, outputs finally: - if not dryrun: - restore_attention(patched) + if not dryrun and patched is not None: + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + restore_attention, + ) + + restore_attention(patched) # type: ignore[no-untyped-call] def _write_sidecar( diff --git a/py/torch_tensorrt/hf/exporters/data.py b/py/torch_tensorrt/hf/exporters/data.py new file mode 100644 index 00000000000..e54eec2dece --- /dev/null +++ b/py/torch_tensorrt/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()) diff --git a/py/torch_tensorrt/hf/exporters/exporter.py b/py/torch_tensorrt/hf/exporters/exporter.py index 8acad2d5074..160c9a119cd 100644 --- a/py/torch_tensorrt/hf/exporters/exporter.py +++ b/py/torch_tensorrt/hf/exporters/exporter.py @@ -140,22 +140,3 @@ def _export_runtime( config.prefer_deferred_runtime_asserts_over_guards ), ) - - def save_engines(self, out_dir: str | Path | None = None) -> dict[str, Path]: - if not self.engines: - raise RuntimeError("save_engines() requires export() first") - if out_dir is None: - return {name: Path(path) for name, path in self.engines.items()} - import shutil - - dest = Path(out_dir) - dest.mkdir(parents=True, exist_ok=True) - written: dict[str, Path] = {} - for name, path in self.engines.items(): - target = dest / name - if Path(path).resolve() != target.resolve(): - if target.exists(): - shutil.rmtree(target) - shutil.copytree(path, target) - written[name] = target - return written diff --git a/py/torch_tensorrt/hf/exporters/helpers/__init__.py b/py/torch_tensorrt/hf/exporters/helpers/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/helpers/__init__.py @@ -0,0 +1 @@ + diff --git a/py/torch_tensorrt/hf/exporters/helpers/groot.py b/py/torch_tensorrt/hf/exporters/helpers/groot.py new file mode 100644 index 00000000000..47d2e3a9cba --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/helpers/groot.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from typing import Any + +import torch + +GROOT_EMBODIMENT_MAPPING = { + "new_embodiment": 31, + "oxe_droid": 17, + "agibot_genie1": 26, + "gr1": 24, +} + + +def make_embodiment_id( + policy: Any, + state: torch.Tensor, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + embodiment_tag = getattr(policy.config, "embodiment_tag", "new_embodiment") + return torch.full( + (state.shape[0],), + GROOT_EMBODIMENT_MAPPING.get(embodiment_tag, 0), + dtype=dtype, + device=device, + ) diff --git a/py/torch_tensorrt/hf/exporters/helpers/pi05.py b/py/torch_tensorrt/hf/exporters/helpers/pi05.py new file mode 100644 index 00000000000..5e15091000a --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/helpers/pi05.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +import torch +from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + + +def build_pi05_prefix_embs( + pi05_model, + img_masks, + tokens, + masks, + image_embs, + images, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compact image+language prefix embeddings for PI05 language prefill.""" + per_camera_batch = int(images[0].shape[0]) + image_embs_list = list( + image_embs.reshape(len(images), per_camera_batch, -1, image_embs.shape[-1]) + ) + + embs: list[torch.Tensor] = [] + pad_masks: list[torch.Tensor] = [] + + for img_emb, img_mask in zip(image_embs_list, img_masks, strict=True): + bsize, num_img_embs = img_emb.shape[:2] + embs.append(img_emb) + img_mask = img_mask.to(device=img_emb.device, dtype=torch.bool) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + + lang_emb = pi05_model.paligemma_with_expert.embed_language_tokens(tokens) + embs.append(lang_emb) + pad_masks.append(masks.to(device=lang_emb.device, dtype=torch.bool)) + + prefix_embs = torch.cat(embs, dim=1) + prefix_pad_masks = torch.cat(pad_masks, dim=1) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + + valid = prefix_pad_masks.to(device=prefix_embs.device, dtype=torch.bool) + valid_counts = valid.sum(dim=1) + if not torch.equal(valid_counts, valid_counts[:1].expand_as(valid_counts)): + raise ValueError( + "build_pi05_prefix_embs requires equal valid token counts across the batch" + ) + + compact_len = int(valid_counts[0].item()) + compact_embs = torch.stack( + [prefix_embs[b, valid[b], :] for b in range(prefix_embs.shape[0])], + dim=0, + ) + compact_position_ids = torch.stack( + [prefix_position_ids[b, valid[b]] for b in range(prefix_position_ids.shape[0])], + dim=0, + ) + compact_pad_mask = torch.ones( + prefix_embs.shape[0], + compact_len, + device=prefix_pad_masks.device, + dtype=torch.bool, + ) + compact_attention_mask = torch.zeros( + prefix_embs.shape[0], + 1, + compact_len, + compact_len, + device=prefix_embs.device, + dtype=torch.float32, + ) + return compact_embs, compact_pad_mask, compact_attention_mask, compact_position_ids + + +def pi05_compact_index( + img_masks, + images: list, + seq_len_per_image: int, + lang_masks: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + """Gather index for leftover ``fuse_prefix`` (same valid rows as ``build_pi05_prefix_embs``).""" + batch = int(lang_masks.shape[0]) + pads: list[torch.Tensor] = [] + for img_mask in img_masks: + mask = img_mask.to(device=device, dtype=torch.bool) + pads.append(mask[:, None].expand(batch, int(seq_len_per_image))) + vision_pad = torch.cat(pads, dim=1) + lang_pad = lang_masks.to(device=device, dtype=torch.bool) + valid = torch.cat([vision_pad, lang_pad], dim=1) + counts = valid.sum(dim=1) + if not torch.equal(counts, counts[:1].expand_as(counts)): + raise ValueError( + "pi05_compact_index requires equal valid token counts across the batch" + ) + return torch.stack( + [torch.nonzero(valid[b], as_tuple=False).squeeze(-1) for b in range(batch)], + dim=0, + ) + + +def pi05_seq_len_per_image(image_embs: torch.Tensor, images: list) -> int: + """Vision token count per image view after PI05 camera reshape.""" + per_camera_batch = int(images[0].shape[0]) + num_images = len(images) + reshaped = image_embs.reshape( + num_images, per_camera_batch, -1, image_embs.shape[-1] + ) + return int(reshaped.shape[2]) + + +def pi05_prefix_max_seq_len( + *, + num_images: int, + seq_len_per_image: int, + tokenizer_max_length: int, +) -> int: + """Upper bound on compact PI05 prefix length (vision slots + language tokens).""" + return int(num_images) * int(seq_len_per_image) + int(tokenizer_max_length) + + +def pi05_compact_prefix_max_seq_len( + image_embs: torch.Tensor, + images: list, + tokenizer_max_length: int, +) -> int: + """Static TRT prefill length for PI05 compact prefix (vision + language slots).""" + return pi05_prefix_max_seq_len( + num_images=len(images), + seq_len_per_image=pi05_seq_len_per_image(image_embs, images), + tokenizer_max_length=tokenizer_max_length, + ) + + +def pad_pi05_compact_prefix( + prefix_embs: torch.Tensor, + prefix_pad_mask: torch.Tensor, + prefix_attention_mask: torch.Tensor, + prefix_position_ids: torch.Tensor, + *, + max_seq_len: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Right-pad a compact PI05 prefix to ``max_seq_len`` for static TRT engines. + + Returns padded tensors and the valid (unpadded) sequence length. + """ + valid_len = int(prefix_embs.shape[1]) + max_seq_len = int(max_seq_len) + if valid_len > max_seq_len: + raise ValueError( + f"PI05 compact prefix length {valid_len} exceeds max_seq_len {max_seq_len}" + ) + if valid_len == max_seq_len: + return ( + prefix_embs, + prefix_pad_mask, + prefix_attention_mask, + prefix_position_ids, + valid_len, + ) + + batch_size = int(prefix_embs.shape[0]) + pad_len = max_seq_len - valid_len + device = prefix_embs.device + dtype = prefix_embs.dtype + + padded_embs = torch.zeros( + batch_size, max_seq_len, prefix_embs.shape[-1], device=device, dtype=dtype + ) + padded_embs[:, :valid_len, :] = prefix_embs + + padded_pad_mask = torch.zeros( + batch_size, max_seq_len, device=device, dtype=torch.bool + ) + padded_pad_mask[:, :valid_len] = prefix_pad_mask.to(device=device, dtype=torch.bool) + + padded_position_ids = torch.zeros( + batch_size, max_seq_len, device=device, dtype=prefix_position_ids.dtype + ) + padded_position_ids[:, :valid_len] = prefix_position_ids + + # Bidirectional padding mask: valid x valid block only. + padded_attention_mask = torch.zeros( + batch_size, + 1, + max_seq_len, + max_seq_len, + device=device, + dtype=prefix_attention_mask.dtype, + ) + padded_attention_mask[:, :, :valid_len, :valid_len] = prefix_attention_mask[ + :, :, :valid_len, :valid_len + ] + + return ( + padded_embs, + padded_pad_mask, + padded_attention_mask, + padded_position_ids, + valid_len, + ) + + +def make_pi05_suffix_position_and_mask(core, prefix_pad_masks, x_t, device): + """Suffix position ids and 4D attention mask for PI05 diffusion.""" + batch_size, suffix_len = x_t.shape[:2] + prefix_pad_masks = prefix_pad_masks.to(device=device) + prefix_len = prefix_pad_masks.shape[1] + + suffix_pad_masks = torch.ones( + batch_size, suffix_len, dtype=torch.bool, device=device + ) + suffix_att_masks = torch.tensor( + [1] + [0] * (suffix_len - 1), + dtype=torch.int64, + device=device, + )[None, :].expand(batch_size, -1) + + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( + batch_size, suffix_len, prefix_len + ) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + + attention_mask = core._prepare_attention_masks_4d(full_att_2d_masks) + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + return position_ids, attention_mask diff --git a/py/torch_tensorrt/hf/exporters/mamba_stub.py b/py/torch_tensorrt/hf/exporters/mamba_stub.py new file mode 100644 index 00000000000..ea0ae9d20bf --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/mamba_stub.py @@ -0,0 +1,88 @@ +"""Import stubs so Hub ``modeling_nemotron_h.py`` can load without ``mamba-ssm``. + +Copied from Edge-LLM ``nemotron_h_patch``: only the sys.modules stubs, not the +ONNX ``from_config`` / dense-MoE forward replacements. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from typing import Any, Dict + +import torch +import torch.nn.functional as F + + +def _rms_norm_ref( + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + upcast=True, +): + dtype = x.dtype + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + z = z.float() if z is not None else z + if z is not None and not norm_before_gate: + x = x * F.silu(z) + if group_size is None: + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + else: + *lead, last = x.shape + x_group = x.reshape(*lead, last // group_size, group_size) + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) + out = (x_group * rstd).reshape(*lead, last) * weight + if bias is not None: + out = out + bias + if z is not None and norm_before_gate: + out *= F.silu(z) + return out.to(dtype) + + +def _stub_layernorm_gated() -> None: + name = "mamba_ssm.ops.triton.layernorm_gated" + stub = types.ModuleType(name) + stub.rmsnorm_fn = _rms_norm_ref + sys.modules[name] = stub + + +def _stub_if_broken(pkg_name: str, sentinel_attrs: Dict[str, Any]) -> None: + try: + spec = importlib.util.find_spec(pkg_name) + except (ValueError, ImportError): + spec = None + if spec is None: + return + stub = types.ModuleType(pkg_name) + stub.__spec__ = importlib.util.spec_from_loader(pkg_name, loader=None) + for attr, value in sentinel_attrs.items(): + setattr(stub, attr, value) + sys.modules[pkg_name] = stub + + +def apply() -> None: + _stub_layernorm_gated() + _stub_if_broken( + "causal_conv1d", + {"causal_conv1d_fn": None, "causal_conv1d_update": None}, + ) + _stub_if_broken( + "mamba_ssm.ops.triton.selective_state_update", + {"selective_state_update": None}, + ) + _stub_if_broken( + "mamba_ssm.ops.triton.ssd_combined", + { + "mamba_chunk_scan_combined": None, + "mamba_split_conv1d_scan_combined": None, + }, + ) diff --git a/py/torch_tensorrt/hf/exporters/patches/__init__.py b/py/torch_tensorrt/hf/exporters/patches/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/patches/__init__.py @@ -0,0 +1 @@ + diff --git a/py/torch_tensorrt/hf/exporters/patches/diffusion.py b/py/torch_tensorrt/hf/exporters/patches/diffusion.py new file mode 100644 index 00000000000..243cdffd1e4 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/patches/diffusion.py @@ -0,0 +1,338 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding +from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache + + +class TRTFixedCategorySpecificLinearPatch(nn.Module): + """Freeze one GR00T embodiment-specific Linear into a normal Linear. + + GR00T stores one weight matrix per robot embodiment and selects it with + embodiment_id at runtime. For TensorRT deployment we compile one robot at a + time, so this wrapper picks that robot's weights once in __init__ and the + forward path becomes a plain static F.linear. + """ + + def __init__(self, layer: nn.Module, embodiment_id: torch.Tensor): + super().__init__() + + cat_id = int(embodiment_id.flatten()[0].item()) + + # Original: [num_embodiments, input_dim, output_dim] + # using cat_id selects the weight matrix for one embodiment/robot -> [input_dim, output_dim] + # nn.functional.linear expects -> weight: [output_dim, input_dim] so we transpose + weight = layer.W[cat_id].transpose(0, 1).contiguous() + bias = layer.b[cat_id].contiguous() + + # detach() breaks any autograd link to the original multi-embodiment + # parameter; clone() gives this fixed wrapper independent storage for + # the selected slice. This copy happens once during wrapper creation, + # not in forward, and lets TensorRT see normal immutable weights. + self.weight = nn.Parameter(weight.detach().clone(), requires_grad=False) + self.bias = nn.Parameter(bias.detach().clone(), requires_grad=False) + self.out_features = int(bias.shape[0]) + + def forward(self, x): + # x: [B, T, input_dim] + # out: [B, T, output_dim] + batch_size = x.shape[0] + seq_len = x.shape[1] + x = x.reshape(batch_size * seq_len, x.shape[-1]) + x = F.linear(x, self.weight, self.bias) + return x.reshape(batch_size, seq_len, self.out_features) + + +class TRTDynamicCategorySpecificLinearPatch(nn.Module): + """TensorRT-friendly dynamic version of GR00T CategorySpecificLinear. + + Unlike the fixed wrapper, this keeps the full embodiment weight bank and + uses runtime embodiment_id values to gather W/b for each batch item. The + math stays equivalent to GR00T category-specific linear: + x [B,T,in] @ W[embodiment] [B,in,out] + b[embodiment]. + """ + + def __init__(self, layer: nn.Module): + super().__init__() + + # Keep the full embodiment weight bank. + # W: [num_embodiments, input_dim, output_dim] + # b: [num_embodiments, output_dim] + self.W = layer.W + self.b = layer.b + + def forward(self, x, cat_ids): + # x: [B, T, input_dim] + # cat_ids: [B] + + cat_ids = cat_ids.to(dtype=torch.long) + + # selected_w: [B, input_dim, output_dim] + # selected_b: [B, output_dim] + selected_w = torch.index_select(self.W, dim=0, index=cat_ids).to(dtype=x.dtype) + selected_b = torch.index_select(self.b, dim=0, index=cat_ids).to(dtype=x.dtype) + + # out: [B, T, output_dim] + out = torch.bmm(x, selected_w) + + # bias: [B, 1, output_dim], broadcast over T + return out + selected_b.unsqueeze(1) + + +class TRTDynamicCategorySpecificMLPPatch(nn.Module): + """Dynamic two-layer category-specific MLP used by GR00T. + + GR00T state encoders and action decoders are CategorySpecificMLP modules: + each contains layer1/layer2 CategorySpecificLinear layers. This wrapper + preserves runtime embodiment selection for both layers. + """ + + def __init__(self, mlp: nn.Module): + super().__init__() + self.layer1 = TRTDynamicCategorySpecificLinearPatch(mlp.layer1) + self.layer2 = TRTDynamicCategorySpecificLinearPatch(mlp.layer2) + + def forward(self, x, embodiment_id): + hidden = F.relu(self.layer1(x, embodiment_id)) + return self.layer2(hidden, embodiment_id) + + +class TRTGrootActionEncoderPatch(nn.Module): + def __init__(self, action_encoder: nn.Module, embodiment_id: torch.Tensor): + super().__init__() + self.W1 = TRTFixedCategorySpecificLinearPatch(action_encoder.W1, embodiment_id) + self.W2 = TRTFixedCategorySpecificLinearPatch(action_encoder.W2, embodiment_id) + self.W3 = TRTFixedCategorySpecificLinearPatch(action_encoder.W3, embodiment_id) + self.pos_encoding = action_encoder.pos_encoding + + def forward(self, actions, timesteps, embodiment_id): + batch_size, action_horizon, _ = actions.shape + + if timesteps.dim() == 1 and timesteps.shape[0] == batch_size: + timesteps = timesteps.unsqueeze(1).expand(-1, action_horizon) + else: + raise ValueError("Expected `timesteps` to have shape (B,).") + + action_emb = self.W1(actions) + timestep_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype) + hidden = torch.cat([action_emb, timestep_emb], dim=-1) + hidden = F.silu(self.W2(hidden)) + return self.W3(hidden) + + +class TRTDynamicGrootActionEncoderPatch(nn.Module): + """Dynamic GR00T noisy-action encoder. + + The original action encoder uses three embodiment-specific linear layers + around the action embedding, timestep positional embedding, and SiLU block. + This wrapper keeps embodiment_id dynamic while spelling the category-specific + pieces as index_select + bmm so Torch-TRT can lower them reliably. + """ + + def __init__(self, action_encoder: nn.Module): + super().__init__() + self.W1 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W1) + self.W2 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W2) + self.W3 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W3) + self.pos_encoding = action_encoder.pos_encoding + + def forward(self, actions, timesteps, embodiment_id): + batch_size, action_horizon, _ = actions.shape + + timesteps = timesteps.unsqueeze(1).expand(-1, action_horizon) + + action_emb = self.W1(actions, embodiment_id) + timestep_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype) + + hidden = torch.cat([action_emb, timestep_emb], dim=-1) + hidden = F.silu(self.W2(hidden, embodiment_id)) + return self.W3(hidden, embodiment_id) + + +class ActionStepEncoderPatch(nn.Module): + """Base contract for model-specific action-step encoding. + + Subclasses implement forward() to turn noisy actions, timestep, and + model-specific context tensors into the args/kwargs consumed by the action + expert and velocity decoder. The default helpers cover common expert output + and velocity shapes, while model-specific encoders can override them. + """ + + def get_action_hidden(self, expert_out, output_tokens: int): + # Default path for experts that return either a standard model output + # with last_hidden_state, a raw hidden-state tensor, or a tuple/list whose + # first item is the hidden-state tensor. + hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out + ) + + if isinstance(hidden, (tuple, list)): + hidden = hidden[0] + + return hidden[:, -output_tokens:] + + def process_velocity(self, velocity): + # Default path for models whose decoder already returns the final action + # velocity shape. Override for models that need reshaping or cropping. + return velocity + + +class StaticActionVelocityStepPatch(nn.Module): + """One static denoising step shared by VLA action diffusion modules. + + The model-specific step_encoder owns the messy part: converting noisy + actions, timestep, and context tensors into the exact action_expert call. + This wrapper only runs the expert, selects action-token hidden states, and + decodes those hidden states into a velocity update. + """ + + def __init__( + self, + *, + step_encoder: ActionStepEncoderPatch, + action_expert: nn.Module, + velocity_decoder: nn.Module, + output_tokens: int, + cast_hidden_fp32: bool = True, + ): + super().__init__() + self.step_encoder = step_encoder + self.action_expert = action_expert + self.velocity_decoder = velocity_decoder + self.output_tokens = int(output_tokens) + self.cast_hidden_fp32 = cast_hidden_fp32 + + def forward(self, x_t, timestep, *inputs): + # Build the action expert inputs and any decoder-specific side inputs. + expert_args, expert_kwargs, decoder_args, decoder_kwargs = self.step_encoder( + x_t, + timestep, + *inputs, + ) + + # Run the model-specific action expert: Gemma expert, DiT, etc. + expert_out = self.action_expert(*expert_args, **expert_kwargs) + + # Most experts return last_hidden_state, but some wrappers return tuples + # or need custom suffix-token selection. + action_hidden = self.step_encoder.get_action_hidden( + expert_out, + self.output_tokens, + ) + + if self.cast_hidden_fp32: + action_hidden = action_hidden.to(dtype=torch.float32) + + # Project action-token hidden states back to action-space velocity. + velocity = self.velocity_decoder( + action_hidden, + *decoder_args, + **decoder_kwargs, + ) + + return self.step_encoder.process_velocity(velocity) + + +class GrootDiTStepEncoderPatch(ActionStepEncoderPatch): + def __init__(self, action_head, embodiment_id: torch.Tensor | None = None): + super().__init__() + if embodiment_id is None: + self.state_encoder = action_head.state_encoder + self.action_encoder = action_head.action_encoder + else: + # Keep embodiment_id as a runtime input while replacing GR00T's + # category-specific modules with Torch-TRT-friendly dynamic wrappers. + self.state_encoder = TRTDynamicCategorySpecificMLPPatch( + action_head.state_encoder + ) + self.action_encoder = TRTDynamicGrootActionEncoderPatch( + action_head.action_encoder + ) + self.future_tokens = action_head.future_tokens + self.position_embedding = getattr(action_head, "position_embedding", None) + self.add_pos_embed = action_head.config.add_pos_embed + + def forward(self, actions, timestep, vl_embs, state, embodiment_id): + state_features = self.state_encoder(state, embodiment_id) + action_features = self.action_encoder(actions, timestep, embodiment_id) + + if self.add_pos_embed: + pos_ids = torch.arange( + action_features.shape[1], + dtype=torch.long, + device=action_features.device, + ) + action_features = action_features + self.position_embedding( + pos_ids + ).unsqueeze(0) + + future_tokens = self.future_tokens.weight.unsqueeze(0).expand( + vl_embs.shape[0], + -1, + -1, + ) + + sa_embs = torch.cat( + (state_features, future_tokens, action_features), + dim=1, + ) + + expert_args = () + expert_kwargs = { + "hidden_states": sa_embs, + "encoder_hidden_states": vl_embs, + "timestep": timestep, + } + + decoder_args = (embodiment_id,) + decoder_kwargs = {} + + return expert_args, expert_kwargs, decoder_args, decoder_kwargs + + +class PI05PrefixKVStepEncoderPatch(ActionStepEncoderPatch): + """PI05 suffix embed + AdaRMS cond, consumed by Gemma action expert.""" + + def __init__(self, core): + super().__init__() + self.action_in_proj = core.action_in_proj + self.time_mlp_in = core.time_mlp_in + self.time_mlp_out = core.time_mlp_out + self.config = core.config + self.hidden_size = core.action_in_proj.out_features + + def forward( + self, + x_t, + timestep, + prefix_k, + prefix_v, + position_ids, + attention_mask, + ): + suffix_embs = self.action_in_proj(x_t) + + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.hidden_size, + min_period=self.config.min_period, + max_period=self.config.max_period, + device=timestep.device, + ).to(dtype=suffix_embs.dtype) + + adarms_cond = self.time_mlp_in(time_emb) + adarms_cond = F.silu(adarms_cond) + adarms_cond = self.time_mlp_out(adarms_cond) + adarms_cond = F.silu(adarms_cond) + + expert_kwargs = { + "inputs_embeds": suffix_embs, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values": PrefixKVCache(prefix_k, prefix_v), + "use_cache": False, + "adarms_cond": adarms_cond, + } + return (), expert_kwargs, (), {} diff --git a/py/torch_tensorrt/hf/exporters/patches/language.py b/py/torch_tensorrt/hf/exporters/patches/language.py new file mode 100644 index 00000000000..0ceffc87640 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/patches/language.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import torch +import torch.nn as nn + + +def _as_tensor(x): + """Unwrap tuple/list outputs from patched attention modules.""" + if isinstance(x, (tuple, list)): + return x[0] + return x + + +def language_decoder(language: nn.Module) -> nn.Module: + """Inner module that owns ``.layers``. + + Paligemma / PI05 store layers on ``language_model`` itself. HF + ``*ForCausalLM`` stores them on ``language_model.model``. Prefer ``.layers`` + so a stray ``.model`` attribute cannot silently pick the wrong submodule. + """ + if hasattr(language, "layers"): + return language + inner = getattr(language, "model", None) + if isinstance(inner, nn.Module) and hasattr(inner, "layers"): + return inner + 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] + + +class CausalLMPatch(nn.Module): + """Edge-LLM causal LM: manual decoder loop -> logits + lm_hidden_states + prefix KV. + + External RoPE and KV-cache controls match ``LLMEngineRunner`` prefill/decode. + ``select_layer=-1`` (default for GR00T TRT) uses final RMSNorm hidden for + context; positive values capture an intermediate layer output pre-norm. + + ``ds_stack`` is always an input (``[num_ds, B, S, H]``). Models without + deepstack pass ``num_ds=0`` (empty leading dim); the add is a no-op. + """ + + def __init__( + self, + lm: nn.Module, + lm_head: nn.Module, + *, + select_layer: int = -1, + ): + super().__init__() + self.lm = lm + self.lm_head = lm_head + self.select_layer = int(select_layer) + + def forward( + self, + inputs_embeds: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + context_lengths: torch.Tensor, + kvcache_start_index: torch.Tensor, + last_token_ids: torch.Tensor, + ds_stack: torch.Tensor, + *past_key_values: torch.Tensor, + ): + lm_dtype = next(self.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 self.select_layer == 0 else None + new_kvs = [] + + for i, layer in enumerate(self.lm.layers): + residual = hidden + hidden = _as_tensor(layer.input_layernorm(hidden)) + hidden, kv = layer.self_attn( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=past_key_values[i], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, + ) + hidden = _as_tensor(hidden) + hidden = residual + hidden + + residual = hidden + hidden = _as_tensor(layer.post_attention_layernorm(hidden)) + hidden = _as_tensor(layer.mlp(hidden)) + hidden = residual + hidden + new_kvs.append(kv) + + if i < num_ds: + hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) + + if self.select_layer > 0 and (i + 1) == self.select_layer: + context_hidden = hidden + + hidden = _as_tensor(self.lm.norm(hidden)) + if context_hidden is None: + context_hidden = hidden + + last_hidden = gather_last_token_hidden(hidden, last_token_ids) + logits = self.lm_head(last_hidden).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 + + +# specific to gr00t, before action another project is required for context embeddings +class ContextProjectionPatch(nn.Module): + """eagle_linear -> vlln -> vl_self_attention (matches eager context path).""" + + def __init__(self, eagle_linear, vlln, vl_self_attention): + super().__init__() + self.eagle_linear = eagle_linear + self.vlln = vlln + self.vl_self_attention = vl_self_attention + + def forward(self, hidden_states: torch.Tensor): + context_embs = self.eagle_linear(hidden_states) + + vlln_weight = getattr(self.vlln, "weight", None) + if vlln_weight is not None: + context_embs = context_embs.to(dtype=vlln_weight.dtype) + + context_embs = self.vlln(context_embs) + context_embs = self.vl_self_attention(context_embs) + return context_embs diff --git a/py/torch_tensorrt/hf/exporters/patches/vision.py b/py/torch_tensorrt/hf/exporters/patches/vision.py new file mode 100644 index 00000000000..152570de40a --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/patches/vision.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def hwc_to_nchw(images: torch.Tensor) -> torch.Tensor: + if images.ndim != 4: + raise ValueError(f"Expected 4D images, got shape {tuple(images.shape)}") + return images.permute(0, 3, 1, 2).contiguous() + + +def is_nchw_pixel_values(pixel_values: torch.Tensor) -> bool: + return ( + pixel_values.ndim == 4 + and pixel_values.shape[1] in (1, 3, 4) + and pixel_values.shape[-1] not in (1, 3, 4) + ) + + +# --------------------------------------------------------------------------- +# Base +# --------------------------------------------------------------------------- + + +class VisionPatch(nn.Module): + """Base TRT trace target: probe static shapes once, flatten output to [N, H].""" + + cast_output_to_input_dtype: bool + output_num_tokens: int + output_hidden_size: int + + def _finalize_output( + self, features: torch.Tensor, out_dtype: torch.dtype + ) -> torch.Tensor: + if self.cast_output_to_input_dtype and features.dtype != out_dtype: + features = features.to(out_dtype) + if features.ndim == 3: + # [B, S, H] -> [B*S, H] + return features.reshape(-1, features.shape[-1]) + if features.ndim == 2: + # already [N, H] (token-pooling encoders) + return features + raise ValueError( + f"Expected 2D or 3D features, got shape {tuple(features.shape)}" + ) + + +# --------------------------------------------------------------------------- +# Grid vision (PI0.5 / GR00T / SmolVLA VitRunner path) +# --------------------------------------------------------------------------- + + +class GridVisionPatch(VisionPatch): + """pixels [B,H,W,C] -> fixed patch grid -> [B*seq_len, lm_hidden]""" + + def __init__( + self, + *, + vision_model: nn.Module, + projector: nn.Module, + sample_pixel_values: torch.Tensor, + select_layer: int = -1, + pixel_shuffle: bool = False, + downsample_ratio: float = 0.5, + force_float32_input: bool = False, + cast_output_to_input_dtype: bool = False, + vision_kwargs: dict[str, Any] | None = None, + ): + super().__init__() + self.vision_model = vision_model + self.projector = projector + self.select_layer = int(select_layer) + self.pixel_shuffle = bool(pixel_shuffle) + self.downsample_ratio = float(downsample_ratio) + self.force_float32_input = bool(force_float32_input) + self.cast_output_to_input_dtype = bool(cast_output_to_input_dtype) + self.vision_kwargs = dict(vision_kwargs or {}) + + with torch.no_grad(): + sample = sample_pixel_values + if self.force_float32_input and sample.dtype != torch.float32: + sample = sample.to(torch.float32) + + vit_embeds = self._select_vision_features(self._run_vision(sample)) + self.seq_len = int(vit_embeds.shape[1]) + self.hidden_size = int(vit_embeds.shape[2]) + + if self.pixel_shuffle: + self._init_pixel_shuffle_shape() + vit_embeds = self._apply_pixel_shuffle(vit_embeds) + + projected = self.projector(self._projector_input(vit_embeds)) + self.batch_size = int(projected.shape[0]) + self.output_seq_len = int(projected.shape[1]) + self.output_hidden_size = int(projected.shape[2]) + self.output_num_tokens = self.batch_size * self.output_seq_len + + def _run_vision(self, images: torch.Tensor): + pixel_values = ( + hwc_to_nchw(images) if not is_nchw_pixel_values(images) else images + ) + kwargs = dict(self.vision_kwargs) + kwargs["pixel_values"] = pixel_values + kwargs["output_hidden_states"] = self.select_layer != -1 + kwargs.setdefault("return_dict", True) + return self.vision_model(**kwargs) + + def _select_vision_features(self, out): + if self.select_layer == -1: + return ( + out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] + ) + return out.hidden_states[self.select_layer] + + def _projector_input(self, vit_embeds: torch.Tensor) -> torch.Tensor: + proj_dtype = next(self.projector.parameters()).dtype + if vit_embeds.dtype != proj_dtype: + return vit_embeds.to(proj_dtype) + return vit_embeds + + def _init_pixel_shuffle_shape(self): + side = int(self.seq_len**0.5) + if side * side != self.seq_len: + raise ValueError( + f"Expected square vision sequence, got seq_len={self.seq_len}" + ) + self.grid_w = side + self.grid_h = side + self.out_w = int(self.grid_w * self.downsample_ratio) + self.out_h = int(self.grid_h * self.downsample_ratio) + self.hidden_after_first_view = int(self.hidden_size / self.downsample_ratio) + self.shuffle_hidden = int( + self.hidden_size / (self.downsample_ratio * self.downsample_ratio) + ) + + def _apply_pixel_shuffle(self, x): + n = x.shape[0] + x = x.reshape(n, self.grid_w, self.out_h, self.hidden_after_first_view) + x = x.permute(0, 2, 1, 3).contiguous() + x = x.reshape(n, self.out_h, self.out_w, self.shuffle_hidden) + x = x.permute(0, 2, 1, 3).contiguous() + return x.reshape(n, -1, self.shuffle_hidden) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + out_dtype = pixel_values.dtype + images = pixel_values + if self.force_float32_input and images.dtype != torch.float32: + images = images.to(torch.float32) + + vit_embeds = self._select_vision_features(self._run_vision(images)) + if self.pixel_shuffle: + vit_embeds = self._apply_pixel_shuffle(vit_embeds) + + features = self.projector(self._projector_input(vit_embeds)) + return self._finalize_output(features, out_dtype) diff --git a/py/torch_tensorrt/hf/exporters/plugin/__init__.py b/py/torch_tensorrt/hf/exporters/plugin/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/__init__.py @@ -0,0 +1 @@ + diff --git a/py/torch_tensorrt/hf/exporters/plugin/attention.py b/py/torch_tensorrt/hf/exporters/plugin/attention.py new file mode 100644 index 00000000000..c2d29b7065d --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/attention.py @@ -0,0 +1,557 @@ +from enum import IntEnum +from typing import Optional, Tuple + +import torch +import torch.nn as nn + + +class ContextAttentionMaskType(IntEnum): + """Context prefill mask type, mirroring the C++ ``ContextAttentionMaskType`` enum. + + The integer values must stay in sync with + ``cpp/kernels/contextAttentionKernels/fmhaParams_v2.h`` since they are passed + directly to the ``AttentionPlugin`` ``context_attention_mask_type`` field. + """ + + PADDING = 0 # Bidirectional full-prefix (attend to all valid tokens) + CAUSAL = 1 + SLIDING_OR_CHUNKED_CAUSAL = 2 + CUSTOM_MASK = 3 + + +class PluginAttention(nn.Module): + """ + Model-agnostic Plugin Attention module that replaces standard attention. + + This module wraps the projection layers from the original attention module + and uses ``trt.attention_plugin`` with separate Q/K/V tensors for the + attention computation. + + Supports: + - Qwen2.5, Llama: Standard attention + - Qwen3: Attention with QK Normalization (q_norm, k_norm) + """ + + def __init__( + self, + original_attn: nn.Module, + *, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + hidden_size: int, + layer_idx: int, + context_attention_mask_type: int = ContextAttentionMaskType.PADDING, + ): + """ + Initialize PluginAttention. + + Args: + original_attn: The original attention module to wrap. + num_attention_heads: Number of query attention heads. + num_key_value_heads: Number of key/value attention heads. + head_dim: Per-head dimension. + hidden_size: Model hidden size. + layer_idx: Index of this layer in the model. + context_attention_mask_type: Context prefill mask type + (``ContextAttentionMaskType`` enum value). + """ + super().__init__() + self.q_proj = original_attn.q_proj + self.k_proj = original_attn.k_proj + self.v_proj = original_attn.v_proj + self.o_proj = original_attn.o_proj + + # Qwen3 has QK Normalization + self.q_norm = getattr(original_attn, "q_norm", None) + self.k_norm = getattr(original_attn, "k_norm", None) + + self.num_heads = int(num_attention_heads) + self.num_key_value_heads = int(num_key_value_heads) + self.head_dim = int(head_dim) + self.attn_hidden_size = self.num_heads * self.head_dim + self.hidden_size = int(hidden_size) + self.layer_idx = layer_idx + self.context_attention_mask_type = int(context_attention_mask_type) + + def forward( + self, + hidden_states: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + past_key_value: Optional[torch.Tensor] = None, + ctx_len: Optional[torch.Tensor] = None, + kvcache_start_index: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass using the plugin attention. + + Args: + hidden_states: Input tensor of shape [batch, seq_len, hidden_size]. + rope_rotary_cos_sin: External RoPE cache of shape + [rope_batch, max_seq_len, rotary_dim] (float32). For standard + RoPE, rope_batch is typically 1 and the plugin broadcasts over + batch. Layout: cos in [:, :, :rotary_dim // 2], sin in + [:, :, rotary_dim // 2:]. Supplied at export as a graph input; + at runtime filled by LLMEngineRunner (not computed here). + attention_mask: Unused (plugin handles masking internally). + position_ids: Unused; RoPE lookup uses rope_rotary_cos_sin and ctx_len. + past_key_value: KV cache tensor of shape [batch, 2, num_kv_heads, capacity, head_dim]. + ctx_len: Context length tensor for each batch item. + kvcache_start_index: External KV cache start indices. Empty tensor + ``[0]`` for fresh prefill; ``[batch]`` for decode/chunked prefill. + + Returns: + Tuple of (output tensor, updated KV cache). + """ + batch_size, seq_len, _ = hidden_states.shape + + # Ensure rope embeddings are FP32 + assert ( + rope_rotary_cos_sin.dtype == torch.float32 + ), "rope_rotary_cos_sin must be FP32" + + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # Qwen3: Apply QK Normalization if available + if self.q_norm is not None: + # Reshape for per-head normalization: [B, S, num_heads, head_dim] + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + q = self.q_norm(q) + q = q.view(batch_size, seq_len, -1) + + if self.k_norm is not None: + # Reshape for per-head normalization: [B, S, num_kv_heads, head_dim] + k = k.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim) + k = self.k_norm(k) + k = k.view(batch_size, seq_len, -1) + + if ctx_len is None: + ctx_len = torch.tensor( + [seq_len], dtype=torch.int32, device=hidden_states.device + ).expand(batch_size) + + if past_key_value is None: + raise ValueError("past_key_value (KV cache tensor) must be provided") + + if kvcache_start_index is None: + raise ValueError("kvcache_start_index must be provided") + + dtype = q.dtype + q = q.to(torch.float16) + k = k.to(torch.float16) + v = v.to(torch.float16) + + attn_out, updated_kv = torch.ops.trt.attention_plugin.default( + q, + k, + v, + past_key_value, + ctx_len, + rope_rotary_cos_sin, + kvcache_start_index, + self.num_heads, + self.num_key_value_heads, + False, + self.head_dim, + False, + -1, + self.context_attention_mask_type, + ) + + # Use attn_hidden_size for reshape (may differ from hidden_size in Qwen3) + attn_out = attn_out.reshape(batch_size, seq_len, self.attn_hidden_size).to( + dtype + ) + output = self.o_proj(attn_out) + return output, updated_kv + + +class ViTPluginAttention(nn.Module): + def __init__( + self, + attn, + *, + batch_size: int, + seq_len: int, + name: str, + allow_attention_mask: bool = False, + ): + super().__init__() + self.q_proj = attn.q_proj + self.k_proj = attn.k_proj + self.v_proj = attn.v_proj + self.out_proj = attn.out_proj + self.num_heads = int(attn.num_heads) + self.head_dim = int(attn.head_dim) + self.name = name + self.allow_attention_mask = bool(allow_attention_mask) + + device = self.q_proj.weight.device + + cu_seqlens = torch.arange( + 0, + (int(batch_size) + 1) * int(seq_len), + int(seq_len), + device=device, + dtype=torch.int32, + ) + max_seqlen_carrier = torch.zeros( + int(seq_len), + device=device, + dtype=torch.int32, + ) + + self.register_buffer("cu_seqlens", cu_seqlens, persistent=False) + self.register_buffer("max_seqlen_carrier", max_seqlen_carrier, persistent=False) + + def forward(self, hidden_states, attention_mask=None, **kwargs): + if attention_mask is not None and not self.allow_attention_mask: + raise RuntimeError( + f"{self.name} ViT plugin path expects no vision attention_mask" + ) + + batch_size, seq_len, _ = hidden_states.shape + + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q = ( + q.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + k = ( + k.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + v = ( + v.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + + attn_output = torch.ops.trt.vit_attention_plugin.default( + q, + k, + v, + self.cu_seqlens, + self.max_seqlen_carrier, + self.num_heads, + self.head_dim, + ) + + attn_output = attn_output.reshape( + batch_size, seq_len, self.num_heads * self.head_dim + ) + attn_output = attn_output.to(dtype=self.out_proj.weight.dtype) + attn_output = self.out_proj(attn_output) + return attn_output, None + + +class MolmoViTPluginAttention(nn.Module): + """TRT ViT attention wrapper for MolmoAct2 vision self-attention. + + Molmo vision uses ``wq/wk/wv/wo`` and returns a tensor directly, unlike the + SigLIP ``q_proj/k_proj/v_proj/out_proj`` modules wrapped by ViTPluginAttention. + """ + + def __init__(self, attn, *, batch_size: int, seq_len: int, name: str): + super().__init__() + self.wq = attn.wq + self.wk = attn.wk + self.wv = attn.wv + self.wo = attn.wo + self.residual_dropout = attn.residual_dropout + + self.num_heads = int(attn.num_heads) + self.num_key_value_heads = int(attn.num_key_value_heads) + self.num_key_value_groups = int(attn.num_key_value_groups) + self.head_dim = int(attn.head_dim) + self.name = name + + device = self.wq.weight.device + self.register_buffer( + "cu_seqlens", + torch.arange( + 0, + (int(batch_size) + 1) * int(seq_len), + int(seq_len), + device=device, + dtype=torch.int32, + ), + persistent=False, + ) + self.register_buffer( + "max_seqlen_carrier", + torch.zeros(int(seq_len), device=device, dtype=torch.int32), + persistent=False, + ) + + def forward( + self, + inputs_q: torch.Tensor, + inputs_kv: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_kv is not None: + raise RuntimeError( + f"{self.name} Molmo ViT plugin path only supports self-attention" + ) + if attn_mask is not None: + raise RuntimeError( + f"{self.name} Molmo ViT plugin path expects no attn_mask" + ) + + batch_size, seq_len, _ = inputs_q.shape + q = self.wq(inputs_q).reshape( + batch_size, seq_len, self.num_heads, self.head_dim + ) + k = self.wk(inputs_q).reshape( + batch_size, seq_len, self.num_key_value_heads, self.head_dim + ) + v = self.wv(inputs_q).reshape( + batch_size, seq_len, self.num_key_value_heads, self.head_dim + ) + + if self.num_heads != self.num_key_value_heads: + k = k.repeat_interleave(self.num_key_value_groups, dim=2) + v = v.repeat_interleave(self.num_key_value_groups, dim=2) + + q = ( + q.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + k = ( + k.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + v = ( + v.reshape(batch_size * seq_len, self.num_heads, self.head_dim) + .to(torch.float16) + .contiguous() + ) + + attn_output = torch.ops.trt.vit_attention_plugin.default( + q, + k, + v, + self.cu_seqlens, + self.max_seqlen_carrier, + self.num_heads, + self.head_dim, + ) + + attn_output = attn_output.reshape( + batch_size, seq_len, self.num_heads * self.head_dim + ) + attn_output = attn_output.to(dtype=self.wo.weight.dtype) + attn_output = self.wo(attn_output) + return self.residual_dropout(attn_output) + + +class MolmoPluginAttention(nn.Module): + """Plugin wrapper for MolmoAct2 fused attention (att_proj + attn_out). + + Exposes the same runtime ABI as PluginAttention: + (hidden_states, rope_rotary_cos_sin, past_key_value, ctx_len, kvcache_start_index) + -> (attn_output, updated_kv) + """ + + def __init__( + self, + original_attn: nn.Module, + *, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + hidden_size: int, + layer_idx: int, + context_attention_mask_type: int = ContextAttentionMaskType.PADDING, + ): + super().__init__() + self.att_proj = original_attn.att_proj + self.attn_out = original_attn.attn_out + self.q_norm = getattr(original_attn, "q_norm", None) + self.k_norm = getattr(original_attn, "k_norm", None) + self.qk_norm_type = getattr(original_attn, "qk_norm_type", None) + + self.num_heads = int(num_attention_heads) + self.num_kv_heads = int(num_key_value_heads) + self.head_dim = int(head_dim) + self.attn_hidden_size = self.num_heads * self.head_dim + self.hidden_size = int(hidden_size) + self.layer_idx = int(layer_idx) + self.context_attention_mask_type = int(context_attention_mask_type) + + self.q_dim = self.num_heads * self.head_dim + self.k_dim = self.num_kv_heads * self.head_dim + self.v_dim = self.num_kv_heads * self.head_dim + + def forward( + self, + hidden_states: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_value: torch.Tensor | None = None, + ctx_len: torch.Tensor | None = None, + kvcache_start_index: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, seq_len, _ = hidden_states.shape + del attention_mask, position_ids, kwargs + assert rope_rotary_cos_sin.dtype == torch.float32 + + qkv = self.att_proj(hidden_states) + q, k, v = qkv.split([self.q_dim, self.k_dim, self.v_dim], dim=-1) + + # Match MolmoAct2Attention norm ordering. + if ( + self.q_norm is not None + and self.k_norm is not None + and self.qk_norm_type != "qwen3" + ): + q = self.q_norm(q) + k = self.k_norm(k) + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + else: + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + + if ( + self.q_norm is not None + and self.k_norm is not None + and self.qk_norm_type == "qwen3" + ): + q = self.q_norm(q) + k = self.k_norm(k) + + q = q.reshape(batch_size, seq_len, -1) + k = k.reshape(batch_size, seq_len, -1) + v = v.reshape(batch_size, seq_len, -1) + + if ctx_len is None: + ctx_len = torch.full( + (batch_size,), seq_len, device=hidden_states.device, dtype=torch.int32 + ) + if past_key_value is None: + raise ValueError("past_key_value must be provided") + if kvcache_start_index is None: + raise ValueError("kvcache_start_index must be provided") + + dtype = q.dtype + q = q.to(torch.float16) + k = k.to(torch.float16) + v = v.to(torch.float16) + + attn_out, updated_kv = torch.ops.trt.attention_plugin.default( + q, + k, + v, + past_key_value, + ctx_len, + rope_rotary_cos_sin, + kvcache_start_index, + self.num_heads, + self.num_kv_heads, + False, # is_cross_attention + self.head_dim, + False, # do_rotary_embedding (RoPE supplied externally) + -1, + self.context_attention_mask_type, + ) + + attn_out = attn_out.reshape(batch_size, seq_len, self.attn_hidden_size).to( + dtype + ) + return self.attn_out(attn_out), updated_kv + + +""" +Below are reference attention implementations to compare math +with plugin stack and kernel implementations +""" + + +class SiglipReferenceAttention(nn.Module): + """ + Hand-written re-implementation of SigLIP multi-head attention. + + Reuses the original module's projection weights and shape params, but computes + QK^T -> softmax -> (attn @ V) explicitly instead of dispatching through HF's + attention_interface. Used to validate that our understanding of the math + matches the stock eager path bit-for-bit (up to fp accumulation). + + Drop-in replacement for SiglipAttention: same forward signature and same + (attn_output, attn_weights) return contract that SiglipEncoderLayer expects. + """ + + def __init__(self, attn: nn.Module): + super().__init__() + # Reuse the trained projection layers directly (no copy). + self.q_proj = attn.q_proj + self.k_proj = attn.k_proj + self.v_proj = attn.v_proj + self.out_proj = attn.out_proj + + self.num_heads = int(attn.num_heads) + self.head_dim = int(attn.head_dim) + self.embed_dim = self.num_heads * self.head_dim + # SiglipAttention.scale == head_dim ** -0.5 + self.scale = float(getattr(attn, "scale", self.head_dim**-0.5)) + + def forward(self, hidden_states, attention_mask=None, **kwargs): + # hidden_states: [B, S, embed_dim] + input_shape = hidden_states.shape[:-1] # (B, S) + hidden_shape = (*input_shape, self.num_heads, self.head_dim) + + # Project and split heads: [B, S, E] -> [B, num_heads, S, head_dim] + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + # Scores: [B, num_heads, S, S] + attn_weights = torch.matmul(q, k.transpose(-1, -2)) * self.scale + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + # Match HF eager: softmax in fp32 then cast back to input dtype. + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q.dtype) + + # Weighted sum of values: [B, num_heads, S, head_dim] + attn_output = torch.matmul(attn_weights, v) + + # Merge heads back: [B, S, E] + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class PluginNemotronAttention(PluginAttention): + """``PluginAttention`` constructed from a Nemotron HF config.""" + + def __init__(self, original_attn: nn.Module, config, layer_idx: int): + head_dim = int(getattr(config, "head_dim", 0) or original_attn.head_dim) + super().__init__( + original_attn, + num_attention_heads=int(config.num_attention_heads), + num_key_value_heads=int(config.num_key_value_heads), + head_dim=head_dim, + hidden_size=int(config.hidden_size), + layer_idx=layer_idx, + context_attention_mask_type=ContextAttentionMaskType.CAUSAL, + ) diff --git a/py/torch_tensorrt/hf/exporters/plugin/mamba.py b/py/torch_tensorrt/hf/exporters/plugin/mamba.py new file mode 100644 index 00000000000..03af59091c7 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/mamba.py @@ -0,0 +1,330 @@ +"""Nemotron Mamba mixer wrapper and ``torch.ops.trt`` custom ops. + +Mirrors ``attention.py`` / ``plugin_utils._register_attention_plugin_op``: +eager stubs + fake kernels for Dynamo, lowered by ``plugin_converter`` onto +Edge-LLM ``causal_conv1d`` and ``update_ssm_state`` IPluginV3. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +_DT_CLAMP = 50.0 + + +def _has_torch_op(namespace: str, name: str) -> bool: + return hasattr(torch.ops, namespace) and hasattr( + getattr(torch.ops, namespace), name + ) + + +def register_mamba_plugin_ops() -> None: + """Register ``trt::causal_conv1d`` and ``trt::update_ssm_state``.""" + if _has_torch_op("trt", "causal_conv1d") and _has_torch_op( + "trt", "update_ssm_state" + ): + return + + @torch.library.custom_op("trt::causal_conv1d", mutates_args=()) + def causal_conv1d( + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + conv_state: torch.Tensor, + context_lengths: torch.Tensor, + stride: int, + padding: int, + dilation: int, + groups: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del stride, padding, dilation, groups + return _causal_conv1d_eager( + hidden_states, weight, bias, conv_state, context_lengths + ) + + @causal_conv1d.register_fake + def _( + hidden_states, + weight, + bias, + conv_state, + context_lengths, + stride, + padding, + dilation, + groups, + ): + del weight, bias, context_lengths, stride, padding, dilation, groups + return torch.empty_like(hidden_states), torch.empty_like(conv_state) + + @torch.library.custom_op("trt::update_ssm_state", mutates_args=()) + def update_ssm_state( + hidden_states: torch.Tensor, + ssm_a: torch.Tensor, + ssm_b: torch.Tensor, + ssm_c: torch.Tensor, + ssm_d: torch.Tensor, + dt: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, + dt_softplus: int, + ngroups: int, + nheads: int, + head_dim: int, + dstate: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del nheads, head_dim, dstate + return _update_ssm_state_eager( + hidden_states, + ssm_a, + ssm_b, + ssm_c, + ssm_d, + dt, + dt_bias, + state, + context_lengths, + bool(dt_softplus), + int(ngroups), + ) + + @update_ssm_state.register_fake + def _( + hidden_states, + ssm_a, + ssm_b, + ssm_c, + ssm_d, + dt, + dt_bias, + state, + context_lengths, + dt_softplus, + ngroups, + nheads, + head_dim, + dstate, + ): + del ssm_a, ssm_b, ssm_c, ssm_d, dt, dt_bias, context_lengths + del dt_softplus, ngroups, nheads, head_dim, dstate + return torch.empty_like(hidden_states), torch.empty_like(state) + + +def _causal_conv1d_eager( + hidden_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + conv_state: torch.Tensor, + context_lengths: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Depthwise causal conv matching Edge-LLM ``causal_conv1d_ref`` / the plugin. + + ``hidden_states`` [B, S, C], ``weight`` [C, 1, K], ``conv_state`` [B, C, K]. + History is ``conv_state[:, :, 1:]`` (left zero-pad on a fresh state). + """ + batch, seq_len, channels = hidden_states.shape + kernel = int(weight.shape[-1]) + x = hidden_states.float().transpose(1, 2) + history = conv_state.float()[:, :, 1:kernel] + padded = torch.cat([history, x], dim=-1) + y = F.conv1d( + padded, + weight.float(), + None if bias is None else bias.float(), + stride=1, + groups=channels, + ) + new_state = padded[:, :, -kernel:] + if context_lengths is not None: + lengths = context_lengths.to(device=y.device, dtype=torch.long) + token = torch.arange(seq_len, device=y.device) + valid = token.unsqueeze(0) < lengths.unsqueeze(1) + y = y.masked_fill(~valid.unsqueeze(1), 0) + return y.transpose(1, 2).to(dtype=hidden_states.dtype), new_state.to( + dtype=conv_state.dtype + ) + + +def _update_ssm_state_eager( + hidden_states: torch.Tensor, + ssm_a: torch.Tensor, + ssm_b: torch.Tensor, + ssm_c: torch.Tensor, + ssm_d: torch.Tensor, + dt: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, + dt_softplus: bool, + ngroups: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Sequential fp32 selective scan matching Edge-LLM ``selective_scan_ref``.""" + decode = hidden_states.dim() == 3 + if decode: + hidden_states = hidden_states[:, None] + ssm_b = ssm_b[:, None] + ssm_c = ssm_c[:, None] + dt = dt[:, None] + batch, seq_len, nheads, head_dim = hidden_states.shape + heads_per_group = nheads // int(ngroups) + x = hidden_states.float() + a = ssm_a.float() + b = ssm_b.float() + c = ssm_c.float() + d = ssm_d.float() if ssm_d is not None else None + dt = dt.float() + dt_bias = dt_bias.float() if dt_bias is not None else None + if context_lengths is None: + lengths = torch.full((batch,), seq_len, device=x.device, dtype=torch.long) + else: + lengths = context_lengths.to(device=x.device, dtype=torch.long) + + y = torch.zeros( + batch, seq_len, nheads, head_dim, device=x.device, dtype=torch.float32 + ) + st = state.float().clone() + for bi in range(batch): + L = int(lengths[bi]) + layer_state = st[bi] + for t in range(L): + dt_t = dt[bi, t] + if dt_bias is not None: + dt_t = dt_t + dt_bias + if dt_softplus: + dt_t = F.softplus(dt_t) + dA = torch.exp(a * dt_t) + x_t = x[bi, t] + b_h = b[bi, t].repeat_interleave(heads_per_group, dim=0) + c_h = c[bi, t].repeat_interleave(heads_per_group, dim=0) + layer_state = ( + layer_state * dA[:, None, None] + + (dt_t[:, None] * x_t)[:, :, None] * b_h[:, None, :] + ) + y_t = (layer_state * c_h[:, None, :]).sum(-1) + if d is not None: + y_t = y_t + d[:, None] * x_t + y[bi, t] = y_t + st[bi] = layer_state + if decode: + y = y[:, 0] + return y.to(dtype=hidden_states.dtype), st.to(dtype=state.dtype) + + +class PluginNemotronMamba(nn.Module): + """Wrap ``NemotronHMamba2Mixer``: native GEMMs, plugin conv + SSM. + + Export contract (not HF ``cache_params``):: + + hidden, conv_state, ssm_state, context_lengths + -> hidden, conv_state_out, ssm_state_out + """ + + def __init__(self, original: nn.Module): + super().__init__() + self.in_proj = original.in_proj + self.out_proj = original.out_proj + self.conv1d = original.conv1d + self.norm = original.norm + self.A_log = original.A_log + self.D = original.D + self.dt_bias = original.dt_bias + + self.num_heads = int(original.num_heads) + self.head_dim = int(original.head_dim) + self.n_groups = int(original.n_groups) + self.ssm_state_size = int(original.ssm_state_size) + self.conv_dim = int(original.conv_dim) + self.conv_kernel = int( + getattr(original, "conv_kernel_size", original.conv1d.kernel_size[0]) + ) + self.layer_idx = getattr(original, "layer_idx", None) + self._group_size = (self.num_heads * self.head_dim) // self.n_groups + self._eps = float( + getattr( + original.norm, "variance_epsilon", getattr(original.norm, "eps", 1e-5) + ) + ) + + def forward( + self, + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_len, _ = hidden_states.shape + d_inner = self.num_heads * self.head_dim + d_state = self.n_groups * self.ssm_state_size + + projected = self.in_proj(hidden_states) + gate, conv_in, dt = projected.split( + [d_inner, self.conv_dim, self.num_heads], dim=-1 + ) + dt = dt.clamp(-_DT_CLAMP, _DT_CLAMP) + + conv_bias = self.conv1d.bias + if conv_bias is None: + conv_bias = torch.zeros( + self.conv_dim, device=conv_in.device, dtype=conv_in.dtype + ) + + conv_out, conv_state_out = torch.ops.trt.causal_conv1d.default( + conv_in, + self.conv1d.weight, + conv_bias, + conv_state, + context_lengths, + 1, + self.conv_kernel - 1, + 1, + self.conv_dim, + ) + conv_out = F.silu(conv_out) + + ssm_input, ssm_b, ssm_c = conv_out.split([d_inner, d_state, d_state], dim=-1) + ssm_input = ssm_input.view(batch_size, seq_len, self.num_heads, self.head_dim) + ssm_b = ssm_b.view(batch_size, seq_len, self.n_groups, self.ssm_state_size) + ssm_c = ssm_c.view(batch_size, seq_len, self.n_groups, self.ssm_state_size) + + ssm_a = -torch.exp(self.A_log.to(torch.float32)) + ssm_out, ssm_state_out = torch.ops.trt.update_ssm_state.default( + ssm_input, + ssm_a, + ssm_b, + ssm_c, + self.D.to(torch.float16), + dt, + self.dt_bias.to(torch.float16), + ssm_state, + context_lengths, + 1, + self.n_groups, + self.num_heads, + self.head_dim, + self.ssm_state_size, + ) + ssm_out = ssm_out.view(batch_size, seq_len, d_inner) + # HF MambaRMSNormGated(x, gate) is gate-then-norm (norm_before_gate=False). + if ( + getattr(self.norm, "forward", None) is not None + and self.norm.__class__.__name__ == "MambaRMSNormGated" + ): + normed = self.norm(ssm_out, gate) + else: + normed = self._gated_rmsnorm(ssm_out, gate) + return self.out_proj(normed), conv_state_out, ssm_state_out + + def _gated_rmsnorm( + self, hidden_states: torch.Tensor, gate: torch.Tensor + ) -> torch.Tensor: + dtype = hidden_states.dtype + gated = (hidden_states * F.silu(gate)).float() + grouped = gated.view(*gated.shape[:-1], -1, self._group_size) + variance = (grouped * grouped).mean(-1, keepdim=True) + normed = grouped * torch.rsqrt(variance + self._eps) + return (normed.view(*hidden_states.shape) * self.norm.weight.float()).to(dtype) diff --git a/py/torch_tensorrt/hf/exporters/plugin/moe.py b/py/torch_tensorrt/hf/exporters/plugin/moe.py new file mode 100644 index 00000000000..bab850b0417 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/moe.py @@ -0,0 +1,423 @@ +"""Nemotron MoE mixer wrapper and ``torch.ops.trt`` NVFP4 MoE custom ops. + +Same pattern as ``attention.py``: eager stub + fake for Dynamo, converter +inserts Edge-LLM ``Nvfp4MoePlugin`` / ``NvFP4MoEPluginGeforce``. +""" + +from __future__ import annotations + +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F + +_NVFP4_ACTIVATION_RELU2 = 4 +_NVFP4_ROUTING_MODE_SIGMOID_GROUP_TOPK = 1 +_NVFP4_MOE_BACKEND_AUTO = 0 +_NVFP4_MOE_IO_DTYPE_FP16 = 1 +_NVFP4_MOE_MAX_ROUTED_ROWS_AUTO = 0 + +_NVFP4_MOE_TARGET_ENV = "EDGELLM_NVFP4_MOE_TARGET" +_NVFP4_MOE_SM12X_ALIASES = frozenset(("sm12x", "sm120", "sm121", "geforce")) +_NVFP4_MOE_SM110_ALIASES = frozenset( + ("sm100", "sm101", "sm110", "blackwell_dc", "thor", "") +) + + +def _has_torch_op(namespace: str, name: str) -> bool: + return hasattr(torch.ops, namespace) and hasattr( + getattr(torch.ops, namespace), name + ) + + +def use_geforce_nvfp4_moe() -> bool: + """True when exporting ``NvFP4MoEPluginGeforce`` (SM12x). Default is SM110.""" + val = os.environ.get(_NVFP4_MOE_TARGET_ENV, "sm110").strip().lower() + if val in _NVFP4_MOE_SM12X_ALIASES: + return True + if val in _NVFP4_MOE_SM110_ALIASES: + return False + raise ValueError( + f"{_NVFP4_MOE_TARGET_ENV}={val!r} is not recognized. " + "Use sm100/sm110 (Nvfp4MoePlugin) or sm12x (NvFP4MoEPluginGeforce)." + ) + + +def _nvfp4_moe_stub( + router_logits: torch.Tensor, + hidden_states: torch.Tensor, + fc1_qweights: torch.Tensor, + fc1_blocks_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_qweights: torch.Tensor, + fc2_blocks_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + input_global_scale: torch.Tensor, + down_input_scale: torch.Tensor, + e_score_correction_bias: torch.Tensor, + num_experts: int, + top_k: int, + hidden_size: int, + moe_inter_size: int, + activation_type: int, + n_group: int, + topk_group: int, + norm_topk_prob: int, + routed_scaling_factor: float, + routing_mode: int, + backend: int, + io_dtype: int, + max_routed_rows: int, +) -> torch.Tensor: + del router_logits, fc1_qweights, fc1_blocks_scale, fc1_alpha + del fc2_qweights, fc2_blocks_scale, fc2_alpha + del input_global_scale, down_input_scale, e_score_correction_bias + del num_experts, top_k, hidden_size, moe_inter_size, activation_type + del n_group, topk_group, norm_topk_prob, routed_scaling_factor + del routing_mode, backend, io_dtype, max_routed_rows + return torch.zeros_like(hidden_states) + + +def register_moe_plugin_ops() -> None: + """Register ``trt::nvfp4_moe_plugin`` and ``trt::nvfp4_moe_plugin_geforce``.""" + if _has_torch_op("trt", "nvfp4_moe_plugin"): + return + + @torch.library.custom_op("trt::nvfp4_moe_plugin", mutates_args=()) + def nvfp4_moe_plugin( + router_logits: torch.Tensor, + hidden_states: torch.Tensor, + fc1_qweights: torch.Tensor, + fc1_blocks_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_qweights: torch.Tensor, + fc2_blocks_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + input_global_scale: torch.Tensor, + down_input_scale: torch.Tensor, + e_score_correction_bias: torch.Tensor, + num_experts: int, + top_k: int, + hidden_size: int, + moe_inter_size: int, + activation_type: int, + n_group: int, + topk_group: int, + norm_topk_prob: int, + routed_scaling_factor: float, + routing_mode: int, + backend: int, + io_dtype: int, + max_routed_rows: int, + ) -> torch.Tensor: + return _nvfp4_moe_stub( + router_logits, + hidden_states, + fc1_qweights, + fc1_blocks_scale, + fc1_alpha, + fc2_qweights, + fc2_blocks_scale, + fc2_alpha, + input_global_scale, + down_input_scale, + e_score_correction_bias, + num_experts, + top_k, + hidden_size, + moe_inter_size, + activation_type, + n_group, + topk_group, + norm_topk_prob, + routed_scaling_factor, + routing_mode, + backend, + io_dtype, + max_routed_rows, + ) + + @nvfp4_moe_plugin.register_fake + def _( + router_logits, + hidden_states, + fc1_qweights, + fc1_blocks_scale, + fc1_alpha, + fc2_qweights, + fc2_blocks_scale, + fc2_alpha, + input_global_scale, + down_input_scale, + e_score_correction_bias, + num_experts, + top_k, + hidden_size, + moe_inter_size, + activation_type, + n_group, + topk_group, + norm_topk_prob, + routed_scaling_factor, + routing_mode, + backend, + io_dtype, + max_routed_rows, + ): + del router_logits, fc1_qweights, fc1_blocks_scale, fc1_alpha + del fc2_qweights, fc2_blocks_scale, fc2_alpha + del input_global_scale, down_input_scale, e_score_correction_bias + del num_experts, top_k, hidden_size, moe_inter_size, activation_type + del n_group, topk_group, norm_topk_prob, routed_scaling_factor + del routing_mode, backend, io_dtype, max_routed_rows + return torch.empty_like(hidden_states) + + @torch.library.custom_op("trt::nvfp4_moe_plugin_geforce", mutates_args=()) + def nvfp4_moe_plugin_geforce( + router_logits: torch.Tensor, + hidden_states: torch.Tensor, + fc1_qweights: torch.Tensor, + fc1_blocks_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_qweights: torch.Tensor, + fc2_blocks_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + input_global_scale: torch.Tensor, + down_input_scale: torch.Tensor, + e_score_correction_bias: torch.Tensor, + num_experts: int, + top_k: int, + hidden_size: int, + moe_inter_size: int, + activation_type: int, + n_group: int, + topk_group: int, + norm_topk_prob: int, + routed_scaling_factor: float, + routing_mode: int, + backend: int, + io_dtype: int, + max_routed_rows: int, + ) -> torch.Tensor: + return _nvfp4_moe_stub( + router_logits, + hidden_states, + fc1_qweights, + fc1_blocks_scale, + fc1_alpha, + fc2_qweights, + fc2_blocks_scale, + fc2_alpha, + input_global_scale, + down_input_scale, + e_score_correction_bias, + num_experts, + top_k, + hidden_size, + moe_inter_size, + activation_type, + n_group, + topk_group, + norm_topk_prob, + routed_scaling_factor, + routing_mode, + backend, + io_dtype, + max_routed_rows, + ) + + @nvfp4_moe_plugin_geforce.register_fake + def _( + router_logits, + hidden_states, + fc1_qweights, + fc1_blocks_scale, + fc1_alpha, + fc2_qweights, + fc2_blocks_scale, + fc2_alpha, + input_global_scale, + down_input_scale, + e_score_correction_bias, + num_experts, + top_k, + hidden_size, + moe_inter_size, + activation_type, + n_group, + topk_group, + norm_topk_prob, + routed_scaling_factor, + routing_mode, + backend, + io_dtype, + max_routed_rows, + ): + del router_logits, fc1_qweights, fc1_blocks_scale, fc1_alpha + del fc2_qweights, fc2_blocks_scale, fc2_alpha + del input_global_scale, down_input_scale, e_score_correction_bias + del num_experts, top_k, hidden_size, moe_inter_size, activation_type + del n_group, topk_group, norm_topk_prob, routed_scaling_factor + del routing_mode, backend, io_dtype, max_routed_rows + return torch.empty_like(hidden_states) + + +class PluginNemotronMoE(nn.Module): + """Wrap ``NemotronHMoE``: native router + shared expert, plugin routed experts. + + Requires packed NVFP4 buffers (``prepare_for_export`` or an already-packed + NVFP4 checkpoint). ReLU2 + sigmoid-group top-k, matching Nemotron-3-30B-A3B. + """ + + def __init__(self, original: nn.Module, config): + super().__init__() + self.gate = original.gate + self.shared_experts = original.shared_experts + self.fc1_latent_proj = getattr(original, "fc1_latent_proj", nn.Identity()) + self.fc2_latent_proj = getattr(original, "fc2_latent_proj", nn.Identity()) + self._hf_experts = original.experts + + self.n_routed_experts = int(config.n_routed_experts) + self.num_experts_per_tok = int(config.num_experts_per_tok) + self.hidden_size = int(config.hidden_size) + self.routed_hidden_size = int( + getattr(config, "moe_latent_size", None) or config.hidden_size + ) + self.moe_intermediate_size = int(config.moe_intermediate_size) + self.group_size = int( + getattr(getattr(config, "quant", None), "group_size", 16) or 16 + ) + + self.n_group = int( + getattr(self.gate, "n_group", getattr(self.gate, "num_group", 1)) + ) + self.topk_group = int(self.gate.topk_group) + self.norm_topk_prob = int(bool(self.gate.norm_topk_prob)) + self.routed_scaling_factor = float(self.gate.routed_scaling_factor) + + self._padded_hidden_size = self.routed_hidden_size + self._padded_moe_intermediate_size = self.moe_intermediate_size + self._export_ready = False + + if hasattr(original, "fc1_qweights"): + self.fc1_qweights = original.fc1_qweights + self.fc1_blocks_scale = original.fc1_blocks_scale + self.fc1_alpha = original.fc1_alpha + self.fc2_qweights = original.fc2_qweights + self.fc2_blocks_scale = original.fc2_blocks_scale + self.fc2_alpha = original.fc2_alpha + self.input_global_scale = original.input_global_scale + self.down_input_scale = original.down_input_scale + self._e_score_correction_bias_fp32 = original._e_score_correction_bias_fp32 + self._padded_hidden_size = int(original._padded_hidden_size) + self._padded_moe_intermediate_size = int( + original._padded_moe_intermediate_size + ) + self._export_ready = True + + def prepare_for_export(self) -> None: + from tensorrt_edgellm.checkpoint.repacking import repack_nvfp4_moe_experts + + experts = self._hf_experts + if not isinstance(experts, nn.ModuleList): + raise TypeError( + "HF NemotronHExperts stores 3D fp16 tensors, not per-expert " + "NVFP4 Linears. Load NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 " + "(or run ModelOpt NVFP4 quant) before the MoE plugin path." + ) + + hidden_align = 256 if use_geforce_nvfp4_moe() else 1 + ( + fc1_q, + fc1_scale, + fc1_alpha, + fc2_q, + fc2_scale, + fc2_alpha, + padded_inter, + padded_h, + ) = repack_nvfp4_moe_experts( + experts, + self.routed_hidden_size, + self.moe_intermediate_size, + self.group_size, + hidden_size_alignment=hidden_align, + ) + device = self.gate.weight.device + self.register_buffer("fc1_qweights", fc1_q.to(device).contiguous()) + self.register_buffer("fc1_blocks_scale", fc1_scale.to(device).contiguous()) + self.register_buffer("fc1_alpha", fc1_alpha.to(device).contiguous()) + self.register_buffer("fc2_qweights", fc2_q.to(device).contiguous()) + self.register_buffer("fc2_blocks_scale", fc2_scale.to(device).contiguous()) + self.register_buffer("fc2_alpha", fc2_alpha.to(device).contiguous()) + self.register_buffer( + "input_global_scale", + torch.ones(self.n_routed_experts, dtype=torch.float32, device=device), + ) + self.register_buffer( + "down_input_scale", + torch.ones(self.n_routed_experts, dtype=torch.float32, device=device), + ) + self.register_buffer( + "_e_score_correction_bias_fp32", + self.gate.e_score_correction_bias.data.to(torch.float32).to(device), + ) + self._padded_moe_intermediate_size = padded_inter + self._padded_hidden_size = padded_h + self._export_ready = True + + def _shared_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + h = self.shared_experts.up_proj(hidden_states) + r = F.relu(h) + return self.shared_experts.down_proj(r * r) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if not self._export_ready: + raise RuntimeError("PluginNemotronMoE.prepare_for_export() was not called") + + router_logits = F.linear( + hidden_states.view(-1, self.hidden_size), self.gate.weight + ).float() + routed = self.fc1_latent_proj(hidden_states) + plugin_hidden = routed + if self._padded_hidden_size != self.routed_hidden_size: + plugin_hidden = F.pad( + routed, (0, self._padded_hidden_size - self.routed_hidden_size) + ) + + moe_op = ( + torch.ops.trt.nvfp4_moe_plugin_geforce.default + if use_geforce_nvfp4_moe() + else torch.ops.trt.nvfp4_moe_plugin.default + ) + moe_out = moe_op( + router_logits, + plugin_hidden, + self.fc1_qweights, + self.fc1_blocks_scale, + self.fc1_alpha, + self.fc2_qweights, + self.fc2_blocks_scale, + self.fc2_alpha, + self.input_global_scale, + self.down_input_scale, + self._e_score_correction_bias_fp32, + self.n_routed_experts, + self.num_experts_per_tok, + self._padded_hidden_size, + self._padded_moe_intermediate_size, + _NVFP4_ACTIVATION_RELU2, + self.n_group, + self.topk_group, + self.norm_topk_prob, + self.routed_scaling_factor, + _NVFP4_ROUTING_MODE_SIGMOID_GROUP_TOPK, + _NVFP4_MOE_BACKEND_AUTO, + _NVFP4_MOE_IO_DTYPE_FP16, + _NVFP4_MOE_MAX_ROUTED_ROWS_AUTO, + ) + if self._padded_hidden_size != self.routed_hidden_size: + moe_out = moe_out[..., : self.routed_hidden_size] + moe_out = self.fc2_latent_proj(moe_out) + return moe_out + self._shared_forward(hidden_states) diff --git a/py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py b/py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py new file mode 100644 index 00000000000..efb6db8ebdd --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py @@ -0,0 +1,366 @@ +""" +TensorRT converters for Edge-LLM plugin custom ops. + +Attention: separate Q/K/V into AttentionPlugin (no fused-qkv slice), plus +``context_attention_mask_type``. Also lowers ``trt::causal_conv1d``, +``trt::update_ssm_state``, and ``trt::nvfp4_moe_plugin`` onto the matching +IPluginV3 creators from libNvInfer_edgellm_plugin.so. +""" + +import numpy as np +import tensorrt as trt +import torch +from torch_tensorrt.dynamo.conversion import ( + ConversionContext, + dynamo_tensorrt_converter, +) +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ConverterPriority +from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor + +from .attention import ContextAttentionMaskType +from .plugin_utils import get_trt_plugin_creator + + +def _creator_is_v3(creator) -> bool: + return "V3" in type(creator).__name__ + + +def _create_trt_plugin( + creator, name: str, field_list: list +) -> trt.IPluginV2 | trt.IPluginV3: + fields = trt.PluginFieldCollection(field_list) + if _creator_is_v3(creator): + return creator.create_plugin(name, fields, trt.TensorRTPhase.BUILD) + return creator.create_plugin(name, fields) + + +def _plugin_is_v3(plugin) -> bool: + return "V3" in type(plugin).__name__ + + +def _add_plugin_layer(ctx: ConversionContext, inputs: list, plugin, name: str): + layer = ( + ctx.net.add_plugin_v3(inputs, [], plugin) + if _plugin_is_v3(plugin) + else ctx.net.add_plugin_v2(inputs, plugin) + ) + layer.name = name + return layer + + +@dynamo_tensorrt_converter( + torch.ops.trt.attention_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_llm_attention_plugin(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + q, k, v, kv, ctx_len, rope, kv_cache_start_idx = args[:7] + num_q_heads = args[7] + num_kv_heads = args[8] + enable_tree_attention = args[9] + head_size = args[10] + enable_fp8_kv_cache = args[11] + sliding_window_size = args[12] if len(args) > 12 else -1 + context_attention_mask_type = ( + int(args[13]) if len(args) > 13 else int(ContextAttentionMaskType.CAUSAL) + ) + attention_mask = args[14] if len(args) > 14 else None + position_ids = args[15] if len(args) > 15 else None + qkv_scales = args[16] if len(args) > 16 else None + + creator = get_trt_plugin_creator("AttentionPlugin", "1", "") + if creator is None: + raise RuntimeError("AttentionPlugin not found in TensorRT plugin registry") + + field_list = [ + trt.PluginField( + field_name, + np.array([field_val], dtype=np.int32), + trt.PluginFieldType.INT32, + ) + for field_name, field_val in [ + ("num_q_heads", int(num_q_heads)), + ("num_kv_heads", int(num_kv_heads)), + ("head_size", int(head_size)), + ("enable_tree_attention", int(enable_tree_attention)), + ("enable_fp8_kv_cache", int(enable_fp8_kv_cache)), + ("sliding_window_size", int(sliding_window_size)), + ("context_attention_mask_type", context_attention_mask_type), + ] + ] + if bool(enable_fp8_kv_cache) and qkv_scales is not None: + field_list.append( + trt.PluginField( + "qkv_scales", + np.array(list(qkv_scales), dtype=np.float32), + trt.PluginFieldType.FLOAT32, + ) + ) + + plugin = _create_trt_plugin(creator, name, field_list) + if plugin is None: + raise RuntimeError("Failed to create AttentionPlugin") + + plugin_inputs = [q, k, v, kv, ctx_len, rope, kv_cache_start_idx] + if bool(enable_tree_attention): + plugin_inputs.extend([attention_mask, position_ids]) + + inputs = [ + ( + get_trt_tensor(ctx, tensor, f"{name}_i{idx}") + if not isinstance(tensor, trt.ITensor) + else tensor + ) + for idx, tensor in enumerate(plugin_inputs) + ] + + kv_cache_start_idx_input_idx = 6 + if ( + len(inputs[kv_cache_start_idx_input_idx].shape) == 2 + and inputs[kv_cache_start_idx_input_idx].shape[1] == 1 + ): + shuffle_layer = ctx.net.add_shuffle(inputs[kv_cache_start_idx_input_idx]) + shuffle_layer.reshape_dims = (inputs[kv_cache_start_idx_input_idx].shape[0],) + inputs[kv_cache_start_idx_input_idx] = shuffle_layer.get_output(0) + + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0), layer.get_output(1) + + +@dynamo_tensorrt_converter( + torch.ops.trt.vit_attention_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_vit_attention_plugin(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + q, k, v, cu_seqlens, max_seqlen_carrier = args[:5] + num_heads = args[5] + head_size = args[6] + + creator = get_trt_plugin_creator("ViTAttentionPlugin", "1", "") + if creator is None: + raise RuntimeError("ViTAttentionPlugin not found in TensorRT plugin registry") + + field_list = [ + trt.PluginField( + "num_heads", + np.array([int(num_heads)], dtype=np.int32), + trt.PluginFieldType.INT32, + ), + trt.PluginField( + "head_size", + np.array([int(head_size)], dtype=np.int32), + trt.PluginFieldType.INT32, + ), + ] + plugin = _create_trt_plugin(creator, name, field_list) + if plugin is None: + raise RuntimeError("Failed to create ViTAttentionPlugin") + + inputs = [] + for idx, tensor in enumerate([q, k, v, cu_seqlens, max_seqlen_carrier]): + tensor_name = f"{name}_i{idx}" + trt_tensor = ( + get_trt_tensor(ctx, tensor, tensor_name) + if not isinstance(tensor, trt.ITensor) + else tensor + ) + if not trt_tensor.name: + trt_tensor.name = tensor_name + inputs.append(trt_tensor) + + layer = _add_plugin_layer(ctx, inputs, plugin, name) + output = layer.get_output(0) + if not output.name: + output.name = f"{name}_output" + return output + + +def _int_field(name: str, value: int) -> trt.PluginField: + return trt.PluginField( + name, np.array([int(value)], dtype=np.int32), trt.PluginFieldType.INT32 + ) + + +def _float_field(name: str, value: float) -> trt.PluginField: + return trt.PluginField( + name, np.array([float(value)], dtype=np.float32), trt.PluginFieldType.FLOAT32 + ) + + +def _as_plugin_inputs(ctx: ConversionContext, tensors: list, name: str) -> list: + inputs = [] + for idx, tensor in enumerate(tensors): + tensor_name = f"{name}_i{idx}" + trt_tensor = ( + get_trt_tensor(ctx, tensor, tensor_name) + if not isinstance(tensor, trt.ITensor) + else tensor + ) + if not getattr(trt_tensor, "name", None): + trt_tensor.name = tensor_name + inputs.append(trt_tensor) + return inputs + + +@dynamo_tensorrt_converter( + torch.ops.trt.causal_conv1d.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_causal_conv1d(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + hidden_states, weight, bias, conv_state, context_lengths = args[:5] + stride, padding, dilation, groups = args[5:9] + + creator = get_trt_plugin_creator("causal_conv1d", "1", "") + if creator is None: + raise RuntimeError("causal_conv1d plugin not found in TensorRT plugin registry") + + plugin = _create_trt_plugin( + creator, + name, + [ + _int_field("stride", stride), + _int_field("padding", padding), + _int_field("dilation", dilation), + _int_field("groups", groups), + _int_field("use_mtp", 0), + _int_field("use_ddtree", 0), + ], + ) + if plugin is None: + raise RuntimeError("Failed to create causal_conv1d plugin") + + inputs = _as_plugin_inputs( + ctx, [hidden_states, weight, bias, conv_state, context_lengths], name + ) + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0), layer.get_output(1) + + +@dynamo_tensorrt_converter( + torch.ops.trt.update_ssm_state.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_update_ssm_state(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + hidden_states, ssm_a, ssm_b, ssm_c, ssm_d, dt, dt_bias, state, context_lengths = ( + args[:9] + ) + dt_softplus, ngroups, nheads, head_dim, dstate = args[9:14] + + creator = get_trt_plugin_creator("update_ssm_state", "1", "") + if creator is None: + raise RuntimeError( + "update_ssm_state plugin not found in TensorRT plugin registry" + ) + + plugin = _create_trt_plugin( + creator, + name, + [ + _int_field("dim", head_dim), + _int_field("dstate", dstate), + _int_field("nheads", nheads), + _int_field("ngroups", ngroups), + _int_field("dt_softplus", dt_softplus), + ], + ) + if plugin is None: + raise RuntimeError("Failed to create update_ssm_state plugin") + + inputs = _as_plugin_inputs( + ctx, + [ + hidden_states, + ssm_a, + ssm_b, + ssm_c, + ssm_d, + dt, + dt_bias, + state, + context_lengths, + ], + name, + ) + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0), layer.get_output(1) + + +def _convert_nvfp4_moe(ctx: ConversionContext, args, name: str, plugin_name: str): + args = list(args) + tensors = args[:11] + num_experts = args[11] + top_k = args[12] + hidden_size = args[13] + moe_inter_size = args[14] + activation_type = args[15] + n_group = args[16] + topk_group = args[17] + norm_topk_prob = args[18] + routed_scaling_factor = args[19] + routing_mode = args[20] + backend = args[21] + io_dtype = args[22] + max_routed_rows = args[23] + + creator = get_trt_plugin_creator(plugin_name, "1", "") + if creator is None: + raise RuntimeError(f"{plugin_name} not found in TensorRT plugin registry") + + plugin = _create_trt_plugin( + creator, + name, + [ + _int_field("num_experts", num_experts), + _int_field("top_k", top_k), + _int_field("hidden_size", hidden_size), + _int_field("moe_inter_size", moe_inter_size), + _int_field("activation_type", activation_type), + _int_field("n_group", n_group), + _int_field("topk_group", topk_group), + _int_field("norm_topk_prob", norm_topk_prob), + _float_field("routed_scaling_factor", routed_scaling_factor), + _int_field("routing_mode", routing_mode), + _int_field("backend", backend), + _int_field("io_dtype", io_dtype), + _int_field("max_routed_rows", max_routed_rows), + ], + ) + if plugin is None: + raise RuntimeError(f"Failed to create {plugin_name}") + + inputs = _as_plugin_inputs(ctx, tensors, name) + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0) + + +@dynamo_tensorrt_converter( + torch.ops.trt.nvfp4_moe_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_nvfp4_moe_plugin(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + return _convert_nvfp4_moe(ctx, args, name, "Nvfp4MoePlugin") + + +@dynamo_tensorrt_converter( + torch.ops.trt.nvfp4_moe_plugin_geforce.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_nvfp4_moe_plugin_geforce( + ctx: ConversionContext, target, args, kwargs, name +): + del target, kwargs + return _convert_nvfp4_moe(ctx, args, name, "NvFP4MoEPluginGeforce") diff --git a/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py b/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py new file mode 100644 index 00000000000..522003330c6 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py @@ -0,0 +1,519 @@ +import ctypes +import os +from typing import Any, List, Optional, Sequence, Tuple + +import tensorrt as trt +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .attention import ( + ContextAttentionMaskType, + MolmoPluginAttention, + MolmoViTPluginAttention, + PluginAttention, + PluginNemotronAttention, + SiglipReferenceAttention, + ViTPluginAttention, +) + +_PLUGIN_CONFIG: dict[str, Any] = {} + + +def get_plugin_config() -> dict[str, Any]: + return _PLUGIN_CONFIG.copy() + + +def set_plugin_config( + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + max_seq_len: int = 2048, + max_batch_size: int = 4, +) -> None: + """Store LM plugin metadata used by Alpamayo plugin compile helpers.""" + global _PLUGIN_CONFIG + _PLUGIN_CONFIG = { + "num_attention_heads": int(num_attention_heads), + "num_key_value_heads": int(num_key_value_heads), + "head_dim": int(head_dim), + "max_seq_len": int(max_seq_len), + "max_batch_size": int(max_batch_size), + } + + +def set_plugin_config_from_model(model_config: Any, max_seq_len: int = 2048) -> None: + """Populate plugin config from a HuggingFace-style model config.""" + if getattr(model_config, "head_dim", None) is not None: + head_dim = int(model_config.head_dim) + else: + head_dim = int(model_config.hidden_size) // int( + model_config.num_attention_heads + ) + set_plugin_config( + num_attention_heads=int(model_config.num_attention_heads), + num_key_value_heads=int(model_config.num_key_value_heads), + head_dim=head_dim, + max_seq_len=int(max_seq_len), + ) + + +def create_kv_caches( + config: Any, + max_seq_len: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype = torch.float16, +) -> List[torch.Tensor]: + """Allocate empty per-layer KV caches ``[B, 2, n_kv, capacity, head_dim]``.""" + num_layers = int(config.num_hidden_layers) + num_kv_heads = int(config.num_key_value_heads) + if getattr(config, "head_dim", None) is not None: + head_dim = int(config.head_dim) + else: + head_dim = int(config.hidden_size) // int(config.num_attention_heads) + return [ + torch.zeros( + int(batch_size), + 2, + num_kv_heads, + int(max_seq_len), + head_dim, + dtype=dtype, + device=device, + ) + for _ in range(num_layers) + ] + + +def _has_torch_op(namespace: str, name: str) -> bool: + return hasattr(torch.ops, namespace) and hasattr( + getattr(torch.ops, namespace), name + ) + + +def _apply_plugin_rope( + q: torch.Tensor, + k: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + head_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply AttentionPlugin cos/sin layout: first half cos, second half sin.""" + half = head_size // 2 + seq_len = q.shape[1] + cos = rope_rotary_cos_sin[:, :seq_len, :half] + sin = rope_rotary_cos_sin[:, :seq_len, half:] + cos = torch.cat([cos, cos], dim=-1).unsqueeze(2) + sin = torch.cat([sin, sin], dim=-1).unsqueeze(2) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., :half] + x2 = x[..., half:] + return torch.cat((-x2, x1), dim=-1) + + return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin + + +def _attention_plugin_eager( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + past_key_value: torch.Tensor, + context_lengths: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + kvcache_start_index: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + head_size: int, + context_attention_mask_type: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Eager SDPA stand-in for ``AttentionPlugin`` (prefill, linear KV).""" + del kvcache_start_index + orig_dtype = q.dtype + batch, seq_len, _ = q.shape + q = q.view(batch, seq_len, num_q_heads, head_size) + k = k.view(batch, seq_len, num_kv_heads, head_size) + v = v.view(batch, seq_len, num_kv_heads, head_size) + q, k = _apply_plugin_rope(q, k, rope_rotary_cos_sin.float(), head_size) + q = q.to(dtype=orig_dtype) + k = k.to(dtype=orig_dtype) + v = v.to(dtype=orig_dtype) + + present = past_key_value.clone() + present[:, 0, :, :seq_len, :] = k.permute(0, 2, 1, 3).to(dtype=present.dtype) + present[:, 1, :, :seq_len, :] = v.permute(0, 2, 1, 3).to(dtype=present.dtype) + + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + if num_q_heads != num_kv_heads: + repeats = num_q_heads // num_kv_heads + k = k.repeat_interleave(repeats, dim=1) + v = v.repeat_interleave(repeats, dim=1) + + is_causal = int(context_attention_mask_type) == int(ContextAttentionMaskType.CAUSAL) + attn = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal and seq_len > 1) + attn = attn.permute(0, 2, 1, 3).contiguous() + if context_lengths is not None: + lengths = context_lengths.to(device=attn.device, dtype=torch.long) + token = torch.arange(seq_len, device=attn.device) + valid = token.unsqueeze(0) < lengths.unsqueeze(1) + attn = attn.masked_fill(~valid[:, :, None, None], 0) + return attn.to(dtype=q.dtype), present + + +# These registrations follow the same custom-op pattern as TensorRT-Edge-LLM's +# ONNX exporter: define a torch.ops.trt operator and register a fake +# implementation for Dynamo shape propagation. The pipelines diverge only +# after capture: Edge-LLM translates the op into a custom ONNX node, while +# Torch-TensorRT lowers it directly through a Dynamo converter. +def _register_attention_plugin_op() -> None: + """Register the LLM attention op using Edge-LLM's ONNX-export pattern.""" + # TODO: Reuse TensorRT-Edge-LLM's canonical attention custom-op registration + # once this wrapper and its converter use the same argument order and + # present-KV output shape. Preserve context_attention_mask_type, which is + # required by VLA models but is not currently exposed by Edge-LLM's schema. + if _has_torch_op("trt", "attention_plugin"): + return + + @torch.library.custom_op("trt::attention_plugin", mutates_args=()) + def attention_plugin( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + past_key_value: torch.Tensor, + context_lengths: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + kvcache_start_index: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + enable_tree_attention: bool, + head_size: int, + enable_fp8_kv_cache: bool, + sliding_window_size: int = -1, + context_attention_mask_type: int = ContextAttentionMaskType.CAUSAL, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + qkv_scales: Optional[Sequence[float]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del enable_tree_attention, enable_fp8_kv_cache, sliding_window_size + del attention_mask, position_ids, qkv_scales + return _attention_plugin_eager( + q, + k, + v, + past_key_value, + context_lengths, + rope_rotary_cos_sin, + kvcache_start_index, + int(num_q_heads), + int(num_kv_heads), + int(head_size), + int(context_attention_mask_type), + ) + + @attention_plugin.register_fake + def _attention_plugin_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + past_key_value: torch.Tensor, + context_lengths: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + kvcache_start_index: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + enable_tree_attention: bool, + head_size: int, + enable_fp8_kv_cache: bool, + sliding_window_size: int = -1, + context_attention_mask_type: int = ContextAttentionMaskType.CAUSAL, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + qkv_scales: Optional[Sequence[float]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del k, v, context_lengths, rope_rotary_cos_sin, kvcache_start_index + del num_kv_heads, enable_tree_attention, enable_fp8_kv_cache + del ( + sliding_window_size, + context_attention_mask_type, + attention_mask, + position_ids, + qkv_scales, + ) + batch_size, seq_len, _ = q.shape + attn_output = torch.empty( + batch_size, + seq_len, + num_q_heads, + head_size, + dtype=q.dtype, + device=q.device, + ) + return attn_output, torch.empty_like(past_key_value) + + +def _register_vit_attention_plugin_op() -> None: + """Register the same ViT attention op and fake used by Edge-LLM ONNX export.""" + if _has_torch_op("trt", "vit_attention_plugin"): + return + + @torch.library.custom_op("trt::vit_attention_plugin", mutates_args=()) + def vit_attention_plugin( + query_states: torch.Tensor, # [T, num_heads, head_size] + key_states: torch.Tensor, # [T, num_heads, head_size] + value_states: torch.Tensor, # [T, num_heads, head_size] + cu_seqlens: torch.Tensor, # [batch+1] int32 + max_seqlen_carrier: torch.Tensor, # [] or [1] int32 (scalar) + num_heads: int, + head_size: int, + ) -> torch.Tensor: + """ViT ragged self-attention. + + In eager mode, implements varlen SDPA using cu_seqlens to process each + sequence segment independently. During dynamo/ONNX tracing the + register_fake shape propagation is used and this body is not executed. + + Unlike AttentionPlugin, ViT attention has no KV cache and takes ragged + input with cu_seqlens instead of context_lengths. RoPE is applied before + this call. + """ + import torch.nn.functional as F + + out = torch.empty_like(query_states) + seqlens = cu_seqlens.tolist() + for i in range(len(seqlens) - 1): + start, end = int(seqlens[i]), int(seqlens[i + 1]) + if start >= end: + continue + # q/k/v: [S, H, D] -> [1, H, S, D] for SDPA + q = query_states[start:end].permute(1, 0, 2).unsqueeze(0) + k = key_states[start:end].permute(1, 0, 2).unsqueeze(0) + v = value_states[start:end].permute(1, 0, 2).unsqueeze(0) + attn = F.scaled_dot_product_attention(q, k, v) # [1, H, S, D] + out[start:end] = attn.squeeze(0).permute(1, 0, 2) + return out + + @vit_attention_plugin.register_fake + def _( + query_states, + key_states, + value_states, + cu_seqlens, + max_seqlen_carrier, + num_heads, + head_size, + ): + return torch.empty_like(query_states) + + +def get_trt_plugin_creator( + plugin_name: str, + version: str = "1", + namespace: str = "", +): + """Return a TensorRT plugin creator, preferring the TRT 10.14+ V3 API.""" + registry = trt.get_plugin_registry() + if hasattr(registry, "get_creator"): + creator = registry.get_creator(plugin_name, version, namespace) + if creator is not None: + return creator + if hasattr(registry, "get_plugin_creator"): + return registry.get_plugin_creator(plugin_name, version, namespace) + return None + + +def load_plugin(): + plugin_so = ( + os.environ.get("EDGE_LLM_PLUGIN_SO") + or os.environ.get("EDGELLM_TRT_PLUGIN_SO") + or os.environ.get("EDGELLM_PLUGIN_PATH") + ) + if not plugin_so: + raise RuntimeError( + "Set EDGE_LLM_PLUGIN_SO (or EDGELLM_PLUGIN_PATH) to libNvInfer_edgellm_plugin.so" + ) + + ctypes.CDLL(plugin_so) + trt.init_libnvinfer_plugins(None, "") + return plugin_so + + +def load_plugins_for_trt(): + from .mamba import register_mamba_plugin_ops + from .moe import register_moe_plugin_ops + + _register_attention_plugin_op() + _register_vit_attention_plugin_op() + register_mamba_plugin_ops() + register_moe_plugin_ops() + load_plugin() + + from . import plugin_converter as _plugin_converter # noqa: F401,E402 + + +def restore_attention(patched): + for item in patched: + if len(item) == 2: + layer, original_attn = item + layer.self_attn = original_attn + else: + module, attr_name, original_attn = item + setattr(module, attr_name, original_attn) + + +def patch_vision_attention( + vision_model, + *, + batch_size: int, + seq_len: int, + name: str, + allow_attention_mask: bool = False, +): + patched = [] + + for layer in vision_model.encoder.layers: + patched.append((layer, layer.self_attn)) + layer.self_attn = ViTPluginAttention( + layer.self_attn, + batch_size=batch_size, + seq_len=seq_len, + name=name, + allow_attention_mask=allow_attention_mask, + ).eval() + + print(f"patched {name} attention modules: {len(patched)}") + return patched + + +def patch_molmo_vision_attention( + vision_backbone, + *, + batch_size: int, + seq_len: int, + name: str = "molmo-vision", +): + patched = [] + resblocks = vision_backbone.image_vit.transformer.resblocks + for i, block in enumerate(resblocks): + patched.append((block, "attention", block.attention)) + block.attention = MolmoViTPluginAttention( + block.attention, + batch_size=batch_size, + seq_len=seq_len, + name=f"{name}.image_vit.block{i}", + ).eval() + + print(f"patched {name} image_vit attention modules: {len(patched)}") + return patched + + +def patch_molmo_language_attention( + transformer: nn.Module, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + context_attention_mask_type: int = ContextAttentionMaskType.PADDING, + name: str = "molmo-language", +): + patched = [] + for i, block in enumerate(transformer.blocks): + patched.append((block, block.self_attn)) + block.self_attn = MolmoPluginAttention( + block.self_attn, + num_attention_heads=int(num_attention_heads), + num_key_value_heads=int(num_key_value_heads), + head_dim=int(head_dim), + hidden_size=int(hidden_size), + layer_idx=i, + context_attention_mask_type=context_attention_mask_type, + ).eval() + print(f"patched {name} attention modules: {len(patched)}") + return patched + + +def patch_vision_attention_reference(vision_model): + """ + Swap every SigLIP encoder-layer self_attn for SiglipReferenceAttention. + `vision_model` must be the INNER transformer (SiglipVisionTransformer), + i.e. eagle_model.vision_model.vision_model, matching patch_vision_attention. + Returns a list of (layer, original_attn) so it can be undone. + """ + patched = [] + for layer in vision_model.encoder.layers: + patched.append((layer, layer.self_attn)) + layer.self_attn = SiglipReferenceAttention(layer.self_attn).eval() + print(f"patched SigLIP reference attention modules: {len(patched)}") + return patched + + +def patch_language_attention( + language_model, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + context_attention_mask_type: int = ContextAttentionMaskType.PADDING, + name: str = "language", +): + patched = [] + + for i, layer in enumerate(language_model.layers): + patched.append((layer, layer.self_attn)) + layer.self_attn = PluginAttention( + layer.self_attn, + num_attention_heads=int(num_attention_heads), + num_key_value_heads=int(num_key_value_heads), + head_dim=int(head_dim), + hidden_size=int(hidden_size), + layer_idx=i, + context_attention_mask_type=context_attention_mask_type, + ).eval() + + print(f"patched {name} attention modules: {len(patched)}") + return patched + + +def patch_nemotron_mixers(model, config): + """Replace Nemotron hybrid mixers with plugin wrappers. MLP stays native.""" + from .mamba import PluginNemotronMamba + from .moe import PluginNemotronMoE + + wrappers = { + "NemotronHAttention": lambda mixer, idx: PluginNemotronAttention( + mixer, config, idx + ), + "NemotronHMamba2Mixer": lambda mixer, idx: PluginNemotronMamba(mixer), + "NemotronHMoE": lambda mixer, idx: PluginNemotronMoE(mixer, config), + } + layers = (getattr(model, "backbone", None) or model.model).layers + patched = [] + for i, block in enumerate(layers): + wrap = wrappers.get(type(block.mixer).__name__) + if wrap is None: + continue + original = block.mixer + block.mixer = wrap(original, i).eval() + patched.append((block, "mixer", original)) + print(f"patched nemotron mixers: {len(patched)}") + return patched + + +@torch.no_grad() +def infer_smolvlm_seq_len(vision_model, image): + patch_size = vision_model.patch_size + patch_attention_mask = torch.ones( + image.shape[0], + image.shape[2] // patch_size, + image.shape[3] // patch_size, + dtype=torch.bool, + device=image.device, + ) + hidden_states = vision_model.embeddings( + pixel_values=image, + patch_attention_mask=patch_attention_mask, + ) + return int(hidden_states.shape[0]), int(hidden_states.shape[1]) diff --git a/py/torch_tensorrt/hf/exporters/prefix_cache.py b/py/torch_tensorrt/hf/exporters/prefix_cache.py new file mode 100644 index 00000000000..3eeff80cc88 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/prefix_cache.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import torch + + +def _is_graph_capture_active() -> bool: + """True when running under Dynamo/export graph capture.""" + try: + import torch._dynamo as _dynamo + + if _dynamo.is_compiling(): + return True + except Exception: + pass + return False + + +def maybe_to( + t: torch.Tensor, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Avoid no-op .to() calls in hot paths.""" + if device is not None and t.device != device: + if dtype is not None and t.dtype != dtype: + return t.to(device=device, dtype=dtype) + return t.to(device=device) + if dtype is not None and t.dtype != dtype: + return t.to(dtype=dtype) + return t + + +class _CacheLayerView: + __slots__ = ("keys", "values") + + def __init__( + self, keys: torch.Tensor | None = None, values: torch.Tensor | None = None + ): + self.keys = keys + self.values = values + + +class _SmolVLAAttnPastEntry: + """Per-layer KV view for ``smolvlm_with_expert`` dict-style cache access.""" + + __slots__ = ("key_states", "value_states") + + def __init__(self, key_states: torch.Tensor, value_states: torch.Tensor): + self.key_states = key_states + self.value_states = value_states + + def __getitem__(self, key: str) -> torch.Tensor: + if key == "key_states": + return self.key_states + if key == "value_states": + return self.value_states + raise KeyError(key) + + +class SmolVLAPrefixPastLayers: + """Stacked prefix KV [L, B, H, S, D] -> SmolVLA ``past_key_values[layer]`` layout.""" + + def __init__(self, prefix_k: torch.Tensor, prefix_v: torch.Tensor): + self._k = prefix_k + self._v = prefix_v + + def __getitem__(self, layer_idx: int) -> _SmolVLAAttnPastEntry: + return _SmolVLAAttnPastEntry( + self._k[layer_idx].transpose(1, 2).contiguous(), + self._v[layer_idx].transpose(1, 2).contiguous(), + ) + + +class PrefixKVCache: + """KV cache backed by stacked tensors with shape [L, B, H, S, D].""" + + __prefix_kv_cache__ = True + + def __init__(self, prefix_k: torch.Tensor, prefix_v: torch.Tensor): + self._k = prefix_k + self._v = prefix_v + self._next_k: torch.Tensor | None = None + self._next_v: torch.Tensor | None = None + self._updated_k: list[torch.Tensor | None] = [None] * prefix_k.shape[0] + self._updated_v: list[torch.Tensor | None] = [None] * prefix_v.shape[0] + self.layers = [_CacheLayerView() for _ in range(prefix_k.shape[0])] + self._sync_layer_views() + + @classmethod + def empty( + cls, + *, + num_layers: int, + batch_size: int, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, + ) -> "PrefixKVCache": + empty_k = torch.zeros( + num_layers, + batch_size, + num_kv_heads, + 0, + head_dim, + dtype=dtype, + device=device, + ) + return cls(empty_k, torch.zeros_like(empty_k)) + + @property + def key_cache(self) -> torch.Tensor: + return self._k + + @property + def value_cache(self) -> torch.Tensor: + return self._v + + def _sync_layer_views(self) -> None: + if len(self.layers) != self._k.shape[0]: + self.layers = [_CacheLayerView() for _ in range(self._k.shape[0])] + for i, layer in enumerate(self.layers): + layer.keys = self._k[i] + layer.values = self._v[i] + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs=None, + ) -> tuple[torch.Tensor, torch.Tensor]: + del cache_kwargs + if _is_graph_capture_active(): + k = torch.cat([self._k[layer_idx], key_states], dim=-2) + v = torch.cat([self._v[layer_idx], value_states], dim=-2) + self._updated_k[layer_idx] = k + self._updated_v[layer_idx] = v + return k, v + + prefix_len = self._k.shape[3] + step_len = key_states.shape[-2] + total_len = prefix_len + step_len + + if self._next_k is None: + self._next_k = torch.empty( + self._k.shape[0], + self._k.shape[1], + self._k.shape[2], + total_len, + self._k.shape[4], + dtype=self._k.dtype, + device=self._k.device, + ) + self._next_v = torch.empty_like(self._next_k) + + k = self._next_k[layer_idx] + v = self._next_v[layer_idx] + k[..., :prefix_len, :] = self._k[layer_idx] + k[..., prefix_len:, :] = key_states + v[..., :prefix_len, :] = self._v[layer_idx] + v[..., prefix_len:, :] = value_states + self._updated_k[layer_idx] = k + self._updated_v[layer_idx] = v + return k, v + + def get_seq_length(self, layer_idx: int = 0) -> int: + del layer_idx + return self._k.shape[3] + + def get_max_cache_shape(self) -> int | None: + return None + + def get_mask_sizes( + self, cache_position: torch.Tensor, layer_idx: int = 0 + ) -> tuple[int, int]: + del layer_idx + kv_offset = 0 + query_length = cache_position.shape[0] + kv_length = self._k.shape[3] + query_length + return kv_length, kv_offset + + def get_updated_stacked(self) -> tuple[torch.Tensor, torch.Tensor]: + if self._next_k is not None and self._next_v is not None: + return self._next_k, self._next_v + if all(k is not None for k in self._updated_k) and all( + v is not None for v in self._updated_v + ): + return torch.stack(self._updated_k, dim=0), torch.stack( + self._updated_v, dim=0 + ) + return self._k, self._v + + def update_stacked( + self, key_cache: torch.Tensor, value_cache: torch.Tensor + ) -> None: + self._k = key_cache + self._v = value_cache + self._next_k = None + self._next_v = None + self._updated_k = [None] * key_cache.shape[0] + self._updated_v = [None] * value_cache.shape[0] + self._sync_layer_views() + + def reorder_cache(self, beam_idx: torch.LongTensor) -> None: + beam_idx = beam_idx.to(self._k.device) + self.update_stacked( + self._k.index_select(1, beam_idx), + self._v.index_select(1, beam_idx), + ) + + def batch_repeat_interleave(self, repeats: int) -> None: + if repeats == 1: + return + self.update_stacked( + self._k.repeat_interleave(repeats, dim=1), + self._v.repeat_interleave(repeats, dim=1), + ) + + def batch_select_indices(self, indices: torch.Tensor) -> None: + indices = indices.to(self._k.device) + self.update_stacked( + self._k[:, indices, ...], + self._v[:, indices, ...], + ) + + def crop(self, max_length: int) -> None: + if max_length < 0: + max_length = self.get_seq_length() - abs(max_length) + if self.get_seq_length() <= max_length: + return + self.update_stacked( + self._k[..., :max_length, :], + self._v[..., :max_length, :], + ) + + def __len__(self) -> int: + return int(self._k.shape[0]) + + def __iter__(self): + for i in range(len(self)): + yield self.layers[i].keys, self.layers[i].values + + def __getitem__(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]: + return self.layers[layer_idx].keys, self.layers[layer_idx].values + + +def stack_prefix_kv_from_cache( + past_key_values, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Convert a fully-initialized cache object into stacked [L, B, H, S, D] tensors. + """ + if isinstance(past_key_values, PrefixKVCache): + return ( + maybe_to(past_key_values.key_cache, device=device, dtype=dtype), + maybe_to(past_key_values.value_cache, device=device, dtype=dtype), + ) + + if isinstance(past_key_values, tuple) and len(past_key_values) == 2: + if past_key_values[0] is None or past_key_values[1] is None: + raise ValueError("Expected concrete stacked KV tensors; got (None, None)") + return ( + maybe_to(past_key_values[0], device=device, dtype=dtype), + maybe_to(past_key_values[1], device=device, dtype=dtype), + ) + + if hasattr(past_key_values, "layers"): + layers = list(past_key_values.layers) + if len(layers) == 0: + raise ValueError("Cache has no initialized layers") + k_tensors = [getattr(layer, "keys", None) for layer in layers] + v_tensors = [getattr(layer, "values", None) for layer in layers] + if any((k is None) != (v is None) for k, v in zip(k_tensors, v_tensors)): + raise ValueError("Inconsistent cache state: one of keys/values is None") + if any(k is None for k in k_tensors): + raise ValueError( + "Cache contains uninitialized layers; cannot infer stacked tensors" + ) + return ( + torch.stack( + [maybe_to(k, device=device, dtype=dtype) for k in k_tensors], dim=0 + ), + torch.stack( + [maybe_to(v, device=device, dtype=dtype) for v in v_tensors], dim=0 + ), + ) + + raise ValueError("Unsupported cache type for stack_prefix_kv_from_cache") diff --git a/py/torch_tensorrt/hf/exporters/rope.py b/py/torch_tensorrt/hf/exporters/rope.py new file mode 100644 index 00000000000..07282b7b51b --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/rope.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import logging + +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +def language_head_dim(config) -> int: + return int( + getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + ) + + +def config_to_dict(config) -> dict: + if hasattr(config, "to_dict"): + return dict(config.to_dict()) + raise TypeError( + f"Expected a HuggingFace-style config with to_dict(), got {type(config)!r}" + ) + + +def select_rope_parameters(config_dict: dict) -> dict: + rope_parameters = config_dict.get("rope_parameters") + if rope_parameters is None: + rope_parameters = config_dict.get("rope_scaling") + if not isinstance(rope_parameters, dict): + return {} + + layer_types = config_dict.get("layer_types") + if layer_types and set(rope_parameters.keys()).issubset(set(layer_types)): + full_attention = rope_parameters.get("full_attention") + if isinstance(full_attention, dict): + return dict(full_attention) + for layer_type in layer_types: + layer_params = rope_parameters.get(layer_type) + if isinstance(layer_params, dict): + return dict(layer_params) + return {} + + full_attention = rope_parameters.get("full_attention") + if isinstance(full_attention, dict): + return dict(full_attention) + + return dict(rope_parameters) + + +def normalize_rope_scaling(rope_params: dict) -> dict: + rope_scaling = dict(rope_params) + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type")) + if rope_type is not None: + rope_scaling.setdefault("rope_type", rope_type) + rope_scaling.setdefault("type", rope_type) + return rope_scaling + + +def export_rope_fields(config_dict: dict) -> dict: + """RoPE fields for LLMEngineRunner::collectRopeConfig / initializeRopeCosSinCache.""" + rope_fields: dict = {} + rope_params = select_rope_parameters(config_dict) + + if "rope_theta" in config_dict: + rope_fields["rope_theta"] = config_dict["rope_theta"] + elif "rope_theta" in rope_params: + rope_fields["rope_theta"] = rope_params["rope_theta"] + else: + raise KeyError("rope_theta not found in language model config") + + if rope_params: + rope_fields["rope_scaling"] = normalize_rope_scaling(rope_params) + else: + rope_fields["rope_scaling"] = None + + if rope_fields["rope_scaling"] is not None: + rope_type = rope_fields["rope_scaling"].get("rope_type") + if rope_type == "longrope": + original_max = rope_fields["rope_scaling"].get( + "original_max_position_embeddings", + config_dict.get("original_max_position_embeddings"), + ) + if original_max is None: + raise KeyError( + "original_max_position_embeddings required for longrope scaling" + ) + rope_fields["original_max_position_embeddings"] = original_max + + if "partial_rotary_factor" in config_dict: + rope_fields["partial_rotary_factor"] = config_dict["partial_rotary_factor"] + elif "partial_rotary_factor" in rope_params: + rope_fields["partial_rotary_factor"] = rope_params["partial_rotary_factor"] + else: + rope_fields["partial_rotary_factor"] = 1.0 + + return rope_fields + + +def rotary_dim_from_config(config) -> int: + head_dim = language_head_dim(config) + partial = float(getattr(config, "partial_rotary_factor", 1.0) or 1.0) + rotary_dim = int(head_dim * partial) + if rotary_dim <= 0 or rotary_dim > head_dim: + rotary_dim = head_dim + return rotary_dim + + +def _resolve_rotary_emb(language_model: nn.Module) -> nn.Module | None: + for module in (language_model, getattr(language_model, "model", None)): + if module is None: + continue + rotary_emb = getattr(module, "rotary_emb", None) + if rotary_emb is not None: + return rotary_emb + return None + + +def _plugin_rope_layout( + cos: torch.Tensor, + sin: torch.Tensor, + *, + max_seq_len: int, + rotary_dim: int, +) -> torch.Tensor: + """Pack cos/sin into AttentionPlugin layout: [:, :, :half]=cos, [:, :, half:]=sin.""" + half = rotary_dim // 2 + cos_half = cos[..., :half].float() + sin_half = sin[..., :half].float() + cache = torch.cat( + [cos_half[:, :max_seq_len], sin_half[:, :max_seq_len]], + dim=-1, + ) + if cache.shape[0] != 1: + cache = cache[:1] + return cache.contiguous() + + +def make_normal_rope_rotary_cos_sin( + max_seq_len: int, + rotary_dim: int, + *, + rope_theta: float, + rotary_scale: float = 1.0, + device: torch.device, +) -> torch.Tensor: + """Build RoPE cache using the same formula as Edge-LLM initializeNormalRopeCosSin.""" + half = rotary_dim // 2 + zid = torch.arange(half, device=device, dtype=torch.float32) + inv_denominator = float(rope_theta) ** (2 * zid / float(rotary_dim)) + positions = torch.arange(max_seq_len, device=device, dtype=torch.float32).unsqueeze( + 1 + ) + angles = positions * float(rotary_scale) / inv_denominator.unsqueeze(0) + cos = angles.cos() + sin = angles.sin() + return torch.cat([cos, sin], dim=-1).unsqueeze(0).to(dtype=torch.float32) + + +def make_rope_rotary_cos_sin_from_config( + config, + max_seq_len: int, + device: torch.device, +) -> torch.Tensor: + """Build RoPE cache from HF config fields (matches LLMEngineRunner default/dynamic RoPE).""" + config_dict = config_to_dict(config) + rope_params = select_rope_parameters(config_dict) + rope_type = None + if rope_params: + rope_type = rope_params.get("rope_type", rope_params.get("type")) + + if rope_type == "longrope": + raise NotImplementedError( + "longrope config cache is not implemented; pass language_model= to " + "make_rope_rotary_cos_sin() instead" + ) + if rope_type not in (None, "default", "dynamic", "llama3"): + raise NotImplementedError( + f"RoPE type {rope_type!r} requires language_model= and position_ids=" + ) + + rotary_dim = rotary_dim_from_config(config) + rope_theta = export_rope_fields(config_dict)["rope_theta"] + return make_normal_rope_rotary_cos_sin( + max_seq_len, + rotary_dim, + rope_theta=float(rope_theta), + rotary_scale=1.0, + device=device, + ) + + +@torch.no_grad() +def make_rope_rotary_cos_sin_from_model( + language_model: nn.Module, + config, + max_seq_len: int, + device: torch.device, + *, + position_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Build plugin RoPE cache from the model's rotary_emb (best eager parity).""" + rotary_emb = _resolve_rotary_emb(language_model) + if rotary_emb is None: + raise AttributeError("language_model has no rotary_emb") + + rotary_dim = rotary_dim_from_config(config) + if position_ids is None: + position_ids = torch.arange( + max_seq_len, device=device, dtype=torch.long + ).unsqueeze(0) + else: + position_ids = position_ids.to(device=device) + + seq_len = int(position_ids.shape[-1]) + dummy = torch.ones( + int(position_ids.shape[0] if position_ids.ndim >= 2 else 1), + seq_len, + 1, + device=device, + dtype=torch.float16, + ) + cos, sin = rotary_emb(dummy, position_ids) + return _plugin_rope_layout( + cos, + sin, + max_seq_len=max_seq_len, + rotary_dim=rotary_dim, + ) + + +def config_is_nope(config) -> bool: + """True when attention has no positional encoding (Nemotron-H / Nano).""" + if getattr(config, "use_rope", None) is False: + return True + model_type = str(getattr(config, "model_type", "") or "").lower() + return model_type.startswith("nemotron_h") + + +def make_nope_rotary_cos_sin( + max_seq_len: int, + rotary_dim: int, + device: torch.device, +) -> torch.Tensor: + """Identity cos/sin cache so AttentionPlugin is a RoPE pass-through. + + Matches Edge-LLM ``initializeNopeCosSinCache``: first half 1.0 (cos), + second half 0.0 (sin). Nemotron-H / Nano omit RoPE; position lives in SSM. + """ + half = int(rotary_dim) // 2 + cos = torch.ones(1, int(max_seq_len), half, dtype=torch.float32, device=device) + sin = torch.zeros(1, int(max_seq_len), half, dtype=torch.float32, device=device) + return torch.cat([cos, sin], dim=-1) + + +@torch.no_grad() +def make_rope_rotary_cos_sin( + config, + max_seq_len: int, + device: torch.device, + *, + language_model: nn.Module | None = None, + position_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Build real RoPE cache for inference/parity (not export tracing). + + Prefers the model's ``rotary_emb`` when available for eager parity. Falls back + to config-based generation that matches Edge-LLM ``initializeNormalRopeCosSin``. + Nemotron-H has no ``rotary_emb`` / ``rope_theta``; uses the NoPE identity cache. + """ + if config_is_nope(config): + return make_nope_rotary_cos_sin( + max_seq_len, rotary_dim_from_config(config), device + ) + if language_model is not None: + try: + return make_rope_rotary_cos_sin_from_model( + language_model, + config, + max_seq_len, + device, + position_ids=position_ids, + ) + except (AttributeError, NotImplementedError, RuntimeError) as exc: + logger.warning( + "Building RoPE from model rotary_emb failed (%s); using config kernel", + exc, + ) + return make_rope_rotary_cos_sin_from_config(config, max_seq_len, device) + + +def make_dummy_rope_rotary_cos_sin( + max_seq_len: int, + head_dim: int, + device: torch.device, +) -> torch.Tensor: + """Placeholder RoPE cache for export/compile tracing (runtime overwrites values).""" + return torch.randn( + 1, + int(max_seq_len), + int(head_dim), + dtype=torch.float32, + device=device, + ) diff --git a/py/torch_tensorrt/hf/exporters/specs/_language.py b/py/torch_tensorrt/hf/exporters/specs/_language.py index 25a8d86154b..a50c6bf0a6b 100644 --- a/py/torch_tensorrt/hf/exporters/specs/_language.py +++ b/py/torch_tensorrt/hf/exporters/specs/_language.py @@ -25,7 +25,7 @@ def causal_lm_flat( num_kv = int(cfg.num_key_value_heads) head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) try: - from trt.rope import make_rope_rotary_cos_sin + from torch_tensorrt.hf.exporters.rope import make_rope_rotary_cos_sin rope = make_rope_rotary_cos_sin( cfg, int(max_seq_len), device, language_model=language diff --git a/py/torch_tensorrt/hf/exporters/specs/groot.py b/py/torch_tensorrt/hf/exporters/specs/groot.py index a9c7d00d7ae..0843db4a7a1 100644 --- a/py/torch_tensorrt/hf/exporters/specs/groot.py +++ b/py/torch_tensorrt/hf/exporters/specs/groot.py @@ -30,8 +30,8 @@ def _groot(model: nn.Module) -> nn.Module: @register_edge_spec("groot", "gr00t") -class GrootSpec(EdgeSpec): - components = ("vision", "language", "action_context", "action") +class GrootSpec(EdgeSpec): # type: ignore[misc] + components = ("vision", "language", "context_projection", "action") def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any @@ -41,8 +41,12 @@ def prepare_sample_inputs( from lerobot.policies.factory import make_pre_post_processors from lerobot.policies.groot.processor_groot import GrootEagleEncodeStep - from trt.data import create_pil_messages, load_test_data, pack_state - from trt.executor.models.groot.helpers import make_embodiment_id + from torch_tensorrt.hf.exporters.data import ( + create_pil_messages, + load_test_data, + pack_state, + ) + from torch_tensorrt.hf.exporters.helpers.groot import make_embodiment_id policy = model device = raw.get( @@ -103,17 +107,17 @@ def prepare_sample_inputs( def wrap( self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any ) -> nn.Module: - from trt.modules.export.diffusion import ( - GrootDiTStepEncoderExportModule, - StaticActionVelocityStepExportModule, - TRTDynamicCategorySpecificMLPExportModule, + from torch_tensorrt.hf.exporters.patches.diffusion import ( + GrootDiTStepEncoderPatch, + StaticActionVelocityStepPatch, + TRTDynamicCategorySpecificMLPPatch, ) - from trt.modules.export.language import ( - CausalLMExportModule, - ContextProjectionExportModule, + from torch_tensorrt.hf.exporters.patches.language import ( + CausalLMPatch, + ContextProjectionPatch, language_decoder, ) - from trt.modules.export.vision import GridVisionExportModule + from torch_tensorrt.hf.exporters.patches.vision import GridVisionPatch found = _groot(model) eagle = found.backbone.eagle_model @@ -123,7 +127,7 @@ def wrap( # The e2e path casts the whole wrapper; without that, sa_embs is # float32 against bf16 attn.to_q weights. if name == "vision": - module = GridVisionExportModule( + module = GridVisionPatch( vision_model=eagle.vision_model, projector=eagle.mlp1, sample_pixel_values=sample["pixel_values"], @@ -134,22 +138,22 @@ def wrap( ) elif name == "language": language = eagle.language_model - module = CausalLMExportModule( + module = CausalLMPatch( language_decoder(language), language.lm_head, select_layer=-1 ) - elif name == "action_context": - module = ContextProjectionExportModule( + elif name == "context_projection": + module = ContextProjectionPatch( # type: ignore[no-untyped-call] found.backbone.eagle_linear, found.action_head.vlln, found.action_head.vl_self_attention, ) elif name == "action": - module = StaticActionVelocityStepExportModule( - step_encoder=GrootDiTStepEncoderExportModule( + module = StaticActionVelocityStepPatch( + step_encoder=GrootDiTStepEncoderPatch( found.action_head, sample.get("embodiment_id") ), action_expert=found.action_head.model, - velocity_decoder=TRTDynamicCategorySpecificMLPExportModule( + velocity_decoder=TRTDynamicCategorySpecificMLPPatch( found.action_head.action_decoder ), output_tokens=int(found.action_head.config.action_horizon), @@ -168,8 +172,10 @@ def prepare( config: Any, module: nn.Module, ) -> ComponentBundle: - from trt.plugin.attention import ContextAttentionMaskType - from trt.plugin.plugin_utils import ( + from torch_tensorrt.hf.exporters.plugin.attention import ( + ContextAttentionMaskType, + ) + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( patch_language_attention, patch_vision_attention, ) @@ -254,7 +260,7 @@ def _patch(mod: nn.Module) -> Any: engine_file="language.engine", ) - if name == "action_context": + if name == "context_projection": hidden = upstream["lm_hidden"].to(dtype=dtype) return ComponentBundle( trace_args=(hidden,), @@ -314,7 +320,7 @@ def capture_upstream( return {"visual_embeds": vis} if name == "language": return {"lm_hidden": outputs[1]} - if name == "action_context": + if name == "context_projection": ctx = outputs[0] if isinstance(outputs, tuple) else outputs return {"context_embs": ctx} return {} @@ -335,7 +341,7 @@ def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: sample["ds_stack"], *kv_kwargs(sample), ) - ctx = call_engine(engines["action_context"], "action_context", lm[1])[0] + ctx = call_engine(engines["context_projection"], "context_projection", lm[1])[0] return call_engine( engines["action"], "action", diff --git a/py/torch_tensorrt/hf/exporters/specs/nemotron.py b/py/torch_tensorrt/hf/exporters/specs/nemotron.py index 8fc8090b080..84887734d1f 100644 --- a/py/torch_tensorrt/hf/exporters/specs/nemotron.py +++ b/py/torch_tensorrt/hf/exporters/specs/nemotron.py @@ -29,7 +29,7 @@ def _kind(mixer: nn.Module) -> str: return "mlp" -class NemotronExportModule(nn.Module): # type: ignore[misc] +class NemotronPatch(nn.Module): # type: ignore[misc] """Hybrid decoder: plugin attention / mamba / moe, native MLP.""" def __init__(self, model: nn.Module): @@ -51,7 +51,9 @@ def forward( last_token_ids: torch.Tensor, *states: torch.Tensor, ) -> tuple[torch.Tensor, ...]: - from trt.modules.export.language import gather_last_token_hidden + from torch_tensorrt.hf.exporters.patches.language import ( + gather_last_token_hidden, + ) na, nm = self.num_attn, self.num_mamba kvs = list(states[:na]) @@ -139,7 +141,7 @@ def allocate_plugin_states( @register_edge_spec("nemotron_h", "nemotron") -class NemotronSpec(EdgeSpec): +class NemotronSpec(EdgeSpec): # type: ignore[misc] components = ("language",) def prepare_sample_inputs( @@ -177,14 +179,16 @@ def prepare_sample_inputs( def wrap( self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any ) -> nn.Module: - from trt.plugin.moe import PluginNemotronMoE - from trt.plugin.plugin_utils import patch_nemotron_mixers + from torch_tensorrt.hf.exporters.plugin.moe import PluginNemotronMoE + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + patch_nemotron_mixers, + ) - patch_nemotron_mixers(model, model.config) + patch_nemotron_mixers(model, model.config) # type: ignore[no-untyped-call] for block in _decoder(model).layers: if isinstance(block.mixer, PluginNemotronMoE): block.mixer.prepare_for_export() - return NemotronExportModule(model).eval() + return NemotronPatch(model).eval() def prepare( self, @@ -195,7 +199,7 @@ def prepare( config: Any, module: nn.Module, ) -> ComponentBundle: - from trt.rope import make_rope_rotary_cos_sin + from torch_tensorrt.hf.exporters.rope import make_rope_rotary_cos_sin embeds = sample["inputs_embeds"] device, dtype = embeds.device, embeds.dtype diff --git a/py/torch_tensorrt/hf/exporters/specs/pi05.py b/py/torch_tensorrt/hf/exporters/specs/pi05.py index 33f013c8c37..5d1baa3b8ad 100644 --- a/py/torch_tensorrt/hf/exporters/specs/pi05.py +++ b/py/torch_tensorrt/hf/exporters/specs/pi05.py @@ -36,7 +36,7 @@ def _nchw_to_hwc(pixel_values: torch.Tensor) -> torch.Tensor: @register_edge_spec("pi05") -class Pi05Spec(EdgeSpec): +class Pi05Spec(EdgeSpec): # type: ignore[misc] components = ("vision", "language", "action") def prepare_sample_inputs( @@ -50,7 +50,10 @@ def prepare_sample_inputs( OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, ) - from trt.data import frame_from_test_data, load_test_data + from torch_tensorrt.hf.exporters.data import ( + frame_from_test_data, + load_test_data, + ) policy = model if hasattr(model, "_preprocess_images") else None if policy is None: @@ -90,19 +93,22 @@ def prepare_sample_inputs( def wrap( self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any ) -> nn.Module: - from trt.modules.export.diffusion import ( - PI05PrefixKVStepEncoderExportModule, - StaticActionVelocityStepExportModule, + from torch_tensorrt.hf.exporters.patches.diffusion import ( + PI05PrefixKVStepEncoderPatch, + StaticActionVelocityStepPatch, + ) + from torch_tensorrt.hf.exporters.patches.language import ( + CausalLMPatch, + language_decoder, ) - from trt.modules.export.language import CausalLMExportModule, language_decoder - from trt.modules.export.vision import GridVisionExportModule + from torch_tensorrt.hf.exporters.patches.vision import GridVisionPatch core = _core(model) paligemma = core.paligemma_with_expert.paligemma.model if name == "vision": px = sample["pixel_values"] sample_px = px.float() if px.dtype != torch.float32 else px - return GridVisionExportModule( + return GridVisionPatch( vision_model=paligemma.vision_tower.float(), projector=paligemma.multi_modal_projector, sample_pixel_values=sample_px, @@ -114,10 +120,10 @@ def wrap( if name == "language": language = paligemma.language_model lm_head = core.paligemma_with_expert.paligemma.lm_head - return CausalLMExportModule(language_decoder(language), lm_head).eval() + return CausalLMPatch(language_decoder(language), lm_head).eval() if name == "action": - return StaticActionVelocityStepExportModule( - step_encoder=PI05PrefixKVStepEncoderExportModule(core), + return StaticActionVelocityStepPatch( + step_encoder=PI05PrefixKVStepEncoderPatch(core), # type: ignore[no-untyped-call] action_expert=core.paligemma_with_expert.gemma_expert.model, velocity_decoder=core.action_out_proj, output_tokens=int(core.config.chunk_size), @@ -134,13 +140,15 @@ def prepare( config: Any, module: nn.Module, ) -> ComponentBundle: - from trt.executor.models.pi05.helpers import ( + from torch_tensorrt.hf.exporters.helpers.pi05 import ( build_pi05_prefix_embs, make_pi05_suffix_position_and_mask, pi05_compact_index, ) - from trt.plugin.attention import ContextAttentionMaskType - from trt.plugin.plugin_utils import ( + from torch_tensorrt.hf.exporters.plugin.attention import ( + ContextAttentionMaskType, + ) + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( patch_language_attention, patch_vision_attention, ) @@ -253,7 +261,7 @@ def _patch(mod: nn.Module) -> Any: 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( + pos, mask = make_pi05_suffix_position_and_mask( # type: ignore[no-untyped-call] core_mod, sample["prefix_pad_mask"], step_actions, device ) sample["suffix_position_ids"] = pos diff --git a/py/torch_tensorrt/hf/exporters/utils.py b/py/torch_tensorrt/hf/exporters/utils.py new file mode 100644 index 00000000000..fb6c04b1779 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/utils.py @@ -0,0 +1,48 @@ +"""Host/device helpers used by the example scripts and attention patches.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch + +_THOR_CUDA_LIB = Path("/usr/local/cuda-13.0/thor/targets/aarch64-linux/lib") + + +def configure_thor_pytorch() -> None: + """Use PyTorch fallbacks for ops whose pip CUDA wheels mismatch DriveOS Thor.""" + on_thor = os.environ.get("TRT_VLA_THOR", "auto") + if on_thor == "auto": + on_thor = "1" if _THOR_CUDA_LIB.is_dir() else "0" + if on_thor == "1": + torch.backends.cudnn.enabled = False + + +def force_hf_attention(module: Any, attn: str, use_cache: bool | None = False) -> None: + """Force HuggingFace attention implementation on a module tree.""" + for m in module.modules(): + cfg = getattr(m, "config", None) + if cfg is None: + continue + if hasattr(cfg, "_attn_implementation"): + cfg._attn_implementation = attn + if hasattr(cfg, "attn_implementation"): + cfg.attn_implementation = attn + if use_cache is not None and hasattr(cfg, "use_cache"): + cfg.use_cache = use_cache + + cfg = getattr(module, "config", None) + if cfg is None: + return + for name in ("vision_config", "text_config"): + sub_cfg = getattr(cfg, name, None) + if sub_cfg is None: + continue + if hasattr(sub_cfg, "_attn_implementation"): + sub_cfg._attn_implementation = attn + if hasattr(sub_cfg, "attn_implementation"): + sub_cfg.attn_implementation = attn + if use_cache is not None and hasattr(sub_cfg, "use_cache"): + sub_cfg.use_cache = use_cache diff --git a/pyproject.toml b/pyproject.toml index f48d7abdc4f..8fe1ce6b7db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -451,6 +451,18 @@ module = "torch_tensorrt.fx.*" ignore_errors = true follow_imports = "skip" +[[tool.mypy.overrides]] +module = [ + "torch_tensorrt.hf.exporters.plugin.*", + "torch_tensorrt.hf.exporters.patches.*", + "torch_tensorrt.hf.exporters.helpers.*", + "torch_tensorrt.hf.exporters.data", + "torch_tensorrt.hf.exporters.rope", + "torch_tensorrt.hf.exporters.prefix_cache", + "torch_tensorrt.hf.exporters.mamba_stub", +] +ignore_errors = true + [tool.typos] files.extend-exclude = [ "docs/**/*", diff --git a/setup.py b/setup.py index 10269cfe912..8b08069e475 100644 --- a/setup.py +++ b/setup.py @@ -618,6 +618,9 @@ def run(self): "torch_tensorrt.executorch", "torch_tensorrt.hf", "torch_tensorrt.hf.exporters", + "torch_tensorrt.hf.exporters.helpers", + "torch_tensorrt.hf.exporters.patches", + "torch_tensorrt.hf.exporters.plugin", "torch_tensorrt.hf.exporters.specs", "torch_tensorrt.runtime", ] @@ -660,6 +663,9 @@ def run(self): "torch_tensorrt.executorch": "py/torch_tensorrt/executorch", "torch_tensorrt.hf": "py/torch_tensorrt/hf", "torch_tensorrt.hf.exporters": "py/torch_tensorrt/hf/exporters", + "torch_tensorrt.hf.exporters.helpers": "py/torch_tensorrt/hf/exporters/helpers", + "torch_tensorrt.hf.exporters.patches": "py/torch_tensorrt/hf/exporters/patches", + "torch_tensorrt.hf.exporters.plugin": "py/torch_tensorrt/hf/exporters/plugin", "torch_tensorrt.hf.exporters.specs": "py/torch_tensorrt/hf/exporters/specs", "torch_tensorrt.runtime": "py/torch_tensorrt/runtime", } From b3e0f1fb4f99d081adcba006e51dc4264e2bde16 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 13:19:08 -0700 Subject: [PATCH 03/11] Group PI05, GR00T, and Nemotron under exporters/models. Each family now has spec, helpers, and patches in one folder, with shared vision/language/action-step code in models/common. --- py/torch_tensorrt/hf/exporters/__init__.py | 12 +- .../hf/exporters/models/__init__.py | 5 + .../{helpers => models/common}/__init__.py | 0 .../_language.py => models/common/helpers.py} | 0 .../hf/exporters/models/common/patches.py | 366 ++++++++++++++++++ .../{patches => models/groot}/__init__.py | 0 .../groot.py => models/groot/helpers.py} | 12 + .../diffusion.py => models/groot/patches.py} | 148 +------ .../{specs/groot.py => models/groot/spec.py} | 49 +-- .../hf/exporters/models/nemotron/__init__.py | 1 + .../hf/exporters/models/nemotron/helpers.py | 69 ++++ .../hf/exporters/models/nemotron/patches.py | 65 ++++ .../nemotron.py => models/nemotron/spec.py} | 136 +------ .../hf/exporters/models/pi05/__init__.py | 1 + .../pi05.py => models/pi05/helpers.py} | 18 + .../hf/exporters/models/pi05/patches.py | 52 +++ .../{specs/pi05.py => models/pi05/spec.py} | 56 +-- .../hf/exporters/patches/language.py | 148 ------- .../hf/exporters/patches/vision.py | 158 -------- .../hf/exporters/specs/__init__.py | 5 - pyproject.toml | 3 +- setup.py | 16 +- 22 files changed, 676 insertions(+), 644 deletions(-) create mode 100644 py/torch_tensorrt/hf/exporters/models/__init__.py rename py/torch_tensorrt/hf/exporters/{helpers => models/common}/__init__.py (100%) rename py/torch_tensorrt/hf/exporters/{specs/_language.py => models/common/helpers.py} (100%) create mode 100644 py/torch_tensorrt/hf/exporters/models/common/patches.py rename py/torch_tensorrt/hf/exporters/{patches => models/groot}/__init__.py (100%) rename py/torch_tensorrt/hf/exporters/{helpers/groot.py => models/groot/helpers.py} (59%) rename py/torch_tensorrt/hf/exporters/{patches/diffusion.py => models/groot/patches.py} (63%) rename py/torch_tensorrt/hf/exporters/{specs/groot.py => models/groot/spec.py} (92%) create mode 100644 py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py create mode 100644 py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py create mode 100644 py/torch_tensorrt/hf/exporters/models/nemotron/patches.py rename py/torch_tensorrt/hf/exporters/{specs/nemotron.py => models/nemotron/spec.py} (51%) create mode 100644 py/torch_tensorrt/hf/exporters/models/pi05/__init__.py rename py/torch_tensorrt/hf/exporters/{helpers/pi05.py => models/pi05/helpers.py} (92%) create mode 100644 py/torch_tensorrt/hf/exporters/models/pi05/patches.py rename py/torch_tensorrt/hf/exporters/{specs/pi05.py => models/pi05/spec.py} (89%) delete mode 100644 py/torch_tensorrt/hf/exporters/patches/language.py delete mode 100644 py/torch_tensorrt/hf/exporters/patches/vision.py delete mode 100644 py/torch_tensorrt/hf/exporters/specs/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/__init__.py b/py/torch_tensorrt/hf/exporters/__init__.py index 5d763b03cb8..01dd627defa 100644 --- a/py/torch_tensorrt/hf/exporters/__init__.py +++ b/py/torch_tensorrt/hf/exporters/__init__.py @@ -1,14 +1,20 @@ from torch_tensorrt.hf.exporters.config import EdgeConfig from torch_tensorrt.hf.exporters.exporter import EdgeExporter +from torch_tensorrt.hf.exporters.models.groot.spec import ( # noqa: F401 + GrootSpec as _GrootSpec, +) +from torch_tensorrt.hf.exporters.models.nemotron.spec import ( # noqa: F401 + NemotronSpec as _NemotronSpec, +) +from torch_tensorrt.hf.exporters.models.pi05.spec import ( # noqa: F401 + Pi05Spec as _Pi05Spec, +) from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, EdgeSpec, get_edge_spec, register_edge_spec, ) -from torch_tensorrt.hf.exporters.specs import groot as _groot # noqa: F401 -from torch_tensorrt.hf.exporters.specs import nemotron as _nemotron # noqa: F401 -from torch_tensorrt.hf.exporters.specs import pi05 as _pi05 # noqa: F401 __all__ = [ "ComponentBundle", diff --git a/py/torch_tensorrt/hf/exporters/models/__init__.py b/py/torch_tensorrt/hf/exporters/models/__init__.py new file mode 100644 index 00000000000..bfb17db20d4 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/__init__.py @@ -0,0 +1,5 @@ +"""Model families. Importing this package registers EdgeSpecs.""" + +from torch_tensorrt.hf.exporters.models.groot import spec as _groot # noqa: F401 +from torch_tensorrt.hf.exporters.models.nemotron import spec as _nemotron # noqa: F401 +from torch_tensorrt.hf.exporters.models.pi05 import spec as _pi05 # noqa: F401 diff --git a/py/torch_tensorrt/hf/exporters/helpers/__init__.py b/py/torch_tensorrt/hf/exporters/models/common/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/helpers/__init__.py rename to py/torch_tensorrt/hf/exporters/models/common/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/specs/_language.py b/py/torch_tensorrt/hf/exporters/models/common/helpers.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/specs/_language.py rename to py/torch_tensorrt/hf/exporters/models/common/helpers.py diff --git a/py/torch_tensorrt/hf/exporters/models/common/patches.py b/py/torch_tensorrt/hf/exporters/models/common/patches.py new file mode 100644 index 00000000000..3d46d8b9985 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/common/patches.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def hwc_to_nchw(images: torch.Tensor) -> torch.Tensor: + if images.ndim != 4: + raise ValueError(f"Expected 4D images, got shape {tuple(images.shape)}") + return images.permute(0, 3, 1, 2).contiguous() + + +def is_nchw_pixel_values(pixel_values: torch.Tensor) -> bool: + return ( + pixel_values.ndim == 4 + and pixel_values.shape[1] in (1, 3, 4) + and pixel_values.shape[-1] not in (1, 3, 4) + ) + + +# --------------------------------------------------------------------------- +# Base +# --------------------------------------------------------------------------- + + +class VisionPatch(nn.Module): + """Base TRT trace target: probe static shapes once, flatten output to [N, H].""" + + cast_output_to_input_dtype: bool + output_num_tokens: int + output_hidden_size: int + + def _finalize_output( + self, features: torch.Tensor, out_dtype: torch.dtype + ) -> torch.Tensor: + if self.cast_output_to_input_dtype and features.dtype != out_dtype: + features = features.to(out_dtype) + if features.ndim == 3: + # [B, S, H] -> [B*S, H] + return features.reshape(-1, features.shape[-1]) + if features.ndim == 2: + # already [N, H] (token-pooling encoders) + return features + raise ValueError( + f"Expected 2D or 3D features, got shape {tuple(features.shape)}" + ) + + +# --------------------------------------------------------------------------- +# Grid vision (PI0.5 / GR00T / SmolVLA VitRunner path) +# --------------------------------------------------------------------------- + + +class GridVisionPatch(VisionPatch): + """pixels [B,H,W,C] -> fixed patch grid -> [B*seq_len, lm_hidden]""" + + def __init__( + self, + *, + vision_model: nn.Module, + projector: nn.Module, + sample_pixel_values: torch.Tensor, + select_layer: int = -1, + pixel_shuffle: bool = False, + downsample_ratio: float = 0.5, + force_float32_input: bool = False, + cast_output_to_input_dtype: bool = False, + vision_kwargs: dict[str, Any] | None = None, + ): + super().__init__() + self.vision_model = vision_model + self.projector = projector + self.select_layer = int(select_layer) + self.pixel_shuffle = bool(pixel_shuffle) + self.downsample_ratio = float(downsample_ratio) + self.force_float32_input = bool(force_float32_input) + self.cast_output_to_input_dtype = bool(cast_output_to_input_dtype) + self.vision_kwargs = dict(vision_kwargs or {}) + + with torch.no_grad(): + sample = sample_pixel_values + if self.force_float32_input and sample.dtype != torch.float32: + sample = sample.to(torch.float32) + + vit_embeds = self._select_vision_features(self._run_vision(sample)) + self.seq_len = int(vit_embeds.shape[1]) + self.hidden_size = int(vit_embeds.shape[2]) + + if self.pixel_shuffle: + self._init_pixel_shuffle_shape() + vit_embeds = self._apply_pixel_shuffle(vit_embeds) + + projected = self.projector(self._projector_input(vit_embeds)) + self.batch_size = int(projected.shape[0]) + self.output_seq_len = int(projected.shape[1]) + self.output_hidden_size = int(projected.shape[2]) + self.output_num_tokens = self.batch_size * self.output_seq_len + + def _run_vision(self, images: torch.Tensor): + pixel_values = ( + hwc_to_nchw(images) if not is_nchw_pixel_values(images) else images + ) + kwargs = dict(self.vision_kwargs) + kwargs["pixel_values"] = pixel_values + kwargs["output_hidden_states"] = self.select_layer != -1 + kwargs.setdefault("return_dict", True) + return self.vision_model(**kwargs) + + def _select_vision_features(self, out): + if self.select_layer == -1: + return ( + out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] + ) + return out.hidden_states[self.select_layer] + + def _projector_input(self, vit_embeds: torch.Tensor) -> torch.Tensor: + proj_dtype = next(self.projector.parameters()).dtype + if vit_embeds.dtype != proj_dtype: + return vit_embeds.to(proj_dtype) + return vit_embeds + + def _init_pixel_shuffle_shape(self): + side = int(self.seq_len**0.5) + if side * side != self.seq_len: + raise ValueError( + f"Expected square vision sequence, got seq_len={self.seq_len}" + ) + self.grid_w = side + self.grid_h = side + self.out_w = int(self.grid_w * self.downsample_ratio) + self.out_h = int(self.grid_h * self.downsample_ratio) + self.hidden_after_first_view = int(self.hidden_size / self.downsample_ratio) + self.shuffle_hidden = int( + self.hidden_size / (self.downsample_ratio * self.downsample_ratio) + ) + + def _apply_pixel_shuffle(self, x): + n = x.shape[0] + x = x.reshape(n, self.grid_w, self.out_h, self.hidden_after_first_view) + x = x.permute(0, 2, 1, 3).contiguous() + x = x.reshape(n, self.out_h, self.out_w, self.shuffle_hidden) + x = x.permute(0, 2, 1, 3).contiguous() + return x.reshape(n, -1, self.shuffle_hidden) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + out_dtype = pixel_values.dtype + images = pixel_values + if self.force_float32_input and images.dtype != torch.float32: + images = images.to(torch.float32) + + vit_embeds = self._select_vision_features(self._run_vision(images)) + if self.pixel_shuffle: + vit_embeds = self._apply_pixel_shuffle(vit_embeds) + + features = self.projector(self._projector_input(vit_embeds)) + return self._finalize_output(features, out_dtype) + + +def _as_tensor(x): + """Unwrap tuple/list outputs from patched attention modules.""" + if isinstance(x, (tuple, list)): + return x[0] + return x + + +def language_decoder(language: nn.Module) -> nn.Module: + """Inner module that owns ``.layers``. + + Paligemma / PI05 store layers on ``language_model`` itself. HF + ``*ForCausalLM`` stores them on ``language_model.model``. Prefer ``.layers`` + so a stray ``.model`` attribute cannot silently pick the wrong submodule. + """ + if hasattr(language, "layers"): + return language + inner = getattr(language, "model", None) + if isinstance(inner, nn.Module) and hasattr(inner, "layers"): + return inner + 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] + + +class CausalLMPatch(nn.Module): + """Edge-LLM causal LM: manual decoder loop -> logits + lm_hidden_states + prefix KV. + + External RoPE and KV-cache controls match ``LLMEngineRunner`` prefill/decode. + ``select_layer=-1`` (default for GR00T TRT) uses final RMSNorm hidden for + context; positive values capture an intermediate layer output pre-norm. + + ``ds_stack`` is always an input (``[num_ds, B, S, H]``). Models without + deepstack pass ``num_ds=0`` (empty leading dim); the add is a no-op. + """ + + def __init__( + self, + lm: nn.Module, + lm_head: nn.Module, + *, + select_layer: int = -1, + ): + super().__init__() + self.lm = lm + self.lm_head = lm_head + self.select_layer = int(select_layer) + + def forward( + self, + inputs_embeds: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + context_lengths: torch.Tensor, + kvcache_start_index: torch.Tensor, + last_token_ids: torch.Tensor, + ds_stack: torch.Tensor, + *past_key_values: torch.Tensor, + ): + lm_dtype = next(self.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 self.select_layer == 0 else None + new_kvs = [] + + for i, layer in enumerate(self.lm.layers): + residual = hidden + hidden = _as_tensor(layer.input_layernorm(hidden)) + hidden, kv = layer.self_attn( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=past_key_values[i], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, + ) + hidden = _as_tensor(hidden) + hidden = residual + hidden + + residual = hidden + hidden = _as_tensor(layer.post_attention_layernorm(hidden)) + hidden = _as_tensor(layer.mlp(hidden)) + hidden = residual + hidden + new_kvs.append(kv) + + if i < num_ds: + hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) + + if self.select_layer > 0 and (i + 1) == self.select_layer: + context_hidden = hidden + + hidden = _as_tensor(self.lm.norm(hidden)) + if context_hidden is None: + context_hidden = hidden + + last_hidden = gather_last_token_hidden(hidden, last_token_ids) + logits = self.lm_head(last_hidden).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 + + +class ActionStepEncoderPatch(nn.Module): + """Base contract for model-specific action-step encoding. + + Subclasses implement forward() to turn noisy actions, timestep, and + model-specific context tensors into the args/kwargs consumed by the action + expert and velocity decoder. The default helpers cover common expert output + and velocity shapes, while model-specific encoders can override them. + """ + + def get_action_hidden(self, expert_out, output_tokens: int): + # Default path for experts that return either a standard model output + # with last_hidden_state, a raw hidden-state tensor, or a tuple/list whose + # first item is the hidden-state tensor. + hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out + ) + + if isinstance(hidden, (tuple, list)): + hidden = hidden[0] + + return hidden[:, -output_tokens:] + + def process_velocity(self, velocity): + # Default path for models whose decoder already returns the final action + # velocity shape. Override for models that need reshaping or cropping. + return velocity + + +class StaticActionVelocityStepPatch(nn.Module): + """One static denoising step shared by VLA action diffusion modules. + + The model-specific step_encoder owns the messy part: converting noisy + actions, timestep, and context tensors into the exact action_expert call. + This wrapper only runs the expert, selects action-token hidden states, and + decodes those hidden states into a velocity update. + """ + + def __init__( + self, + *, + step_encoder: ActionStepEncoderPatch, + action_expert: nn.Module, + velocity_decoder: nn.Module, + output_tokens: int, + cast_hidden_fp32: bool = True, + ): + super().__init__() + self.step_encoder = step_encoder + self.action_expert = action_expert + self.velocity_decoder = velocity_decoder + self.output_tokens = int(output_tokens) + self.cast_hidden_fp32 = cast_hidden_fp32 + + def forward(self, x_t, timestep, *inputs): + # Build the action expert inputs and any decoder-specific side inputs. + expert_args, expert_kwargs, decoder_args, decoder_kwargs = self.step_encoder( + x_t, + timestep, + *inputs, + ) + + # Run the model-specific action expert: Gemma expert, DiT, etc. + expert_out = self.action_expert(*expert_args, **expert_kwargs) + + # Most experts return last_hidden_state, but some wrappers return tuples + # or need custom suffix-token selection. + action_hidden = self.step_encoder.get_action_hidden( + expert_out, + self.output_tokens, + ) + + if self.cast_hidden_fp32: + action_hidden = action_hidden.to(dtype=torch.float32) + + # Project action-token hidden states back to action-space velocity. + velocity = self.velocity_decoder( + action_hidden, + *decoder_args, + **decoder_kwargs, + ) + + return self.step_encoder.process_velocity(velocity) diff --git a/py/torch_tensorrt/hf/exporters/patches/__init__.py b/py/torch_tensorrt/hf/exporters/models/groot/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/patches/__init__.py rename to py/torch_tensorrt/hf/exporters/models/groot/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/helpers/groot.py b/py/torch_tensorrt/hf/exporters/models/groot/helpers.py similarity index 59% rename from py/torch_tensorrt/hf/exporters/helpers/groot.py rename to py/torch_tensorrt/hf/exporters/models/groot/helpers.py index 47d2e3a9cba..6db2f9e56b5 100644 --- a/py/torch_tensorrt/hf/exporters/helpers/groot.py +++ b/py/torch_tensorrt/hf/exporters/models/groot/helpers.py @@ -3,6 +3,7 @@ from typing import Any import torch +import torch.nn as nn GROOT_EMBODIMENT_MAPPING = { "new_embodiment": 31, @@ -25,3 +26,14 @@ def make_embodiment_id( dtype=dtype, device=device, ) + + +def _groot(model: nn.Module) -> nn.Module: + if hasattr(model, "_groot_model"): + return model._groot_model + backbone = getattr(model, "backbone", None) + if backbone is not None and hasattr(backbone, "eagle_model"): + return model + raise RuntimeError( + "GR00T spec expected GrootPolicy or a module with backbone.eagle_model" + ) diff --git a/py/torch_tensorrt/hf/exporters/patches/diffusion.py b/py/torch_tensorrt/hf/exporters/models/groot/patches.py similarity index 63% rename from py/torch_tensorrt/hf/exporters/patches/diffusion.py rename to py/torch_tensorrt/hf/exporters/models/groot/patches.py index 243cdffd1e4..e882f1e25ed 100644 --- a/py/torch_tensorrt/hf/exporters/patches/diffusion.py +++ b/py/torch_tensorrt/hf/exporters/models/groot/patches.py @@ -1,8 +1,9 @@ +from __future__ import annotations + import torch import torch.nn as nn import torch.nn.functional as F -from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding -from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache +from torch_tensorrt.hf.exporters.models.common.patches import ActionStepEncoderPatch class TRTFixedCategorySpecificLinearPatch(nn.Module): @@ -149,92 +150,6 @@ def forward(self, actions, timesteps, embodiment_id): return self.W3(hidden, embodiment_id) -class ActionStepEncoderPatch(nn.Module): - """Base contract for model-specific action-step encoding. - - Subclasses implement forward() to turn noisy actions, timestep, and - model-specific context tensors into the args/kwargs consumed by the action - expert and velocity decoder. The default helpers cover common expert output - and velocity shapes, while model-specific encoders can override them. - """ - - def get_action_hidden(self, expert_out, output_tokens: int): - # Default path for experts that return either a standard model output - # with last_hidden_state, a raw hidden-state tensor, or a tuple/list whose - # first item is the hidden-state tensor. - hidden = ( - expert_out.last_hidden_state - if hasattr(expert_out, "last_hidden_state") - else expert_out - ) - - if isinstance(hidden, (tuple, list)): - hidden = hidden[0] - - return hidden[:, -output_tokens:] - - def process_velocity(self, velocity): - # Default path for models whose decoder already returns the final action - # velocity shape. Override for models that need reshaping or cropping. - return velocity - - -class StaticActionVelocityStepPatch(nn.Module): - """One static denoising step shared by VLA action diffusion modules. - - The model-specific step_encoder owns the messy part: converting noisy - actions, timestep, and context tensors into the exact action_expert call. - This wrapper only runs the expert, selects action-token hidden states, and - decodes those hidden states into a velocity update. - """ - - def __init__( - self, - *, - step_encoder: ActionStepEncoderPatch, - action_expert: nn.Module, - velocity_decoder: nn.Module, - output_tokens: int, - cast_hidden_fp32: bool = True, - ): - super().__init__() - self.step_encoder = step_encoder - self.action_expert = action_expert - self.velocity_decoder = velocity_decoder - self.output_tokens = int(output_tokens) - self.cast_hidden_fp32 = cast_hidden_fp32 - - def forward(self, x_t, timestep, *inputs): - # Build the action expert inputs and any decoder-specific side inputs. - expert_args, expert_kwargs, decoder_args, decoder_kwargs = self.step_encoder( - x_t, - timestep, - *inputs, - ) - - # Run the model-specific action expert: Gemma expert, DiT, etc. - expert_out = self.action_expert(*expert_args, **expert_kwargs) - - # Most experts return last_hidden_state, but some wrappers return tuples - # or need custom suffix-token selection. - action_hidden = self.step_encoder.get_action_hidden( - expert_out, - self.output_tokens, - ) - - if self.cast_hidden_fp32: - action_hidden = action_hidden.to(dtype=torch.float32) - - # Project action-token hidden states back to action-space velocity. - velocity = self.velocity_decoder( - action_hidden, - *decoder_args, - **decoder_kwargs, - ) - - return self.step_encoder.process_velocity(velocity) - - class GrootDiTStepEncoderPatch(ActionStepEncoderPatch): def __init__(self, action_head, embodiment_id: torch.Tensor | None = None): super().__init__() @@ -292,47 +207,22 @@ def forward(self, actions, timestep, vl_embs, state, embodiment_id): return expert_args, expert_kwargs, decoder_args, decoder_kwargs -class PI05PrefixKVStepEncoderPatch(ActionStepEncoderPatch): - """PI05 suffix embed + AdaRMS cond, consumed by Gemma action expert.""" +class ContextProjectionPatch(nn.Module): + """eagle_linear -> vlln -> vl_self_attention (matches eager context path).""" - def __init__(self, core): + def __init__(self, eagle_linear, vlln, vl_self_attention): super().__init__() - self.action_in_proj = core.action_in_proj - self.time_mlp_in = core.time_mlp_in - self.time_mlp_out = core.time_mlp_out - self.config = core.config - self.hidden_size = core.action_in_proj.out_features - - def forward( - self, - x_t, - timestep, - prefix_k, - prefix_v, - position_ids, - attention_mask, - ): - suffix_embs = self.action_in_proj(x_t) - - time_emb = create_sinusoidal_pos_embedding( - timestep, - self.hidden_size, - min_period=self.config.min_period, - max_period=self.config.max_period, - device=timestep.device, - ).to(dtype=suffix_embs.dtype) - - adarms_cond = self.time_mlp_in(time_emb) - adarms_cond = F.silu(adarms_cond) - adarms_cond = self.time_mlp_out(adarms_cond) - adarms_cond = F.silu(adarms_cond) + self.eagle_linear = eagle_linear + self.vlln = vlln + self.vl_self_attention = vl_self_attention - expert_kwargs = { - "inputs_embeds": suffix_embs, - "attention_mask": attention_mask, - "position_ids": position_ids, - "past_key_values": PrefixKVCache(prefix_k, prefix_v), - "use_cache": False, - "adarms_cond": adarms_cond, - } - return (), expert_kwargs, (), {} + def forward(self, hidden_states: torch.Tensor): + context_embs = self.eagle_linear(hidden_states) + + vlln_weight = getattr(self.vlln, "weight", None) + if vlln_weight is not None: + context_embs = context_embs.to(dtype=vlln_weight.dtype) + + context_embs = self.vlln(context_embs) + context_embs = self.vl_self_attention(context_embs) + return context_embs diff --git a/py/torch_tensorrt/hf/exporters/specs/groot.py b/py/torch_tensorrt/hf/exporters/models/groot/spec.py similarity index 92% rename from py/torch_tensorrt/hf/exporters/specs/groot.py rename to py/torch_tensorrt/hf/exporters/models/groot/spec.py index 0843db4a7a1..b39529c23fe 100644 --- a/py/torch_tensorrt/hf/exporters/specs/groot.py +++ b/py/torch_tensorrt/hf/exporters/models/groot/spec.py @@ -5,28 +5,32 @@ import torch import torch.nn as nn +from torch_tensorrt.hf.exporters.models.common.helpers import ( + causal_lm_flat, + kv_kwargs, + split_flat_to_kwargs, +) +from torch_tensorrt.hf.exporters.models.common.patches import ( + CausalLMPatch, + GridVisionPatch, + StaticActionVelocityStepPatch, + language_decoder, +) +from torch_tensorrt.hf.exporters.models.groot.helpers import ( + _groot, + make_embodiment_id, +) +from torch_tensorrt.hf.exporters.models.groot.patches import ( + ContextProjectionPatch, + GrootDiTStepEncoderPatch, + TRTDynamicCategorySpecificMLPPatch, +) from torch_tensorrt.hf.exporters.ops import call_engine, scatter_image_tokens from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, EdgeSpec, register_edge_spec, ) -from torch_tensorrt.hf.exporters.specs._language import ( - causal_lm_flat, - kv_kwargs, - split_flat_to_kwargs, -) - - -def _groot(model: nn.Module) -> nn.Module: - if hasattr(model, "_groot_model"): - return model._groot_model - backbone = getattr(model, "backbone", None) - if backbone is not None and hasattr(backbone, "eagle_model"): - return model - raise RuntimeError( - "GR00T spec expected GrootPolicy or a module with backbone.eagle_model" - ) @register_edge_spec("groot", "gr00t") @@ -46,7 +50,6 @@ def prepare_sample_inputs( load_test_data, pack_state, ) - from torch_tensorrt.hf.exporters.helpers.groot import make_embodiment_id policy = model device = raw.get( @@ -107,18 +110,6 @@ def prepare_sample_inputs( def wrap( self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any ) -> nn.Module: - from torch_tensorrt.hf.exporters.patches.diffusion import ( - GrootDiTStepEncoderPatch, - StaticActionVelocityStepPatch, - TRTDynamicCategorySpecificMLPPatch, - ) - from torch_tensorrt.hf.exporters.patches.language import ( - CausalLMPatch, - ContextProjectionPatch, - language_decoder, - ) - from torch_tensorrt.hf.exporters.patches.vision import GridVisionPatch - found = _groot(model) eagle = found.backbone.eagle_model device = sample["pixel_values"].device diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py b/py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py @@ -0,0 +1 @@ + diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py b/py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py new file mode 100644 index 00000000000..67a80f49c54 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def _decoder(model: nn.Module) -> nn.Module: + return getattr(model, "backbone", None) or model.model + + +def _kind(mixer: nn.Module) -> str: + name = type(mixer).__name__ + if "Mamba" in name: + return "mamba" + if "Attention" in name: + return "attention" + if "MoE" in name or "Moe" in name: + return "moe" + return "mlp" + + +def allocate_plugin_states( + model: nn.Module, + config: Any, + batch: int, + max_seq_len: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + kinds = [_kind(block.mixer) for block in _decoder(model).layers] + head_dim = int( + getattr(config, "head_dim", 0) + or config.hidden_size // config.num_attention_heads + ) + conv_dim = int(config.mamba_num_heads) * int(config.mamba_head_dim) + 2 * int( + config.n_groups + ) * int(config.ssm_state_size) + conv_kernel = int(getattr(config, "conv_kernel", 4)) + kvs, convs, ssms = [], [], [] + for kind in kinds: + if kind == "attention": + kvs.append( + torch.zeros( + batch, + 2, + int(config.num_key_value_heads), + max_seq_len, + head_dim, + device=device, + dtype=dtype, + ) + ) + elif kind == "mamba": + convs.append( + torch.zeros(batch, conv_dim, conv_kernel, device=device, dtype=dtype) + ) + ssms.append( + torch.zeros( + batch, + int(config.mamba_num_heads), + int(config.mamba_head_dim), + int(config.ssm_state_size), + device=device, + dtype=dtype, + ) + ) + return kvs, convs, ssms diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py b/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py new file mode 100644 index 00000000000..c0850837685 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.models.common.patches import gather_last_token_hidden +from torch_tensorrt.hf.exporters.models.nemotron.helpers import _decoder, _kind + + +class NemotronPatch(nn.Module): # type: ignore[misc] + """Hybrid decoder: plugin attention / mamba / moe, native MLP.""" + + def __init__(self, model: nn.Module): + super().__init__() + decoder = _decoder(model) + self.layers = decoder.layers + self.norm = decoder.norm_f + self.lm_head = model.lm_head + self.kinds = [_kind(block.mixer) for block in self.layers] + self.num_attn = self.kinds.count("attention") + self.num_mamba = self.kinds.count("mamba") + + def forward( + self, + inputs_embeds: torch.Tensor, + rope_rotary_cos_sin: torch.Tensor, + context_lengths: torch.Tensor, + kvcache_start_index: torch.Tensor, + last_token_ids: torch.Tensor, + *states: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + na, nm = self.num_attn, self.num_mamba + kvs = list(states[:na]) + convs = list(states[na : na + nm]) + ssms = list(states[na + nm : na + 2 * nm]) + kv_i = conv_i = 0 + hidden = inputs_embeds + present_kv, present_conv, present_ssm = [], [], [] + for block, kind in zip(self.layers, self.kinds): + residual = hidden + hidden = block.norm(hidden) + mixer = block.mixer + if kind == "attention": + hidden, kv = mixer( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=kvs[kv_i], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, + ) + present_kv.append(kv) + kv_i += 1 + elif kind == "mamba": + hidden, conv_out, ssm_out = mixer( + hidden, convs[conv_i], ssms[conv_i], context_lengths + ) + present_conv.append(conv_out) + present_ssm.append(ssm_out) + conv_i += 1 + else: + hidden = mixer(hidden) + hidden = residual + hidden + hidden = self.norm(hidden) + last = gather_last_token_hidden(hidden, last_token_ids) + logits = self.lm_head(last).float() + return (logits, *present_kv, *present_conv, *present_ssm) diff --git a/py/torch_tensorrt/hf/exporters/specs/nemotron.py b/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py similarity index 51% rename from py/torch_tensorrt/hf/exporters/specs/nemotron.py rename to py/torch_tensorrt/hf/exporters/models/nemotron/spec.py index 84887734d1f..36010b65780 100644 --- a/py/torch_tensorrt/hf/exporters/specs/nemotron.py +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py @@ -5,139 +5,21 @@ import torch import torch.nn as nn +from torch_tensorrt.hf.exporters.models.common.helpers import ( + kv_kwargs, + split_flat_to_kwargs, +) +from torch_tensorrt.hf.exporters.models.nemotron.helpers import ( + _decoder, + allocate_plugin_states, +) +from torch_tensorrt.hf.exporters.models.nemotron.patches import NemotronPatch from torch_tensorrt.hf.exporters.ops import call_engine from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, EdgeSpec, register_edge_spec, ) -from torch_tensorrt.hf.exporters.specs._language import kv_kwargs, split_flat_to_kwargs - - -def _decoder(model: nn.Module) -> nn.Module: - return getattr(model, "backbone", None) or model.model - - -def _kind(mixer: nn.Module) -> str: - name = type(mixer).__name__ - if "Mamba" in name: - return "mamba" - if "Attention" in name: - return "attention" - if "MoE" in name or "Moe" in name: - return "moe" - return "mlp" - - -class NemotronPatch(nn.Module): # type: ignore[misc] - """Hybrid decoder: plugin attention / mamba / moe, native MLP.""" - - def __init__(self, model: nn.Module): - super().__init__() - decoder = _decoder(model) - self.layers = decoder.layers - self.norm = decoder.norm_f - self.lm_head = model.lm_head - self.kinds = [_kind(block.mixer) for block in self.layers] - self.num_attn = self.kinds.count("attention") - self.num_mamba = self.kinds.count("mamba") - - def forward( - self, - inputs_embeds: torch.Tensor, - rope_rotary_cos_sin: torch.Tensor, - context_lengths: torch.Tensor, - kvcache_start_index: torch.Tensor, - last_token_ids: torch.Tensor, - *states: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - from torch_tensorrt.hf.exporters.patches.language import ( - gather_last_token_hidden, - ) - - na, nm = self.num_attn, self.num_mamba - kvs = list(states[:na]) - convs = list(states[na : na + nm]) - ssms = list(states[na + nm : na + 2 * nm]) - kv_i = conv_i = 0 - hidden = inputs_embeds - present_kv, present_conv, present_ssm = [], [], [] - for block, kind in zip(self.layers, self.kinds): - residual = hidden - hidden = block.norm(hidden) - mixer = block.mixer - if kind == "attention": - hidden, kv = mixer( - hidden_states=hidden, - rope_rotary_cos_sin=rope_rotary_cos_sin, - past_key_value=kvs[kv_i], - ctx_len=context_lengths, - kvcache_start_index=kvcache_start_index, - ) - present_kv.append(kv) - kv_i += 1 - elif kind == "mamba": - hidden, conv_out, ssm_out = mixer( - hidden, convs[conv_i], ssms[conv_i], context_lengths - ) - present_conv.append(conv_out) - present_ssm.append(ssm_out) - conv_i += 1 - else: - hidden = mixer(hidden) - hidden = residual + hidden - hidden = self.norm(hidden) - last = gather_last_token_hidden(hidden, last_token_ids) - logits = self.lm_head(last).float() - return (logits, *present_kv, *present_conv, *present_ssm) - - -def allocate_plugin_states( - model: nn.Module, - config: Any, - batch: int, - max_seq_len: int, - device: torch.device, - dtype: torch.dtype, -) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: - kinds = [_kind(block.mixer) for block in _decoder(model).layers] - head_dim = int( - getattr(config, "head_dim", 0) - or config.hidden_size // config.num_attention_heads - ) - conv_dim = int(config.mamba_num_heads) * int(config.mamba_head_dim) + 2 * int( - config.n_groups - ) * int(config.ssm_state_size) - conv_kernel = int(getattr(config, "conv_kernel", 4)) - kvs, convs, ssms = [], [], [] - for kind in kinds: - if kind == "attention": - kvs.append( - torch.zeros( - batch, - 2, - int(config.num_key_value_heads), - max_seq_len, - head_dim, - device=device, - dtype=dtype, - ) - ) - elif kind == "mamba": - convs.append( - torch.zeros(batch, conv_dim, conv_kernel, device=device, dtype=dtype) - ) - ssms.append( - torch.zeros( - batch, - int(config.mamba_num_heads), - int(config.mamba_head_dim), - int(config.ssm_state_size), - device=device, - dtype=dtype, - ) - ) - return kvs, convs, ssms @register_edge_spec("nemotron_h", "nemotron") diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/__init__.py b/py/torch_tensorrt/hf/exporters/models/pi05/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/pi05/__init__.py @@ -0,0 +1 @@ + diff --git a/py/torch_tensorrt/hf/exporters/helpers/pi05.py b/py/torch_tensorrt/hf/exporters/models/pi05/helpers.py similarity index 92% rename from py/torch_tensorrt/hf/exporters/helpers/pi05.py rename to py/torch_tensorrt/hf/exporters/models/pi05/helpers.py index 5e15091000a..5a1608837c6 100644 --- a/py/torch_tensorrt/hf/exporters/helpers/pi05.py +++ b/py/torch_tensorrt/hf/exporters/models/pi05/helpers.py @@ -1,6 +1,7 @@ from __future__ import annotations import torch +import torch.nn as nn from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks @@ -222,3 +223,20 @@ def make_pi05_suffix_position_and_mask(core, prefix_pad_masks, x_t, device): prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 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 + if pixel_values.shape[1] in (1, 3, 4) and pixel_values.shape[-1] not in (1, 3, 4): + return pixel_values.permute(0, 2, 3, 1).contiguous() + return pixel_values diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/patches.py b/py/torch_tensorrt/hf/exporters/models/pi05/patches.py new file mode 100644 index 00000000000..55e59467fc2 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/models/pi05/patches.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import torch.nn.functional as F +from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding +from torch_tensorrt.hf.exporters.models.common.patches import ActionStepEncoderPatch +from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache + + +class PI05PrefixKVStepEncoderPatch(ActionStepEncoderPatch): + """PI05 suffix embed + AdaRMS cond, consumed by Gemma action expert.""" + + def __init__(self, core): + super().__init__() + self.action_in_proj = core.action_in_proj + self.time_mlp_in = core.time_mlp_in + self.time_mlp_out = core.time_mlp_out + self.config = core.config + self.hidden_size = core.action_in_proj.out_features + + def forward( + self, + x_t, + timestep, + prefix_k, + prefix_v, + position_ids, + attention_mask, + ): + suffix_embs = self.action_in_proj(x_t) + + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.hidden_size, + min_period=self.config.min_period, + max_period=self.config.max_period, + device=timestep.device, + ).to(dtype=suffix_embs.dtype) + + adarms_cond = self.time_mlp_in(time_emb) + adarms_cond = F.silu(adarms_cond) + adarms_cond = self.time_mlp_out(adarms_cond) + adarms_cond = F.silu(adarms_cond) + + expert_kwargs = { + "inputs_embeds": suffix_embs, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values": PrefixKVCache(prefix_k, prefix_v), + "use_cache": False, + "adarms_cond": adarms_cond, + } + return (), expert_kwargs, (), {} diff --git a/py/torch_tensorrt/hf/exporters/specs/pi05.py b/py/torch_tensorrt/hf/exporters/models/pi05/spec.py similarity index 89% rename from py/torch_tensorrt/hf/exporters/specs/pi05.py rename to py/torch_tensorrt/hf/exporters/models/pi05/spec.py index 5d1baa3b8ad..63feb4670f8 100644 --- a/py/torch_tensorrt/hf/exporters/specs/pi05.py +++ b/py/torch_tensorrt/hf/exporters/models/pi05/spec.py @@ -5,34 +5,31 @@ import torch import torch.nn as nn +from torch_tensorrt.hf.exporters.models.common.helpers import ( + causal_lm_flat, + kv_kwargs, + split_flat_to_kwargs, +) +from torch_tensorrt.hf.exporters.models.common.patches import ( + CausalLMPatch, + GridVisionPatch, + StaticActionVelocityStepPatch, + language_decoder, +) +from torch_tensorrt.hf.exporters.models.pi05.helpers import ( + _core, + _nchw_to_hwc, + build_pi05_prefix_embs, + make_pi05_suffix_position_and_mask, + pi05_compact_index, +) +from torch_tensorrt.hf.exporters.models.pi05.patches import PI05PrefixKVStepEncoderPatch from torch_tensorrt.hf.exporters.ops import call_engine, fuse_prefix from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, EdgeSpec, register_edge_spec, ) -from torch_tensorrt.hf.exporters.specs._language import ( - causal_lm_flat, - kv_kwargs, - split_flat_to_kwargs, -) - - -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: torch.Tensor) -> torch.Tensor: - if pixel_values.ndim != 4: - return pixel_values - if pixel_values.shape[1] in (1, 3, 4) and pixel_values.shape[-1] not in (1, 3, 4): - return pixel_values.permute(0, 2, 3, 1).contiguous() - return pixel_values @register_edge_spec("pi05") @@ -93,16 +90,6 @@ def prepare_sample_inputs( def wrap( self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any ) -> nn.Module: - from torch_tensorrt.hf.exporters.patches.diffusion import ( - PI05PrefixKVStepEncoderPatch, - StaticActionVelocityStepPatch, - ) - from torch_tensorrt.hf.exporters.patches.language import ( - CausalLMPatch, - language_decoder, - ) - from torch_tensorrt.hf.exporters.patches.vision import GridVisionPatch - core = _core(model) paligemma = core.paligemma_with_expert.paligemma.model if name == "vision": @@ -140,11 +127,6 @@ def prepare( config: Any, module: nn.Module, ) -> ComponentBundle: - from torch_tensorrt.hf.exporters.helpers.pi05 import ( - build_pi05_prefix_embs, - make_pi05_suffix_position_and_mask, - pi05_compact_index, - ) from torch_tensorrt.hf.exporters.plugin.attention import ( ContextAttentionMaskType, ) diff --git a/py/torch_tensorrt/hf/exporters/patches/language.py b/py/torch_tensorrt/hf/exporters/patches/language.py deleted file mode 100644 index 0ceffc87640..00000000000 --- a/py/torch_tensorrt/hf/exporters/patches/language.py +++ /dev/null @@ -1,148 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn as nn - - -def _as_tensor(x): - """Unwrap tuple/list outputs from patched attention modules.""" - if isinstance(x, (tuple, list)): - return x[0] - return x - - -def language_decoder(language: nn.Module) -> nn.Module: - """Inner module that owns ``.layers``. - - Paligemma / PI05 store layers on ``language_model`` itself. HF - ``*ForCausalLM`` stores them on ``language_model.model``. Prefer ``.layers`` - so a stray ``.model`` attribute cannot silently pick the wrong submodule. - """ - if hasattr(language, "layers"): - return language - inner = getattr(language, "model", None) - if isinstance(inner, nn.Module) and hasattr(inner, "layers"): - return inner - 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] - - -class CausalLMPatch(nn.Module): - """Edge-LLM causal LM: manual decoder loop -> logits + lm_hidden_states + prefix KV. - - External RoPE and KV-cache controls match ``LLMEngineRunner`` prefill/decode. - ``select_layer=-1`` (default for GR00T TRT) uses final RMSNorm hidden for - context; positive values capture an intermediate layer output pre-norm. - - ``ds_stack`` is always an input (``[num_ds, B, S, H]``). Models without - deepstack pass ``num_ds=0`` (empty leading dim); the add is a no-op. - """ - - def __init__( - self, - lm: nn.Module, - lm_head: nn.Module, - *, - select_layer: int = -1, - ): - super().__init__() - self.lm = lm - self.lm_head = lm_head - self.select_layer = int(select_layer) - - def forward( - self, - inputs_embeds: torch.Tensor, - rope_rotary_cos_sin: torch.Tensor, - context_lengths: torch.Tensor, - kvcache_start_index: torch.Tensor, - last_token_ids: torch.Tensor, - ds_stack: torch.Tensor, - *past_key_values: torch.Tensor, - ): - lm_dtype = next(self.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 self.select_layer == 0 else None - new_kvs = [] - - for i, layer in enumerate(self.lm.layers): - residual = hidden - hidden = _as_tensor(layer.input_layernorm(hidden)) - hidden, kv = layer.self_attn( - hidden_states=hidden, - rope_rotary_cos_sin=rope_rotary_cos_sin, - past_key_value=past_key_values[i], - ctx_len=context_lengths, - kvcache_start_index=kvcache_start_index, - ) - hidden = _as_tensor(hidden) - hidden = residual + hidden - - residual = hidden - hidden = _as_tensor(layer.post_attention_layernorm(hidden)) - hidden = _as_tensor(layer.mlp(hidden)) - hidden = residual + hidden - new_kvs.append(kv) - - if i < num_ds: - hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) - - if self.select_layer > 0 and (i + 1) == self.select_layer: - context_hidden = hidden - - hidden = _as_tensor(self.lm.norm(hidden)) - if context_hidden is None: - context_hidden = hidden - - last_hidden = gather_last_token_hidden(hidden, last_token_ids) - logits = self.lm_head(last_hidden).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 - - -# specific to gr00t, before action another project is required for context embeddings -class ContextProjectionPatch(nn.Module): - """eagle_linear -> vlln -> vl_self_attention (matches eager context path).""" - - def __init__(self, eagle_linear, vlln, vl_self_attention): - super().__init__() - self.eagle_linear = eagle_linear - self.vlln = vlln - self.vl_self_attention = vl_self_attention - - def forward(self, hidden_states: torch.Tensor): - context_embs = self.eagle_linear(hidden_states) - - vlln_weight = getattr(self.vlln, "weight", None) - if vlln_weight is not None: - context_embs = context_embs.to(dtype=vlln_weight.dtype) - - context_embs = self.vlln(context_embs) - context_embs = self.vl_self_attention(context_embs) - return context_embs diff --git a/py/torch_tensorrt/hf/exporters/patches/vision.py b/py/torch_tensorrt/hf/exporters/patches/vision.py deleted file mode 100644 index 152570de40a..00000000000 --- a/py/torch_tensorrt/hf/exporters/patches/vision.py +++ /dev/null @@ -1,158 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import torch -import torch.nn as nn - - -def hwc_to_nchw(images: torch.Tensor) -> torch.Tensor: - if images.ndim != 4: - raise ValueError(f"Expected 4D images, got shape {tuple(images.shape)}") - return images.permute(0, 3, 1, 2).contiguous() - - -def is_nchw_pixel_values(pixel_values: torch.Tensor) -> bool: - return ( - pixel_values.ndim == 4 - and pixel_values.shape[1] in (1, 3, 4) - and pixel_values.shape[-1] not in (1, 3, 4) - ) - - -# --------------------------------------------------------------------------- -# Base -# --------------------------------------------------------------------------- - - -class VisionPatch(nn.Module): - """Base TRT trace target: probe static shapes once, flatten output to [N, H].""" - - cast_output_to_input_dtype: bool - output_num_tokens: int - output_hidden_size: int - - def _finalize_output( - self, features: torch.Tensor, out_dtype: torch.dtype - ) -> torch.Tensor: - if self.cast_output_to_input_dtype and features.dtype != out_dtype: - features = features.to(out_dtype) - if features.ndim == 3: - # [B, S, H] -> [B*S, H] - return features.reshape(-1, features.shape[-1]) - if features.ndim == 2: - # already [N, H] (token-pooling encoders) - return features - raise ValueError( - f"Expected 2D or 3D features, got shape {tuple(features.shape)}" - ) - - -# --------------------------------------------------------------------------- -# Grid vision (PI0.5 / GR00T / SmolVLA VitRunner path) -# --------------------------------------------------------------------------- - - -class GridVisionPatch(VisionPatch): - """pixels [B,H,W,C] -> fixed patch grid -> [B*seq_len, lm_hidden]""" - - def __init__( - self, - *, - vision_model: nn.Module, - projector: nn.Module, - sample_pixel_values: torch.Tensor, - select_layer: int = -1, - pixel_shuffle: bool = False, - downsample_ratio: float = 0.5, - force_float32_input: bool = False, - cast_output_to_input_dtype: bool = False, - vision_kwargs: dict[str, Any] | None = None, - ): - super().__init__() - self.vision_model = vision_model - self.projector = projector - self.select_layer = int(select_layer) - self.pixel_shuffle = bool(pixel_shuffle) - self.downsample_ratio = float(downsample_ratio) - self.force_float32_input = bool(force_float32_input) - self.cast_output_to_input_dtype = bool(cast_output_to_input_dtype) - self.vision_kwargs = dict(vision_kwargs or {}) - - with torch.no_grad(): - sample = sample_pixel_values - if self.force_float32_input and sample.dtype != torch.float32: - sample = sample.to(torch.float32) - - vit_embeds = self._select_vision_features(self._run_vision(sample)) - self.seq_len = int(vit_embeds.shape[1]) - self.hidden_size = int(vit_embeds.shape[2]) - - if self.pixel_shuffle: - self._init_pixel_shuffle_shape() - vit_embeds = self._apply_pixel_shuffle(vit_embeds) - - projected = self.projector(self._projector_input(vit_embeds)) - self.batch_size = int(projected.shape[0]) - self.output_seq_len = int(projected.shape[1]) - self.output_hidden_size = int(projected.shape[2]) - self.output_num_tokens = self.batch_size * self.output_seq_len - - def _run_vision(self, images: torch.Tensor): - pixel_values = ( - hwc_to_nchw(images) if not is_nchw_pixel_values(images) else images - ) - kwargs = dict(self.vision_kwargs) - kwargs["pixel_values"] = pixel_values - kwargs["output_hidden_states"] = self.select_layer != -1 - kwargs.setdefault("return_dict", True) - return self.vision_model(**kwargs) - - def _select_vision_features(self, out): - if self.select_layer == -1: - return ( - out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] - ) - return out.hidden_states[self.select_layer] - - def _projector_input(self, vit_embeds: torch.Tensor) -> torch.Tensor: - proj_dtype = next(self.projector.parameters()).dtype - if vit_embeds.dtype != proj_dtype: - return vit_embeds.to(proj_dtype) - return vit_embeds - - def _init_pixel_shuffle_shape(self): - side = int(self.seq_len**0.5) - if side * side != self.seq_len: - raise ValueError( - f"Expected square vision sequence, got seq_len={self.seq_len}" - ) - self.grid_w = side - self.grid_h = side - self.out_w = int(self.grid_w * self.downsample_ratio) - self.out_h = int(self.grid_h * self.downsample_ratio) - self.hidden_after_first_view = int(self.hidden_size / self.downsample_ratio) - self.shuffle_hidden = int( - self.hidden_size / (self.downsample_ratio * self.downsample_ratio) - ) - - def _apply_pixel_shuffle(self, x): - n = x.shape[0] - x = x.reshape(n, self.grid_w, self.out_h, self.hidden_after_first_view) - x = x.permute(0, 2, 1, 3).contiguous() - x = x.reshape(n, self.out_h, self.out_w, self.shuffle_hidden) - x = x.permute(0, 2, 1, 3).contiguous() - return x.reshape(n, -1, self.shuffle_hidden) - - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - out_dtype = pixel_values.dtype - images = pixel_values - if self.force_float32_input and images.dtype != torch.float32: - images = images.to(torch.float32) - - vit_embeds = self._select_vision_features(self._run_vision(images)) - if self.pixel_shuffle: - vit_embeds = self._apply_pixel_shuffle(vit_embeds) - - features = self.projector(self._projector_input(vit_embeds)) - return self._finalize_output(features, out_dtype) diff --git a/py/torch_tensorrt/hf/exporters/specs/__init__.py b/py/torch_tensorrt/hf/exporters/specs/__init__.py deleted file mode 100644 index ed8ebec6682..00000000000 --- a/py/torch_tensorrt/hf/exporters/specs/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Family specs. Import is a side effect: ``@register_edge_spec`` attaches.""" - -from torch_tensorrt.hf.exporters.specs import groot, nemotron, pi05 - -__all__ = ["groot", "nemotron", "pi05"] diff --git a/pyproject.toml b/pyproject.toml index 8fe1ce6b7db..d580cf19753 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -454,8 +454,7 @@ follow_imports = "skip" [[tool.mypy.overrides]] module = [ "torch_tensorrt.hf.exporters.plugin.*", - "torch_tensorrt.hf.exporters.patches.*", - "torch_tensorrt.hf.exporters.helpers.*", + "torch_tensorrt.hf.exporters.models.*", "torch_tensorrt.hf.exporters.data", "torch_tensorrt.hf.exporters.rope", "torch_tensorrt.hf.exporters.prefix_cache", diff --git a/setup.py b/setup.py index 8b08069e475..64a039f61fa 100644 --- a/setup.py +++ b/setup.py @@ -618,10 +618,12 @@ def run(self): "torch_tensorrt.executorch", "torch_tensorrt.hf", "torch_tensorrt.hf.exporters", - "torch_tensorrt.hf.exporters.helpers", - "torch_tensorrt.hf.exporters.patches", + "torch_tensorrt.hf.exporters.models", + "torch_tensorrt.hf.exporters.models.common", + "torch_tensorrt.hf.exporters.models.groot", + "torch_tensorrt.hf.exporters.models.nemotron", + "torch_tensorrt.hf.exporters.models.pi05", "torch_tensorrt.hf.exporters.plugin", - "torch_tensorrt.hf.exporters.specs", "torch_tensorrt.runtime", ] @@ -663,10 +665,12 @@ def run(self): "torch_tensorrt.executorch": "py/torch_tensorrt/executorch", "torch_tensorrt.hf": "py/torch_tensorrt/hf", "torch_tensorrt.hf.exporters": "py/torch_tensorrt/hf/exporters", - "torch_tensorrt.hf.exporters.helpers": "py/torch_tensorrt/hf/exporters/helpers", - "torch_tensorrt.hf.exporters.patches": "py/torch_tensorrt/hf/exporters/patches", + "torch_tensorrt.hf.exporters.models": "py/torch_tensorrt/hf/exporters/models", + "torch_tensorrt.hf.exporters.models.common": "py/torch_tensorrt/hf/exporters/models/common", + "torch_tensorrt.hf.exporters.models.groot": "py/torch_tensorrt/hf/exporters/models/groot", + "torch_tensorrt.hf.exporters.models.nemotron": "py/torch_tensorrt/hf/exporters/models/nemotron", + "torch_tensorrt.hf.exporters.models.pi05": "py/torch_tensorrt/hf/exporters/models/pi05", "torch_tensorrt.hf.exporters.plugin": "py/torch_tensorrt/hf/exporters/plugin", - "torch_tensorrt.hf.exporters.specs": "py/torch_tensorrt/hf/exporters/specs", "torch_tensorrt.runtime": "py/torch_tensorrt/runtime", } From d95036e29c756efb237722b34090331a5ee44d6f Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 15:52:55 -0700 Subject: [PATCH 04/11] Replace Edge wrapper modules with family setattr patches. Compile the original HF/LeRobot submodules under apply_patches and stitch them with execute_engine, matching HF DynamoExporter. --- py/torch_tensorrt/hf/exporters/compile.py | 18 +- py/torch_tensorrt/hf/exporters/exporter.py | 43 +- .../hf/exporters/models/common/patches.py | 375 +++--------------- .../hf/exporters/models/groot/patches.py | 357 +++++++---------- .../hf/exporters/models/groot/spec.py | 119 ++---- .../hf/exporters/models/nemotron/patches.py | 79 ++-- .../hf/exporters/models/nemotron/spec.py | 28 +- .../hf/exporters/models/pi05/patches.py | 157 ++++++-- .../hf/exporters/models/pi05/spec.py | 88 +--- py/torch_tensorrt/hf/exporters/ops.py | 3 + .../hf/exporters/plugin/attn_patches.py | 287 ++++++++++++++ .../hf/exporters/plugin/plugin_utils.py | 2 + py/torch_tensorrt/hf/exporters/spec.py | 33 +- tests/py/dynamo/hf/test_edge_exporter.py | 309 ++++++++++++++- 14 files changed, 1109 insertions(+), 789 deletions(-) create mode 100644 py/torch_tensorrt/hf/exporters/plugin/attn_patches.py diff --git a/py/torch_tensorrt/hf/exporters/compile.py b/py/torch_tensorrt/hf/exporters/compile.py index df5c5c5f7a4..a8b9dee20d6 100644 --- a/py/torch_tensorrt/hf/exporters/compile.py +++ b/py/torch_tensorrt/hf/exporters/compile.py @@ -5,7 +5,6 @@ from typing import Any import torch -import torch.nn as nn import torch_tensorrt from torch_tensorrt.hf.exporters.ops import _as_tuple, record_engine from torch_tensorrt.hf.exporters.spec import ComponentBundle @@ -28,7 +27,6 @@ def compile_component( - module: nn.Module, bundle: ComponentBundle, *, name: str, @@ -36,16 +34,19 @@ def compile_component( dryrun: bool = False, trt_settings: dict[str, Any] | None = None, ) -> tuple[str, tuple[torch.Tensor, ...]]: - """Export one wrapper, compile it, write ``engine_dir//``. + """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. - ``dryrun`` records the patched eager module for ``execute_engine`` and - leaves the patch in place. A real compile restores HF attention after the - TensorRT module is recorded. + Family setattr is owned by ``EdgeSpec.apply_patches``, not this helper. + ``dryrun`` records the patched eager module for ``execute_engine``. """ - module = module.eval() + from torch_tensorrt.hf.exporters.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) @@ -53,6 +54,9 @@ def compile_component( 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) + patched = bundle.patch_fn(module) if bundle.patch_fn is not None else None try: with torch.no_grad(): diff --git a/py/torch_tensorrt/hf/exporters/exporter.py b/py/torch_tensorrt/hf/exporters/exporter.py index 160c9a119cd..f619ac01b23 100644 --- a/py/torch_tensorrt/hf/exporters/exporter.py +++ b/py/torch_tensorrt/hf/exporters/exporter.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextlib import copy import inspect import logging @@ -41,6 +42,7 @@ def __init__(self) -> None: self.engines: dict[str, str] = {} self.runtime: EdgeRuntimeModule | None = None self.sample: dict[str, Any] = {} + self._dryrun_patches: contextlib.ExitStack | None = None def export( self, @@ -53,11 +55,15 @@ def export( 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) engine_dir = Path(config.engine_dir or "edge_engines") @@ -66,19 +72,31 @@ def export( engines: dict[str, str] = {} upstream: dict[str, Any] = {} - for name in names: - module = spec.wrap(name, model, sample, config) - bundle = spec.prepare(name, model, sample, upstream, config, module) - engines[name], outs = compile_component( - module, - bundle, - name=name, - engine_dir=engine_dir, - dryrun=config.dryrun, - trt_settings=config.trt_settings, - ) - upstream.update(spec.capture_upstream(name, outs, sample, bundle)) + def _compile_components() -> None: + for name in names: + bundle = spec.prepare(name, model, sample, upstream, config) + engines[name], outs = compile_component( + bundle, + name=name, + engine_dir=engine_dir, + dryrun=config.dryrun, + trt_settings=config.trt_settings, + ) + upstream.update(spec.capture_upstream(name, outs, sample, bundle)) + + # 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)) @@ -88,6 +106,7 @@ def export( 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( diff --git a/py/torch_tensorrt/hf/exporters/models/common/patches.py b/py/torch_tensorrt/hf/exporters/models/common/patches.py index 3d46d8b9985..c9434f0a278 100644 --- a/py/torch_tensorrt/hf/exporters/models/common/patches.py +++ b/py/torch_tensorrt/hf/exporters/models/common/patches.py @@ -1,161 +1,8 @@ from __future__ import annotations -from typing import Any - import torch import torch.nn as nn - - -def hwc_to_nchw(images: torch.Tensor) -> torch.Tensor: - if images.ndim != 4: - raise ValueError(f"Expected 4D images, got shape {tuple(images.shape)}") - return images.permute(0, 3, 1, 2).contiguous() - - -def is_nchw_pixel_values(pixel_values: torch.Tensor) -> bool: - return ( - pixel_values.ndim == 4 - and pixel_values.shape[1] in (1, 3, 4) - and pixel_values.shape[-1] not in (1, 3, 4) - ) - - -# --------------------------------------------------------------------------- -# Base -# --------------------------------------------------------------------------- - - -class VisionPatch(nn.Module): - """Base TRT trace target: probe static shapes once, flatten output to [N, H].""" - - cast_output_to_input_dtype: bool - output_num_tokens: int - output_hidden_size: int - - def _finalize_output( - self, features: torch.Tensor, out_dtype: torch.dtype - ) -> torch.Tensor: - if self.cast_output_to_input_dtype and features.dtype != out_dtype: - features = features.to(out_dtype) - if features.ndim == 3: - # [B, S, H] -> [B*S, H] - return features.reshape(-1, features.shape[-1]) - if features.ndim == 2: - # already [N, H] (token-pooling encoders) - return features - raise ValueError( - f"Expected 2D or 3D features, got shape {tuple(features.shape)}" - ) - - -# --------------------------------------------------------------------------- -# Grid vision (PI0.5 / GR00T / SmolVLA VitRunner path) -# --------------------------------------------------------------------------- - - -class GridVisionPatch(VisionPatch): - """pixels [B,H,W,C] -> fixed patch grid -> [B*seq_len, lm_hidden]""" - - def __init__( - self, - *, - vision_model: nn.Module, - projector: nn.Module, - sample_pixel_values: torch.Tensor, - select_layer: int = -1, - pixel_shuffle: bool = False, - downsample_ratio: float = 0.5, - force_float32_input: bool = False, - cast_output_to_input_dtype: bool = False, - vision_kwargs: dict[str, Any] | None = None, - ): - super().__init__() - self.vision_model = vision_model - self.projector = projector - self.select_layer = int(select_layer) - self.pixel_shuffle = bool(pixel_shuffle) - self.downsample_ratio = float(downsample_ratio) - self.force_float32_input = bool(force_float32_input) - self.cast_output_to_input_dtype = bool(cast_output_to_input_dtype) - self.vision_kwargs = dict(vision_kwargs or {}) - - with torch.no_grad(): - sample = sample_pixel_values - if self.force_float32_input and sample.dtype != torch.float32: - sample = sample.to(torch.float32) - - vit_embeds = self._select_vision_features(self._run_vision(sample)) - self.seq_len = int(vit_embeds.shape[1]) - self.hidden_size = int(vit_embeds.shape[2]) - - if self.pixel_shuffle: - self._init_pixel_shuffle_shape() - vit_embeds = self._apply_pixel_shuffle(vit_embeds) - - projected = self.projector(self._projector_input(vit_embeds)) - self.batch_size = int(projected.shape[0]) - self.output_seq_len = int(projected.shape[1]) - self.output_hidden_size = int(projected.shape[2]) - self.output_num_tokens = self.batch_size * self.output_seq_len - - def _run_vision(self, images: torch.Tensor): - pixel_values = ( - hwc_to_nchw(images) if not is_nchw_pixel_values(images) else images - ) - kwargs = dict(self.vision_kwargs) - kwargs["pixel_values"] = pixel_values - kwargs["output_hidden_states"] = self.select_layer != -1 - kwargs.setdefault("return_dict", True) - return self.vision_model(**kwargs) - - def _select_vision_features(self, out): - if self.select_layer == -1: - return ( - out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] - ) - return out.hidden_states[self.select_layer] - - def _projector_input(self, vit_embeds: torch.Tensor) -> torch.Tensor: - proj_dtype = next(self.projector.parameters()).dtype - if vit_embeds.dtype != proj_dtype: - return vit_embeds.to(proj_dtype) - return vit_embeds - - def _init_pixel_shuffle_shape(self): - side = int(self.seq_len**0.5) - if side * side != self.seq_len: - raise ValueError( - f"Expected square vision sequence, got seq_len={self.seq_len}" - ) - self.grid_w = side - self.grid_h = side - self.out_w = int(self.grid_w * self.downsample_ratio) - self.out_h = int(self.grid_h * self.downsample_ratio) - self.hidden_after_first_view = int(self.hidden_size / self.downsample_ratio) - self.shuffle_hidden = int( - self.hidden_size / (self.downsample_ratio * self.downsample_ratio) - ) - - def _apply_pixel_shuffle(self, x): - n = x.shape[0] - x = x.reshape(n, self.grid_w, self.out_h, self.hidden_after_first_view) - x = x.permute(0, 2, 1, 3).contiguous() - x = x.reshape(n, self.out_h, self.out_w, self.shuffle_hidden) - x = x.permute(0, 2, 1, 3).contiguous() - return x.reshape(n, -1, self.shuffle_hidden) - - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - out_dtype = pixel_values.dtype - images = pixel_values - if self.force_float32_input and images.dtype != torch.float32: - images = images.to(torch.float32) - - vit_embeds = self._select_vision_features(self._run_vision(images)) - if self.pixel_shuffle: - vit_embeds = self._apply_pixel_shuffle(vit_embeds) - - features = self.projector(self._projector_input(vit_embeds)) - return self._finalize_output(features, out_dtype) +import torch.nn.functional as F def _as_tensor(x): @@ -197,170 +44,68 @@ def gather_last_token_hidden( return hidden_states[batch_idx, indices] -class CausalLMPatch(nn.Module): - """Edge-LLM causal LM: manual decoder loop -> logits + lm_hidden_states + prefix KV. - - External RoPE and KV-cache controls match ``LLMEngineRunner`` prefill/decode. - ``select_layer=-1`` (default for GR00T TRT) uses final RMSNorm hidden for - context; positive values capture an intermediate layer output pre-norm. - - ``ds_stack`` is always an input (``[num_ds, B, S, H]``). Models without - deepstack pass ``num_ds=0`` (empty leading dim); the add is a no-op. - """ - - def __init__( - self, - lm: nn.Module, - lm_head: nn.Module, - *, - select_layer: int = -1, - ): - super().__init__() - self.lm = lm - self.lm_head = lm_head - self.select_layer = int(select_layer) - - def forward( - self, - inputs_embeds: torch.Tensor, - rope_rotary_cos_sin: torch.Tensor, - context_lengths: torch.Tensor, - kvcache_start_index: torch.Tensor, - last_token_ids: torch.Tensor, - ds_stack: torch.Tensor, - *past_key_values: torch.Tensor, - ): - lm_dtype = next(self.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 self.select_layer == 0 else None - new_kvs = [] - - for i, layer in enumerate(self.lm.layers): - residual = hidden - hidden = _as_tensor(layer.input_layernorm(hidden)) - hidden, kv = layer.self_attn( - hidden_states=hidden, - rope_rotary_cos_sin=rope_rotary_cos_sin, - past_key_value=past_key_values[i], - ctx_len=context_lengths, - kvcache_start_index=kvcache_start_index, - ) - hidden = _as_tensor(hidden) - hidden = residual + hidden - - residual = hidden - hidden = _as_tensor(layer.post_attention_layernorm(hidden)) - hidden = _as_tensor(layer.mlp(hidden)) - hidden = residual + hidden - new_kvs.append(kv) - - if i < num_ds: - hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) - - if self.select_layer > 0 and (i + 1) == self.select_layer: - context_hidden = hidden - - hidden = _as_tensor(self.lm.norm(hidden)) - if context_hidden is None: - context_hidden = hidden - - last_hidden = gather_last_token_hidden(hidden, last_token_ids) - logits = self.lm_head(last_hidden).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 - - -class ActionStepEncoderPatch(nn.Module): - """Base contract for model-specific action-step encoding. - - Subclasses implement forward() to turn noisy actions, timestep, and - model-specific context tensors into the args/kwargs consumed by the action - expert and velocity decoder. The default helpers cover common expert output - and velocity shapes, while model-specific encoders can override them. - """ - - def get_action_hidden(self, expert_out, output_tokens: int): - # Default path for experts that return either a standard model output - # with last_hidden_state, a raw hidden-state tensor, or a tuple/list whose - # first item is the hidden-state tensor. - hidden = ( - expert_out.last_hidden_state - if hasattr(expert_out, "last_hidden_state") - else expert_out - ) - - if isinstance(hidden, (tuple, list)): - hidden = hidden[0] - - return hidden[:, -output_tokens:] - - def process_velocity(self, velocity): - # Default path for models whose decoder already returns the final action - # velocity shape. Override for models that need reshaping or cropping. - return velocity - - -class StaticActionVelocityStepPatch(nn.Module): - """One static denoising step shared by VLA action diffusion modules. - - The model-specific step_encoder owns the messy part: converting noisy - actions, timestep, and context tensors into the exact action_expert call. - This wrapper only runs the expert, selects action-token hidden states, and - decodes those hidden states into a velocity update. - """ - - def __init__( - self, - *, - step_encoder: ActionStepEncoderPatch, - action_expert: nn.Module, - velocity_decoder: nn.Module, - output_tokens: int, - cast_hidden_fp32: bool = True, - ): - super().__init__() - self.step_encoder = step_encoder - self.action_expert = action_expert - self.velocity_decoder = velocity_decoder - self.output_tokens = int(output_tokens) - self.cast_hidden_fp32 = cast_hidden_fp32 - - def forward(self, x_t, timestep, *inputs): - # Build the action expert inputs and any decoder-specific side inputs. - expert_args, expert_kwargs, decoder_args, decoder_kwargs = self.step_encoder( - x_t, - timestep, - *inputs, +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, + rope_rotary_cos_sin: torch.Tensor, + context_lengths: torch.Tensor, + kvcache_start_index: torch.Tensor, + last_token_ids: torch.Tensor, + ds_stack: torch.Tensor, + *past_key_values: torch.Tensor, + lm_head: nn.Module | None = None, + select_layer: int = -1, +): + """Prefill loop used by Edge language engines (plugin attention + prefix KV).""" + 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 = [] + + for i, layer in enumerate(lm.layers): + residual = hidden + hidden = _as_tensor(layer.input_layernorm(hidden)) + hidden, kv = layer.self_attn( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=past_key_values[i], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, ) + hidden = _as_tensor(hidden) + hidden = residual + hidden - # Run the model-specific action expert: Gemma expert, DiT, etc. - expert_out = self.action_expert(*expert_args, **expert_kwargs) + residual = hidden + hidden = _as_tensor(layer.post_attention_layernorm(hidden)) + hidden = _as_tensor(layer.mlp(hidden)) + hidden = residual + hidden + new_kvs.append(kv) - # Most experts return last_hidden_state, but some wrappers return tuples - # or need custom suffix-token selection. - action_hidden = self.step_encoder.get_action_hidden( - expert_out, - self.output_tokens, - ) + if i < num_ds: + hidden = hidden + ds_stack[i, :, :seq_len, :].to(dtype=hidden.dtype) - if self.cast_hidden_fp32: - action_hidden = action_hidden.to(dtype=torch.float32) + if select_layer > 0 and (i + 1) == select_layer: + context_hidden = hidden - # Project action-token hidden states back to action-space velocity. - velocity = self.velocity_decoder( - action_hidden, - *decoder_args, - **decoder_kwargs, - ) + hidden = _as_tensor(lm.norm(hidden)) + if context_hidden is None: + context_hidden = hidden - return self.step_encoder.process_velocity(velocity) + last_hidden = gather_last_token_hidden(hidden, last_token_ids) + logits = _lm_head_logits(lm, lm_head, last_hidden) + 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/py/torch_tensorrt/hf/exporters/models/groot/patches.py b/py/torch_tensorrt/hf/exporters/models/groot/patches.py index e882f1e25ed..b0eb154d095 100644 --- a/py/torch_tensorrt/hf/exporters/models/groot/patches.py +++ b/py/torch_tensorrt/hf/exporters/models/groot/patches.py @@ -1,179 +1,134 @@ -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch_tensorrt.hf.exporters.models.common.patches import ActionStepEncoderPatch - - -class TRTFixedCategorySpecificLinearPatch(nn.Module): - """Freeze one GR00T embodiment-specific Linear into a normal Linear. - - GR00T stores one weight matrix per robot embodiment and selects it with - embodiment_id at runtime. For TensorRT deployment we compile one robot at a - time, so this wrapper picks that robot's weights once in __init__ and the - forward path becomes a plain static F.linear. - """ - - def __init__(self, layer: nn.Module, embodiment_id: torch.Tensor): - super().__init__() - - cat_id = int(embodiment_id.flatten()[0].item()) - - # Original: [num_embodiments, input_dim, output_dim] - # using cat_id selects the weight matrix for one embodiment/robot -> [input_dim, output_dim] - # nn.functional.linear expects -> weight: [output_dim, input_dim] so we transpose - weight = layer.W[cat_id].transpose(0, 1).contiguous() - bias = layer.b[cat_id].contiguous() - - # detach() breaks any autograd link to the original multi-embodiment - # parameter; clone() gives this fixed wrapper independent storage for - # the selected slice. This copy happens once during wrapper creation, - # not in forward, and lets TensorRT see normal immutable weights. - self.weight = nn.Parameter(weight.detach().clone(), requires_grad=False) - self.bias = nn.Parameter(bias.detach().clone(), requires_grad=False) - self.out_features = int(bias.shape[0]) - - def forward(self, x): - # x: [B, T, input_dim] - # out: [B, T, output_dim] - batch_size = x.shape[0] - seq_len = x.shape[1] - x = x.reshape(batch_size * seq_len, x.shape[-1]) - x = F.linear(x, self.weight, self.bias) - return x.reshape(batch_size, seq_len, self.out_features) - - -class TRTDynamicCategorySpecificLinearPatch(nn.Module): - """TensorRT-friendly dynamic version of GR00T CategorySpecificLinear. - - Unlike the fixed wrapper, this keeps the full embodiment weight bank and - uses runtime embodiment_id values to gather W/b for each batch item. The - math stays equivalent to GR00T category-specific linear: - x [B,T,in] @ W[embodiment] [B,in,out] + b[embodiment]. - """ - - def __init__(self, layer: nn.Module): - super().__init__() - - # Keep the full embodiment weight bank. - # W: [num_embodiments, input_dim, output_dim] - # b: [num_embodiments, output_dim] - self.W = layer.W - self.b = layer.b - - def forward(self, x, cat_ids): - # x: [B, T, input_dim] - # cat_ids: [B] - - cat_ids = cat_ids.to(dtype=torch.long) - - # selected_w: [B, input_dim, output_dim] - # selected_b: [B, output_dim] - selected_w = torch.index_select(self.W, dim=0, index=cat_ids).to(dtype=x.dtype) - selected_b = torch.index_select(self.b, dim=0, index=cat_ids).to(dtype=x.dtype) - - # out: [B, T, output_dim] - out = torch.bmm(x, selected_w) - - # bias: [B, 1, output_dim], broadcast over T - return out + selected_b.unsqueeze(1) - - -class TRTDynamicCategorySpecificMLPPatch(nn.Module): - """Dynamic two-layer category-specific MLP used by GR00T. - - GR00T state encoders and action decoders are CategorySpecificMLP modules: - each contains layer1/layer2 CategorySpecificLinear layers. This wrapper - preserves runtime embodiment selection for both layers. - """ - - def __init__(self, mlp: nn.Module): - super().__init__() - self.layer1 = TRTDynamicCategorySpecificLinearPatch(mlp.layer1) - self.layer2 = TRTDynamicCategorySpecificLinearPatch(mlp.layer2) - - def forward(self, x, embodiment_id): - hidden = F.relu(self.layer1(x, embodiment_id)) - return self.layer2(hidden, embodiment_id) - - -class TRTGrootActionEncoderPatch(nn.Module): - def __init__(self, action_encoder: nn.Module, embodiment_id: torch.Tensor): - super().__init__() - self.W1 = TRTFixedCategorySpecificLinearPatch(action_encoder.W1, embodiment_id) - self.W2 = TRTFixedCategorySpecificLinearPatch(action_encoder.W2, embodiment_id) - self.W3 = TRTFixedCategorySpecificLinearPatch(action_encoder.W3, embodiment_id) - self.pos_encoding = action_encoder.pos_encoding - - def forward(self, actions, timesteps, embodiment_id): - batch_size, action_horizon, _ = actions.shape - - if timesteps.dim() == 1 and timesteps.shape[0] == batch_size: - timesteps = timesteps.unsqueeze(1).expand(-1, action_horizon) - else: - raise ValueError("Expected `timesteps` to have shape (B,).") - - action_emb = self.W1(actions) - timestep_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype) - hidden = torch.cat([action_emb, timestep_emb], dim=-1) - hidden = F.silu(self.W2(hidden)) - return self.W3(hidden) - - -class TRTDynamicGrootActionEncoderPatch(nn.Module): - """Dynamic GR00T noisy-action encoder. +"""GR00T setattr replacements. Installed for the whole export via ``GrootSpec.apply_patches``.""" - The original action encoder uses three embodiment-specific linear layers - around the action embedding, timestep positional embedding, and SiLU block. - This wrapper keeps embodiment_id dynamic while spelling the category-specific - pieces as index_select + bmm so Torch-TRT can lower them reliably. - """ - - def __init__(self, action_encoder: nn.Module): - super().__init__() - self.W1 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W1) - self.W2 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W2) - self.W3 = TRTDynamicCategorySpecificLinearPatch(action_encoder.W3) - self.pos_encoding = action_encoder.pos_encoding - - def forward(self, actions, timesteps, embodiment_id): - batch_size, action_horizon, _ = actions.shape - - timesteps = timesteps.unsqueeze(1).expand(-1, action_horizon) - - action_emb = self.W1(actions, embodiment_id) - timestep_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype) - - hidden = torch.cat([action_emb, timestep_emb], dim=-1) - hidden = F.silu(self.W2(hidden, embodiment_id)) - return self.W3(hidden, embodiment_id) +from __future__ import annotations +from typing import Any, Callable -class GrootDiTStepEncoderPatch(ActionStepEncoderPatch): - def __init__(self, action_head, embodiment_id: torch.Tensor | None = None): - super().__init__() - if embodiment_id is None: - self.state_encoder = action_head.state_encoder - self.action_encoder = action_head.action_encoder - else: - # Keep embodiment_id as a runtime input while replacing GR00T's - # category-specific modules with Torch-TRT-friendly dynamic wrappers. - self.state_encoder = TRTDynamicCategorySpecificMLPPatch( - action_head.state_encoder - ) - self.action_encoder = TRTDynamicGrootActionEncoderPatch( - action_head.action_encoder - ) - self.future_tokens = action_head.future_tokens - self.position_embedding = getattr(action_head, "position_embedding", None) - self.add_pos_embed = action_head.config.add_pos_embed +import torch +from torch_tensorrt.hf.exporters.models.common.patches import causal_lm_plugin_forward +from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + _patch_language_attention, + _patch_vision_attention, + register_patch, +) + +GROOT = "groot" + +register_patch( + GROOT, + "transformers.models.internvl.modeling_internvl.InternVLVisionAttention.forward", + "transformers.models.siglip.modeling_siglip.SiglipAttention.forward", + "transformers.models.siglip.modeling_siglip.SiglipSdpaAttention.forward", + "transformers.models.siglip.modeling_siglip.SiglipFlashAttention2.forward", +)(_patch_vision_attention) + +register_patch( + GROOT, + "transformers.models.llama.modeling_llama.LlamaAttention.forward", + "transformers.models.qwen2.modeling_qwen2.Qwen2Attention.forward", + "transformers.models.qwen3.modeling_qwen3.Qwen3Attention.forward", +)(_patch_language_attention) + + +@register_patch( + GROOT, + "lerobot.policies.groot.eagle2_hg_model.modeling_eagle2_5_vl.Eagle25VLForConditionalGeneration.forward", +) +def _patch_eagle_image_features(original: Callable) -> Callable: + """Vision-only compile calls ``eagle(pixel_values)``; otherwise the full VLM forward.""" + + def forward(self, pixel_values, input_ids=None, **kwargs: Any): + if input_ids is None: + return self.extract_feature(pixel_values) + return original(self, pixel_values, input_ids, **kwargs) + + return forward + + +@register_patch( + GROOT, + "transformers.models.llama.modeling_llama.LlamaForCausalLM.forward", + "transformers.models.qwen2.modeling_qwen2.Qwen2ForCausalLM.forward", + "transformers.models.qwen3.modeling_qwen3.Qwen3ForCausalLM.forward", + "transformers.models.llama.modeling_llama.LlamaModel.forward", + "transformers.models.qwen2.modeling_qwen2.Qwen2Model.forward", + "transformers.models.qwen3.modeling_qwen3.Qwen3Model.forward", +) +def _patch_groot_language_model(original: Callable) -> Callable: + """Edge prefill when rope is present; otherwise HF causal-LM forward.""" + + def forward( + self, + inputs_embeds=None, + rope_rotary_cos_sin=None, + context_lengths=None, + kvcache_start_index=None, + last_token_ids=None, + ds_stack=None, + *past_key_values, + **kwargs: Any, + ): + if rope_rotary_cos_sin is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + decoder = self if hasattr(self, "layers") else self.model + return causal_lm_plugin_forward( + decoder, + inputs_embeds, + rope_rotary_cos_sin, + context_lengths, + kvcache_start_index, + last_token_ids, + ds_stack, + *past_key_values, + lm_head=getattr(self, "lm_head", None), + ) - def forward(self, actions, timestep, vl_embs, state, embodiment_id): + return forward + + +@register_patch( + GROOT, + "lerobot.policies.groot.groot_n1.GR00TN15.forward", +) +def _patch_groot_context_projection(original: Callable) -> Callable: + """One linear + VLLN + VL attention when a single hidden tensor is present.""" + + def forward(self, hidden_states, *args, **kwargs: Any): + if args: + return original(self, hidden_states, *args, **kwargs) + context_embs = self.backbone.eagle_linear(hidden_states) + vlln = self.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) + return self.action_head.vl_self_attention(context_embs) + + return forward + + +@register_patch( + GROOT, + "lerobot.policies.groot.action_head.flow_matching_action_head.FlowmatchingActionHead.forward", +) +def _patch_groot_action_step_forward(original: Callable) -> Callable: + """One DiT velocity step when Edge action I/O is present; otherwise training.""" + + def forward( + self, + actions, + timestep=None, + context_embs=None, + state=None, + embodiment_id=None, + *args, + **kwargs: Any, + ): + if context_embs is None: + return original(self, actions, timestep, *args, **kwargs) state_features = self.state_encoder(state, embodiment_id) action_features = self.action_encoder(actions, timestep, embodiment_id) - - if self.add_pos_embed: + if self.config.add_pos_embed: pos_ids = torch.arange( action_features.shape[1], dtype=torch.long, @@ -182,47 +137,41 @@ def forward(self, actions, timestep, vl_embs, state, embodiment_id): action_features = action_features + self.position_embedding( pos_ids ).unsqueeze(0) - future_tokens = self.future_tokens.weight.unsqueeze(0).expand( - vl_embs.shape[0], + context_embs.shape[0], -1, -1, ) - - sa_embs = torch.cat( - (state_features, future_tokens, action_features), - dim=1, + sa_embs = torch.cat((state_features, future_tokens, action_features), dim=1) + expert_out = self.model( + hidden_states=sa_embs, + encoder_hidden_states=context_embs, + timestep=timestep, ) + hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out + ) + if isinstance(hidden, (tuple, list)): + hidden = hidden[0] + action_hidden = hidden[:, -int(self.config.action_horizon) :] + return self.action_decoder(action_hidden, embodiment_id) - expert_args = () - expert_kwargs = { - "hidden_states": sa_embs, - "encoder_hidden_states": vl_embs, - "timestep": timestep, - } - - decoder_args = (embodiment_id,) - decoder_kwargs = {} - - return expert_args, expert_kwargs, decoder_args, decoder_kwargs - - -class ContextProjectionPatch(nn.Module): - """eagle_linear -> vlln -> vl_self_attention (matches eager context path).""" + return forward - def __init__(self, eagle_linear, vlln, vl_self_attention): - super().__init__() - self.eagle_linear = eagle_linear - self.vlln = vlln - self.vl_self_attention = vl_self_attention - def forward(self, hidden_states: torch.Tensor): - context_embs = self.eagle_linear(hidden_states) +@register_patch( + GROOT, + "lerobot.policies.groot.action_head.flow_matching_action_head.CategorySpecificLinear.forward", +) +def _patch_category_specific_linear(_original: Callable) -> Callable: + """``index_select`` + ``bmm`` is the TensorRT-friendly form of ``W[cat_ids]``.""" - vlln_weight = getattr(self.vlln, "weight", None) - if vlln_weight is not None: - context_embs = context_embs.to(dtype=vlln_weight.dtype) + def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor: + cat_ids = cat_ids.to(dtype=torch.long) + selected_w = torch.index_select(self.W, dim=0, index=cat_ids).to(dtype=x.dtype) + selected_b = torch.index_select(self.b, dim=0, index=cat_ids).to(dtype=x.dtype) + return torch.bmm(x, selected_w) + selected_b.unsqueeze(1) - context_embs = self.vlln(context_embs) - context_embs = self.vl_self_attention(context_embs) - return context_embs + return forward diff --git a/py/torch_tensorrt/hf/exporters/models/groot/spec.py b/py/torch_tensorrt/hf/exporters/models/groot/spec.py index b39529c23fe..59b8a97e139 100644 --- a/py/torch_tensorrt/hf/exporters/models/groot/spec.py +++ b/py/torch_tensorrt/hf/exporters/models/groot/spec.py @@ -10,21 +10,11 @@ kv_kwargs, split_flat_to_kwargs, ) -from torch_tensorrt.hf.exporters.models.common.patches import ( - CausalLMPatch, - GridVisionPatch, - StaticActionVelocityStepPatch, - language_decoder, -) from torch_tensorrt.hf.exporters.models.groot.helpers import ( _groot, make_embodiment_id, ) -from torch_tensorrt.hf.exporters.models.groot.patches import ( - ContextProjectionPatch, - GrootDiTStepEncoderPatch, - TRTDynamicCategorySpecificMLPPatch, -) +from torch_tensorrt.hf.exporters.models.groot.patches import GROOT from torch_tensorrt.hf.exporters.ops import call_engine, scatter_image_tokens from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, @@ -33,10 +23,32 @@ ) +def _export_module(module: nn.Module, sample: Mapping[str, Any]) -> nn.Module: + device = sample["pixel_values"].device + dtype = sample["pixel_values"].dtype + return module.eval().to(device=device, dtype=dtype) + + +def _causal_lm(language: nn.Module) -> nn.Module: + get_base = getattr(language, "get_base_model", None) + if callable(get_base): + try: + return get_base() + except Exception: + pass + return language + + @register_edge_spec("groot", "gr00t") class GrootSpec(EdgeSpec): # type: ignore[misc] components = ("vision", "language", "context_projection", "action") + def apply_patches(self, model=None): + del model + from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + + return apply_patches(GROOT) + def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any ) -> MutableMapping[str, Any]: @@ -107,53 +119,6 @@ def prepare_sample_inputs( "embodiment_id": make_embodiment_id(policy, state, device, torch.long), } - def wrap( - self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any - ) -> nn.Module: - found = _groot(model) - eagle = found.backbone.eagle_model - device = sample["pixel_values"].device - dtype = sample["pixel_values"].dtype - # GR00T loads DiT / category MLPs in bf16. Export tensors are fp16. - # The e2e path casts the whole wrapper; without that, sa_embs is - # float32 against bf16 attn.to_q weights. - if name == "vision": - module = GridVisionPatch( - vision_model=eagle.vision_model, - projector=eagle.mlp1, - sample_pixel_values=sample["pixel_values"], - select_layer=eagle.select_layer, - pixel_shuffle=eagle.use_pixel_shuffle, - downsample_ratio=getattr(eagle, "downsample_ratio", 0.5), - vision_kwargs={}, - ) - elif name == "language": - language = eagle.language_model - module = CausalLMPatch( - language_decoder(language), language.lm_head, select_layer=-1 - ) - elif name == "context_projection": - module = ContextProjectionPatch( # type: ignore[no-untyped-call] - found.backbone.eagle_linear, - found.action_head.vlln, - found.action_head.vl_self_attention, - ) - elif name == "action": - module = StaticActionVelocityStepPatch( - step_encoder=GrootDiTStepEncoderPatch( - found.action_head, sample.get("embodiment_id") - ), - action_expert=found.action_head.model, - velocity_decoder=TRTDynamicCategorySpecificMLPPatch( - found.action_head.action_decoder - ), - output_tokens=int(found.action_head.config.action_horizon), - cast_hidden_fp32=False, - ) - else: - raise KeyError(name) - return module.eval().to(device=device, dtype=dtype) - def prepare( self, name: str, @@ -161,15 +126,10 @@ def prepare( sample: MutableMapping[str, Any], upstream: Mapping[str, Any], config: Any, - module: nn.Module, ) -> ComponentBundle: from torch_tensorrt.hf.exporters.plugin.attention import ( ContextAttentionMaskType, ) - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( - patch_language_attention, - patch_vision_attention, - ) found = _groot(model) eagle = found.backbone.eagle_model @@ -178,27 +138,18 @@ def prepare( if name == "vision": px = sample["pixel_values"] - seq_len = int(getattr(module, "seq_len", 1) or 1) - batch = int(getattr(module, "batch_size", 1) or 1) - - def _patch(mod: nn.Module) -> Any: - vision = getattr(mod.vision_model, "vision_model", mod.vision_model) - return patch_vision_attention( - vision, batch_size=batch, seq_len=seq_len, name="SigLIP" - ) - return ComponentBundle( + module=_export_module(eagle, sample), trace_args=(px,), save_args=(px,), input_names=["pixel_values"], output_names=["visual_embeds"], - patch_fn=_patch, model_type="vit", engine_file="visual.engine", ) if name == "language": - language = eagle.language_model + language = _causal_lm(eagle.language_model) input_ids = sample["input_ids"] input_embs = language.get_input_embeddings()(input_ids) image_token_index = getattr( @@ -225,28 +176,14 @@ def _patch(mod: nn.Module) -> Any: dtype=dtype, ) sample.update(split_flat_to_kwargs(packed, meta["input_names"])) - cfg = language.config - head_dim = int( - getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) - ) - - def _patch(mod: nn.Module) -> Any: - decoder = getattr(mod, "lm", mod) - return patch_language_attention( - decoder, - hidden_size=int(cfg.hidden_size), - num_attention_heads=int(cfg.num_attention_heads), - num_key_value_heads=int(cfg.num_key_value_heads), - head_dim=head_dim, - context_attention_mask_type=ContextAttentionMaskType.CAUSAL, - ) 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"], - patch_fn=_patch, + context_attention_mask_type=int(ContextAttentionMaskType.CAUSAL), model_type="language", engine_file="language.engine", ) @@ -254,6 +191,7 @@ def _patch(mod: nn.Module) -> Any: 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"], @@ -284,6 +222,7 @@ def _patch(mod: nn.Module) -> Any: sample["embodiment_id"], ) return ComponentBundle( + module=_export_module(found.action_head, sample), trace_args=args, save_args=args, input_names=[ diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py b/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py index c0850837685..e9856dc285e 100644 --- a/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py @@ -1,41 +1,49 @@ +"""Nemotron setattr replacements. Installed for the whole export via ``NemotronSpec.apply_patches``.""" + from __future__ import annotations -import torch -import torch.nn as nn +from contextlib import contextmanager +from typing import Any, Callable, Iterator + from torch_tensorrt.hf.exporters.models.common.patches import gather_last_token_hidden from torch_tensorrt.hf.exporters.models.nemotron.helpers import _decoder, _kind +from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + apply_patches, + register_patch, +) +NEMOTRON = "nemotron" -class NemotronPatch(nn.Module): # type: ignore[misc] - """Hybrid decoder: plugin attention / mamba / moe, native MLP.""" - def __init__(self, model: nn.Module): - super().__init__() - decoder = _decoder(model) - self.layers = decoder.layers - self.norm = decoder.norm_f - self.lm_head = model.lm_head - self.kinds = [_kind(block.mixer) for block in self.layers] - self.num_attn = self.kinds.count("attention") - self.num_mamba = self.kinds.count("mamba") +@register_patch( + NEMOTRON, + "transformers.models.nemotron_h.modeling_nemotron_h.NemotronHForCausalLM.forward", +) +def _patch_nemotron_causal_lm(original: Callable) -> Callable: + """Hybrid plugin prefill when rope is present; otherwise HF causal-LM forward.""" def forward( self, - inputs_embeds: torch.Tensor, - rope_rotary_cos_sin: torch.Tensor, - context_lengths: torch.Tensor, - kvcache_start_index: torch.Tensor, - last_token_ids: torch.Tensor, - *states: torch.Tensor, - ) -> tuple[torch.Tensor, ...]: - na, nm = self.num_attn, self.num_mamba + inputs_embeds=None, + rope_rotary_cos_sin=None, + context_lengths=None, + kvcache_start_index=None, + last_token_ids=None, + *states, + **kwargs: Any, + ): + if rope_rotary_cos_sin is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + decoder = _decoder(self) + kinds = [_kind(block.mixer) for block in decoder.layers] + na, nm = kinds.count("attention"), kinds.count("mamba") kvs = list(states[:na]) convs = list(states[na : na + nm]) ssms = list(states[na + nm : na + 2 * nm]) kv_i = conv_i = 0 hidden = inputs_embeds present_kv, present_conv, present_ssm = [], [], [] - for block, kind in zip(self.layers, self.kinds): + for block, kind in zip(decoder.layers, kinds): residual = hidden hidden = block.norm(hidden) mixer = block.mixer @@ -59,7 +67,32 @@ def forward( else: hidden = mixer(hidden) hidden = residual + hidden - hidden = self.norm(hidden) + hidden = decoder.norm_f(hidden) last = gather_last_token_hidden(hidden, last_token_ids) logits = self.lm_head(last).float() return (logits, *present_kv, *present_conv, *present_ssm) + + return forward + + +@contextmanager +def apply_nemotron_patches(model: Any | None = None) -> Iterator[None]: + """Class setattr plus mixer plugin wrappers (MoE packing needs the instance).""" + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + patch_nemotron_mixers, + restore_attention, + ) + + restore = [] + try: + if model is not None: + restore = patch_nemotron_mixers(model, model.config) + for block in _decoder(model).layers: + prepare = getattr(block.mixer, "prepare_for_export", None) + if callable(prepare): + prepare() + with apply_patches(NEMOTRON): + yield + finally: + if restore: + restore_attention(restore) diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py b/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py index 36010b65780..413a5b21e85 100644 --- a/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py +++ b/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py @@ -11,9 +11,12 @@ ) from torch_tensorrt.hf.exporters.models.nemotron.helpers import ( _decoder, + _kind, allocate_plugin_states, ) -from torch_tensorrt.hf.exporters.models.nemotron.patches import NemotronPatch +from torch_tensorrt.hf.exporters.models.nemotron.patches import ( + apply_nemotron_patches, +) from torch_tensorrt.hf.exporters.ops import call_engine from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, @@ -26,6 +29,9 @@ class NemotronSpec(EdgeSpec): # type: ignore[misc] components = ("language",) + def apply_patches(self, model=None): + return apply_nemotron_patches(model) + def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any ) -> MutableMapping[str, Any]: @@ -58,20 +64,6 @@ def prepare_sample_inputs( "bsz": embeddings.shape[0], } - def wrap( - self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any - ) -> nn.Module: - from torch_tensorrt.hf.exporters.plugin.moe import PluginNemotronMoE - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( - patch_nemotron_mixers, - ) - - patch_nemotron_mixers(model, model.config) # type: ignore[no-untyped-call] - for block in _decoder(model).layers: - if isinstance(block.mixer, PluginNemotronMoE): - block.mixer.prepare_for_export() - return NemotronPatch(model).eval() - def prepare( self, name: str, @@ -79,10 +71,10 @@ def prepare( sample: MutableMapping[str, Any], upstream: Mapping[str, Any], config: Any, - module: nn.Module, ) -> ComponentBundle: from torch_tensorrt.hf.exporters.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 @@ -101,7 +93,8 @@ def prepare( model, model.config, bsz, int(config.max_seq_len), device, dtype ) flat = (embeds, rope, ctx_len, kv_start, last_token_ids, *kvs, *convs, *ssms) - na, nm = module.num_attn, module.num_mamba + kinds = [_kind(block.mixer) for block in _decoder(model).layers] + na, nm = kinds.count("attention"), kinds.count("mamba") names = [ "inputs_embeds", "rope_rotary_cos_sin", @@ -114,6 +107,7 @@ def prepare( ] sample.update(split_flat_to_kwargs(flat, names)) return ComponentBundle( + module=model.eval(), trace_args=flat, save_args=flat, input_names=names, diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/patches.py b/py/torch_tensorrt/hf/exporters/models/pi05/patches.py index 55e59467fc2..1ef966509a2 100644 --- a/py/torch_tensorrt/hf/exporters/models/pi05/patches.py +++ b/py/torch_tensorrt/hf/exporters/models/pi05/patches.py @@ -1,21 +1,77 @@ +"""PI05 setattr replacements. Installed for the whole export via ``Pi05Spec.apply_patches``.""" + from __future__ import annotations +from typing import Any, Callable + +import torch import torch.nn.functional as F -from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding -from torch_tensorrt.hf.exporters.models.common.patches import ActionStepEncoderPatch -from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache +from torch_tensorrt.hf.exporters.models.common.patches import causal_lm_plugin_forward +from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + _patch_language_attention, + _patch_vision_attention, + register_patch, +) + +PI05 = "pi05" +register_patch( + PI05, + "transformers.models.siglip.modeling_siglip.SiglipAttention.forward", + "transformers.models.siglip.modeling_siglip.SiglipSdpaAttention.forward", + "transformers.models.siglip.modeling_siglip.SiglipFlashAttention2.forward", +)(_patch_vision_attention) -class PI05PrefixKVStepEncoderPatch(ActionStepEncoderPatch): - """PI05 suffix embed + AdaRMS cond, consumed by Gemma action expert.""" +register_patch( + PI05, + "transformers.models.gemma.modeling_gemma.GemmaAttention.forward", + "transformers.models.gemma2.modeling_gemma2.Gemma2Attention.forward", +)(_patch_language_attention) - def __init__(self, core): - super().__init__() - self.action_in_proj = core.action_in_proj - self.time_mlp_in = core.time_mlp_in - self.time_mlp_out = core.time_mlp_out - self.config = core.config - self.hidden_size = core.action_in_proj.out_features + +@register_patch( + PI05, + "lerobot.policies.pi_gemma.PiGemmaModel.forward", + "transformers.models.gemma.modeling_gemma.GemmaModel.forward", + "transformers.models.gemma2.modeling_gemma2.Gemma2Model.forward", +) +def _patch_pi05_language_model(original: Callable) -> Callable: + """Edge prefill when rope is present; otherwise HF / action-expert forward.""" + + def forward( + self, + inputs_embeds=None, + rope_rotary_cos_sin=None, + context_lengths=None, + kvcache_start_index=None, + last_token_ids=None, + ds_stack=None, + *past_key_values, + **kwargs: Any, + ): + if rope_rotary_cos_sin is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + return causal_lm_plugin_forward( + self, + inputs_embeds, + rope_rotary_cos_sin, + context_lengths, + kvcache_start_index, + last_token_ids, + ds_stack, + *past_key_values, + lm_head=getattr(self, "lm_head", None), + ) + + return forward + + +@register_patch( + PI05, + "lerobot.policies.pi05.modeling_pi05.PI05Pytorch.forward", +) +def _patch_pi05_action_step_forward(original: Callable) -> Callable: + """One diffusion step when prefix KV is present; otherwise training forward.""" def forward( self, @@ -25,28 +81,73 @@ def forward( prefix_v, position_ids, attention_mask, + *args, + **kwargs: Any, ): - suffix_embs = self.action_in_proj(x_t) + if getattr(prefix_k, "ndim", 0) != 5: + return original( + self, + x_t, + timestep, + prefix_k, + prefix_v, + position_ids, + attention_mask, + *args, + **kwargs, + ) + from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding + from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache + suffix_embs = self.action_in_proj(x_t) time_emb = create_sinusoidal_pos_embedding( timestep, - self.hidden_size, + self.action_in_proj.out_features, min_period=self.config.min_period, max_period=self.config.max_period, device=timestep.device, ).to(dtype=suffix_embs.dtype) + adarms_cond = F.silu(self.time_mlp_out(F.silu(self.time_mlp_in(time_emb)))) + expert_out = self.paligemma_with_expert.gemma_expert.model( + inputs_embeds=suffix_embs, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=PrefixKVCache(prefix_k, prefix_v), + use_cache=False, + adarms_cond=adarms_cond, + ) + hidden = ( + expert_out.last_hidden_state + if hasattr(expert_out, "last_hidden_state") + else expert_out + ) + if isinstance(hidden, (tuple, list)): + hidden = hidden[0] + return self.action_out_proj(hidden[:, -int(self.config.chunk_size) :]) + + return forward + + +@register_patch( + PI05, + "transformers.models.paligemma.modeling_paligemma.PaliGemmaModel.forward", +) +def _patch_paligemma_image_features(_original: Callable) -> Callable: + """Match LeRobot ``embed_image`` / HF ``get_image_features``, return a tensor.""" + + def forward(self, pixel_values: torch.Tensor, **kwargs: Any) -> torch.Tensor: + out_dtype = pixel_values.dtype + if pixel_values.dtype != torch.float32: + pixel_values = pixel_values.to(torch.float32) + image_outputs = self.vision_tower(pixel_values, **kwargs) + hidden = ( + image_outputs.last_hidden_state + if hasattr(image_outputs, "last_hidden_state") + else image_outputs[0] + ) + features = self.multi_modal_projector(hidden) + if features.dtype != out_dtype: + features = features.to(out_dtype) + return features - adarms_cond = self.time_mlp_in(time_emb) - adarms_cond = F.silu(adarms_cond) - adarms_cond = self.time_mlp_out(adarms_cond) - adarms_cond = F.silu(adarms_cond) - - expert_kwargs = { - "inputs_embeds": suffix_embs, - "attention_mask": attention_mask, - "position_ids": position_ids, - "past_key_values": PrefixKVCache(prefix_k, prefix_v), - "use_cache": False, - "adarms_cond": adarms_cond, - } - return (), expert_kwargs, (), {} + return forward diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/spec.py b/py/torch_tensorrt/hf/exporters/models/pi05/spec.py index 63feb4670f8..7fb9381f449 100644 --- a/py/torch_tensorrt/hf/exporters/models/pi05/spec.py +++ b/py/torch_tensorrt/hf/exporters/models/pi05/spec.py @@ -10,20 +10,14 @@ kv_kwargs, split_flat_to_kwargs, ) -from torch_tensorrt.hf.exporters.models.common.patches import ( - CausalLMPatch, - GridVisionPatch, - StaticActionVelocityStepPatch, - language_decoder, -) +from torch_tensorrt.hf.exporters.models.common.patches import language_decoder from torch_tensorrt.hf.exporters.models.pi05.helpers import ( _core, - _nchw_to_hwc, build_pi05_prefix_embs, make_pi05_suffix_position_and_mask, pi05_compact_index, ) -from torch_tensorrt.hf.exporters.models.pi05.patches import PI05PrefixKVStepEncoderPatch +from torch_tensorrt.hf.exporters.models.pi05.patches import PI05 from torch_tensorrt.hf.exporters.ops import call_engine, fuse_prefix from torch_tensorrt.hf.exporters.spec import ( ComponentBundle, @@ -36,6 +30,13 @@ 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 + from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + + return apply_patches(PI05) + def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any ) -> MutableMapping[str, Any]: @@ -87,37 +88,6 @@ def prepare_sample_inputs( "lang_embeds": lang_embeds.to(device=device, dtype=dtype).contiguous(), } - def wrap( - self, name: str, model: nn.Module, sample: Mapping[str, Any], config: Any - ) -> nn.Module: - core = _core(model) - paligemma = core.paligemma_with_expert.paligemma.model - if name == "vision": - px = sample["pixel_values"] - sample_px = px.float() if px.dtype != torch.float32 else px - return GridVisionPatch( - vision_model=paligemma.vision_tower.float(), - projector=paligemma.multi_modal_projector, - sample_pixel_values=sample_px, - select_layer=-1, - pixel_shuffle=False, - downsample_ratio=0.5, - force_float32_input=True, - ).eval() - if name == "language": - language = paligemma.language_model - lm_head = core.paligemma_with_expert.paligemma.lm_head - return CausalLMPatch(language_decoder(language), lm_head).eval() - if name == "action": - return StaticActionVelocityStepPatch( - step_encoder=PI05PrefixKVStepEncoderPatch(core), # type: ignore[no-untyped-call] - action_expert=core.paligemma_with_expert.gemma_expert.model, - velocity_decoder=core.action_out_proj, - output_tokens=int(core.config.chunk_size), - cast_hidden_fp32=False, - ).eval() - raise KeyError(name) - def prepare( self, name: str, @@ -125,36 +95,24 @@ def prepare( sample: MutableMapping[str, Any], upstream: Mapping[str, Any], config: Any, - module: nn.Module, ) -> ComponentBundle: from torch_tensorrt.hf.exporters.plugin.attention import ( ContextAttentionMaskType, ) - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( - patch_language_attention, - patch_vision_attention, - ) core = _core(model) + paligemma = core.paligemma_with_expert.paligemma.model device = sample["pixel_values"].device dtype = sample["pixel_values"].dtype if name == "vision": px = sample["pixel_values"] - seq_len = int(getattr(module, "seq_len", 1) or 1) - batch = int(getattr(module, "batch_size", 1) or 1) - - def _patch(mod: nn.Module) -> Any: - return patch_vision_attention( - mod.vision_model, batch_size=batch, seq_len=seq_len, name="SigLIP" - ) - return ComponentBundle( + module=paligemma.eval(), trace_args=(px,), - save_args=(_nchw_to_hwc(px),), + save_args=(px,), input_names=["pixel_values"], output_names=["visual_embeds"], - patch_fn=_patch, model_type="vit", engine_file="visual.engine", ) @@ -196,28 +154,14 @@ def _patch(mod: nn.Module) -> Any: seq_len=compact_len, ) sample.update(split_flat_to_kwargs(flat, meta["input_names"])) - cfg = language.config - head_dim = int( - getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads) - ) - - def _patch(mod: nn.Module) -> Any: - decoder = getattr(mod, "lm", mod) - return patch_language_attention( - decoder, - hidden_size=int(cfg.hidden_size), - num_attention_heads=int(cfg.num_attention_heads), - num_key_value_heads=int(cfg.num_key_value_heads), - head_dim=head_dim, - context_attention_mask_type=ContextAttentionMaskType.PADDING, - ) 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"], - patch_fn=_patch, + context_attention_mask_type=int(ContextAttentionMaskType.PADDING), extra_config={"prefix_pad_mask_len": compact_len}, model_type="language", engine_file="language.engine", @@ -250,6 +194,7 @@ def _patch(mod: nn.Module) -> Any: 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=[ @@ -275,6 +220,9 @@ def capture_upstream( ) -> 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]} diff --git a/py/torch_tensorrt/hf/exporters/ops.py b/py/torch_tensorrt/hf/exporters/ops.py index 7623c8b43ee..9755d451aff 100644 --- a/py/torch_tensorrt/hf/exporters/ops.py +++ b/py/torch_tensorrt/hf/exporters/ops.py @@ -83,6 +83,9 @@ def fuse_prefix( hidden = lang_embeds.shape[-1] batch = lang_embeds.shape[0] vis = vision_tokens + # Vision engines may emit [B, S, H] (HF image features) or flattened [N, H]. + if vis.ndim == 3: + vis = vis.reshape(-1, vis.shape[-1]) if vis.ndim == 2: vis = vis.reshape(batch, -1, hidden) embs = torch.cat([vis, lang_embeds], dim=1) diff --git a/py/torch_tensorrt/hf/exporters/plugin/attn_patches.py b/py/torch_tensorrt/hf/exporters/plugin/attn_patches.py new file mode 100644 index 00000000000..cea2ea5b6c5 --- /dev/null +++ b/py/torch_tensorrt/hf/exporters/plugin/attn_patches.py @@ -0,0 +1,287 @@ +"""HF-style class ``setattr`` patches that swap attention ``forward`` for Edge plugins. + +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). + +Language dispatch: Edge prefill calls ``self_attn(..., rope_rotary_cos_sin=...)``. +The PI05 action expert is often the same class (GemmaAttention / PiGemmaModel) +but uses HF ``past_key_values``. If ``rope_rotary_cos_sin`` is absent, the +original forward runs. +""" + +from __future__ import annotations + +import contextlib +import importlib +import inspect +from typing import Any, Callable + +import torch +import torch.nn as nn +from torch_tensorrt.hf.exporters.plugin.attention import ContextAttentionMaskType + +_PATCHES: dict[str, list[tuple[str, Callable]]] = {} +_LANGUAGE_MASK_TYPE = int(ContextAttentionMaskType.PADDING) + +VISION_BACKEND = "edge_vision" +LANGUAGE_BACKEND = "edge_language" + + +def set_language_mask_type(mask_type: int) -> None: + """Prefill mask for ``trt::attention_plugin`` (PADDING vs CAUSAL).""" + global _LANGUAGE_MASK_TYPE + _LANGUAGE_MASK_TYPE = int(mask_type) + + +def language_mask_type() -> int: + return _LANGUAGE_MASK_TYPE + + +@contextlib.contextmanager +def patch_attribute(obj: Any, attribute: str, factory: Callable): + original = getattr(obj, attribute) + setattr(obj, attribute, factory(original)) + try: + yield + finally: + setattr(obj, attribute, original) + + +@contextlib.contextmanager +def apply_patches(backend: str | None): + """Resolve dotted paths and install those class attributes for the block. + + Paths are resolved here so the policy's modeling modules + are already loaded and we do not import every HF family up front. + """ + if not backend: + yield + return + with contextlib.ExitStack() as stack: + for path, factory in _PATCHES.get(backend, []): + obj_path, _, attribute = path.rpartition(".") + obj = _resolve_dotted_path(obj_path) + if obj is None: + continue + stack.enter_context(patch_attribute(obj, attribute, factory)) + yield + + +def register_patch(backend: str, *paths: str): + """Record ``factory(original)`` for each dotted ``Class.attribute`` path.""" + + def decorator(fn: Callable) -> Callable: + for path in paths: + _PATCHES.setdefault(backend, []).append((path, fn)) + return fn + + return decorator + + +def _resolve_dotted_path(path: str) -> Any | None: + parts = path.split(".") + try: + obj: Any = importlib.import_module(parts[0]) + for part in parts[1:]: + try: + obj = importlib.import_module(f"{obj.__name__}.{part}") + except (ImportError, AttributeError): + obj = getattr(obj, part) + return obj + except Exception: + # Missing family, or a modeling import that fails for unrelated reasons + # (e.g. torchaudio CUDA mismatch). Skip that class path. + return None + + +def _returns_tuple(original: Callable) -> bool: + try: + src = inspect.getsource(original) + except (OSError, TypeError): + return True + return ( + "return attn_output, attn_weights" in src + or "return attn_output, attn_weight" in src + or "return attn_output, None" in src + or "return output, attn_weights" in src + ) + + +def _out_proj(module: nn.Module) -> nn.Module: + proj = getattr(module, "out_proj", None) or getattr( + module, "projection_layer", None + ) + if proj is None: + raise AttributeError( + f"{type(module).__name__} has no out_proj/projection_layer" + ) + return proj + + +def _language_dims(module: nn.Module) -> tuple[int, int, int]: + cfg = getattr(module, "config", None) + num_heads = int( + getattr(module, "num_heads", None) + or getattr(module, "num_attention_heads", None) + or cfg.num_attention_heads + ) + num_kv = int( + getattr(module, "num_key_value_heads", None) + or getattr(cfg, "num_key_value_heads", num_heads) + ) + head_dim = int( + getattr(module, "head_dim", None) + or (cfg.hidden_size // cfg.num_attention_heads) + ) + return num_heads, num_kv, head_dim + + +@register_patch( + VISION_BACKEND, + "transformers.models.internvl.modeling_internvl.InternVLVisionAttention.forward", +) +def _patch_vision_attention(original: Callable) -> Callable: + returns_tuple = _returns_tuple(original) + + def forward(self, hidden_states, attention_mask=None, **kwargs): + del kwargs + if attention_mask is not None: + raise RuntimeError( + f"{type(self).__name__} Edge vision plugin expects no attention_mask" + ) + + batch_size, seq_len, _ = hidden_states.shape + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q_norm = getattr(self, "q_norm", None) + k_norm = getattr(self, "k_norm", None) + if q_norm is not None: + q = q_norm(q) + if k_norm is not None: + k = k_norm(k) + + num_heads = int(self.num_heads) + head_dim = int(self.head_dim) + q = ( + q.reshape(batch_size * seq_len, num_heads, head_dim) + .to(torch.float16) + .contiguous() + ) + k = ( + k.reshape(batch_size * seq_len, num_heads, head_dim) + .to(torch.float16) + .contiguous() + ) + v = ( + v.reshape(batch_size * seq_len, num_heads, head_dim) + .to(torch.float16) + .contiguous() + ) + + cu_seqlens = torch.arange( + 0, + (batch_size + 1) * seq_len, + seq_len, + device=q.device, + dtype=torch.int32, + ) + max_seqlen_carrier = torch.zeros(seq_len, device=q.device, dtype=torch.int32) + + attn_output = torch.ops.trt.vit_attention_plugin.default( + q, + k, + v, + cu_seqlens, + max_seqlen_carrier, + num_heads, + head_dim, + ) + attn_output = attn_output.reshape(batch_size, seq_len, num_heads * head_dim) + out_proj = _out_proj(self) + attn_output = attn_output.to(dtype=out_proj.weight.dtype) + attn_output = out_proj(attn_output) + drop = getattr(self, "projection_dropout", None) + if drop is not None: + attn_output = drop(attn_output) + if returns_tuple: + return attn_output, None + return attn_output + + return forward + + +@register_patch( + LANGUAGE_BACKEND, + "transformers.models.gemma.modeling_gemma.GemmaAttention.forward", + "transformers.models.gemma2.modeling_gemma2.Gemma2Attention.forward", + "transformers.models.llama.modeling_llama.LlamaAttention.forward", + "transformers.models.qwen2.modeling_qwen2.Qwen2Attention.forward", + "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") + if rope_rotary_cos_sin is None: + return original(self, hidden_states, *args, **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: + raise ValueError("past_key_value (KV cache tensor) must be provided") + if kvcache_start_index is None: + raise ValueError("kvcache_start_index must be provided") + + batch_size, seq_len, _ = hidden_states.shape + num_heads, num_kv, head_dim = _language_dims(self) + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q_norm = getattr(self, "q_norm", None) + k_norm = getattr(self, "k_norm", None) + if q_norm is not None: + q = q_norm(q.view(batch_size, seq_len, num_heads, head_dim)).view( + batch_size, seq_len, -1 + ) + if k_norm is not None: + k = k_norm(k.view(batch_size, seq_len, num_kv, head_dim)).view( + batch_size, seq_len, -1 + ) + + dtype = q.dtype + attn_out, updated_kv = torch.ops.trt.attention_plugin.default( + q.to(torch.float16), + k.to(torch.float16), + v.to(torch.float16), + past_key_value, + ctx_len, + rope_rotary_cos_sin, + kvcache_start_index, + num_heads, + num_kv, + False, + head_dim, + False, + -1, + language_mask_type(), + ) + attn_hidden = num_heads * head_dim + attn_out = attn_out.reshape(batch_size, seq_len, attn_hidden).to(dtype) + o_proj = getattr(self, "o_proj", None) or getattr(self, "out_proj", None) + if o_proj is None: + raise AttributeError(f"{type(self).__name__} has no o_proj/out_proj") + return o_proj(attn_out), updated_kv + + return forward + + +def registered_backends() -> dict[str, int]: + """Test helper: how many class paths are registered per backend.""" + return {name: len(items) for name, items in _PATCHES.items()} diff --git a/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py b/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py index 522003330c6..31a3c58a82c 100644 --- a/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py +++ b/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py @@ -347,6 +347,8 @@ def load_plugins_for_trt(): _register_vit_attention_plugin_op() register_mamba_plugin_ops() register_moe_plugin_ops() + from . import attn_patches as _attn_patches # noqa: F401,E402 + load_plugin() from . import plugin_converter as _plugin_converter # noqa: F401,E402 diff --git a/py/torch_tensorrt/hf/exporters/spec.py b/py/torch_tensorrt/hf/exporters/spec.py index 98f94a9e406..0186cdf2600 100644 --- a/py/torch_tensorrt/hf/exporters/spec.py +++ b/py/torch_tensorrt/hf/exporters/spec.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Mapping, MutableMapping +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from typing import Any @@ -29,21 +30,22 @@ def registered_specs() -> dict[str, type[EdgeSpec]]: class ComponentBundle: """Everything :func:`compile_component` needs to build one engine.""" + module: nn.Module trace_args: tuple[Any, ...] save_args: tuple[Any, ...] input_names: list[str] output_names: list[str] - input_specs: Any = 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 model_type: str = "edge" engine_file: str = "engine.engine" class EdgeSpec(ABC): - """Per-family wrap / flatten / runtime wiring. + """Per-family flatten / runtime wiring. ``EdgeExporter.export`` never branches on PI05 vs Nemotron. It only loops ``spec.components``. @@ -51,6 +53,18 @@ class EdgeSpec(ABC): components: tuple[str, ...] = () + def apply_patches( + self, model: nn.Module | None = None + ) -> AbstractContextManager[None]: + """Install this family's setattr replacements for the whole ``export()``. + + 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() + @abstractmethod def prepare_sample_inputs( self, @@ -58,17 +72,7 @@ def prepare_sample_inputs( raw: Mapping[str, Any], config: Any, ) -> MutableMapping[str, Any]: - """Caller payload → stem dict used by wrap/prepare/run.""" - - @abstractmethod - def wrap( - self, - name: str, - model: nn.Module, - sample: Mapping[str, Any], - config: Any, - ) -> nn.Module: - """Replace a submodule with an export wrapper (not an HF forward patch).""" + """Caller payload → stem dict used by prepare/run.""" @abstractmethod def prepare( @@ -78,9 +82,8 @@ def prepare( sample: MutableMapping[str, Any], upstream: Mapping[str, Any], config: Any, - module: nn.Module, ) -> ComponentBundle: - """Build the trace/save tuple for ``module``.""" + """Select the original submodule and build its trace/save tuple.""" def capture_upstream( self, diff --git a/tests/py/dynamo/hf/test_edge_exporter.py b/tests/py/dynamo/hf/test_edge_exporter.py index bd9ae2aacbe..43bd6fec71f 100644 --- a/tests/py/dynamo/hf/test_edge_exporter.py +++ b/tests/py/dynamo/hf/test_edge_exporter.py @@ -18,12 +18,10 @@ class DummySpec(EdgeSpec): def prepare_sample_inputs(self, model, raw, config): return {"x": raw["x"]} - def wrap(self, name, model, sample, config) -> nn.Module: - return model.eval() - - def prepare(self, name, model, sample, upstream, config, module) -> ComponentBundle: + def prepare(self, name, model, sample, upstream, config) -> ComponentBundle: x = sample["x"] return ComponentBundle( + module=model.eval(), trace_args=(x,), save_args=(x,), input_names=["x"], @@ -132,10 +130,7 @@ class _PatchSpec(EdgeSpec): def prepare_sample_inputs(self, model, raw, config): return {"x": raw["x"]} - def wrap(self, name, model, sample, config) -> nn.Module: - return model.eval() - - def prepare(self, name, model, sample, upstream, config, module) -> ComponentBundle: + def prepare(self, name, model, sample, upstream, config) -> ComponentBundle: x = sample["x"] def _patch(mod): @@ -144,6 +139,7 @@ def _patch(mod): return [(mod.layer, orig)] return ComponentBundle( + module=model.eval(), trace_args=(x,), save_args=(x,), input_names=["x"], @@ -177,3 +173,300 @@ def test_edge_exporter_dryrun_keeps_attention_patch(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 torch_tensorrt.hf.exporters.plugin.attn_patches import patch_attribute + + class Owner: + def go(self): + return 1 + + def factory(original): + def go(self): + return original(self) + 1 + + return go + + with patch_attribute(Owner, "go", factory): + assert Owner().go() == 2 + assert Owner().go() == 1 + + +@pytest.mark.unit +def test_language_attn_keeps_hf_forward_without_rope(): + from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + _patch_language_attention, + ) + + class Dummy(nn.Module): + def forward(self, hidden_states, past_key_values=None, **kwargs): + del past_key_values, kwargs + return hidden_states * 2, None + + Dummy.forward = _patch_language_attention(Dummy.forward) + hidden = torch.ones(1, 2, 4) + out, extra = Dummy()(hidden, past_key_values="cache") + torch.testing.assert_close(out, hidden * 2) + assert extra is None + + +@pytest.mark.unit +def test_pi05_backend_registers_vision_and_language(): + from torch_tensorrt.hf.exporters.models.pi05.patches import PI05 + from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + + paths = [p for p, _ in _PATCHES[PI05]] + assert any("SiglipAttention.forward" in p for p in paths) + assert any("PaliGemmaModel.forward" in p for p in paths) + assert any("GemmaAttention.forward" in p for p in paths) + assert any("PiGemmaModel.forward" in p for p in paths) + assert any("PI05Pytorch.forward" in p for p in paths) + + +@pytest.mark.unit +def test_paligemma_image_features_patch_returns_tensor(): + from torch_tensorrt.hf.exporters.models.pi05.patches import ( + _patch_paligemma_image_features, + ) + + class _Out: + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state + + class Tower(nn.Module): + def forward(self, pixel_values, **kwargs): + del kwargs + return _Out(pixel_values.new_ones(pixel_values.shape[0], 4, 8)) + + class Proj(nn.Module): + def forward(self, hidden): + return hidden + + class DummyPaliGemma(nn.Module): + def __init__(self): + super().__init__() + self.vision_tower = Tower() + self.multi_modal_projector = Proj() + + def forward(self, *args, **kwargs): + raise AssertionError("original PaliGemmaModel.forward should not run") + + DummyPaliGemma.forward = _patch_paligemma_image_features(DummyPaliGemma.forward) + pixel_values = torch.randn(2, 3, 8, 8, dtype=torch.float16) + out = DummyPaliGemma()(pixel_values) + assert out.shape == (2, 4, 8) + assert out.dtype == torch.float16 + + +@pytest.mark.unit +def test_pi05_language_model_keeps_hf_forward_without_rope(): + from torch_tensorrt.hf.exporters.models.pi05.patches import ( + _patch_pi05_language_model, + ) + + class Dummy(nn.Module): + def forward(self, inputs_embeds=None, past_key_values=None, **kwargs): + del past_key_values, kwargs + return inputs_embeds * 2 + + Dummy.forward = _patch_pi05_language_model(Dummy.forward) + hidden = torch.ones(1, 2, 4) + out = Dummy()(inputs_embeds=hidden, past_key_values="cache") + torch.testing.assert_close(out, hidden * 2) + + +@pytest.mark.unit +def test_pi05_action_keeps_training_forward_without_prefix_kv(): + from torch_tensorrt.hf.exporters.models.pi05.patches import ( + _patch_pi05_action_step_forward, + ) + + class Dummy(nn.Module): + def forward(self, images, img_masks, tokens, masks, actions, noise, time): + del img_masks, tokens, masks, actions, noise, time + return images + + Dummy.forward = _patch_pi05_action_step_forward(Dummy.forward) + assert Dummy()(7, 0, 0, 0, 0, 0, 0) == 7 + + +@pytest.mark.unit +def test_language_attn_plugin_when_rope_present(): + from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + _patch_language_attention, + ) + from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + _register_attention_plugin_op, + ) + + _register_attention_plugin_op() + + class Dummy(nn.Module): + def __init__(self): + super().__init__() + self.num_heads = 2 + self.num_key_value_heads = 2 + self.head_dim = 4 + self.q_proj = nn.Linear(8, 8) + self.k_proj = nn.Linear(8, 8) + self.v_proj = nn.Linear(8, 8) + self.o_proj = nn.Linear(8, 8) + + def forward(self, hidden_states, **kwargs): + raise AssertionError("HF forward should not run for plugin kwargs") + + Dummy.forward = _patch_language_attention(Dummy.forward) + hidden = torch.randn(1, 3, 8) + rope = torch.randn(1, 3, 4, dtype=torch.float32) + kv = torch.zeros(1, 2, 2, 8, 4) + ctx = torch.tensor([3], dtype=torch.int32) + start = torch.empty(0, dtype=torch.int32) + out, present = Dummy()( + hidden, + rope_rotary_cos_sin=rope, + past_key_value=kv, + ctx_len=ctx, + kvcache_start_index=start, + ) + assert out.shape == hidden.shape + assert present.shape == kv.shape + + +@pytest.mark.unit +def test_groot_backend_registers_components(): + from torch_tensorrt.hf.exporters.models.groot.patches import GROOT + from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + + paths = [p for p, _ in _PATCHES[GROOT]] + assert any("SiglipAttention.forward" in p for p in paths) + assert any("Qwen3Attention.forward" in p for p in paths) + assert any("Eagle25VLForConditionalGeneration.forward" in p for p in paths) + assert any("Qwen3ForCausalLM.forward" in p for p in paths) + assert any("GR00TN15.forward" in p for p in paths) + assert any("FlowmatchingActionHead.forward" in p for p in paths) + assert any("CategorySpecificLinear.forward" in p for p in paths) + + +@pytest.mark.unit +def test_nemotron_backend_registers_causal_lm(): + from torch_tensorrt.hf.exporters.models.nemotron.patches import NEMOTRON + from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + + paths = [p for p, _ in _PATCHES[NEMOTRON]] + assert any("NemotronHForCausalLM.forward" in p for p in paths) + + +@pytest.mark.unit +def test_eagle_vision_patch_extracts_features(): + from torch_tensorrt.hf.exporters.models.groot.patches import ( + _patch_eagle_image_features, + ) + + class Dummy: + def extract_feature(self, pixel_values): + return pixel_values + 1 + + def forward(self, *args, **kwargs): + raise AssertionError("full VLM forward should not run") + + Dummy.forward = _patch_eagle_image_features(Dummy.forward) + pixel_values = torch.zeros(1, 3, 4, 4) + torch.testing.assert_close(Dummy()(pixel_values), pixel_values + 1) + + +@pytest.mark.unit +def test_eagle_vision_keeps_vlm_forward_with_input_ids(): + from torch_tensorrt.hf.exporters.models.groot.patches import ( + _patch_eagle_image_features, + ) + + class Dummy: + def extract_feature(self, pixel_values): + raise AssertionError("extract_feature should not run") + + def forward(self, pixel_values, input_ids=None, **kwargs): + del kwargs + return (pixel_values, input_ids) + + Dummy.forward = _patch_eagle_image_features(Dummy.forward) + pixel_values = torch.zeros(1, 3, 4, 4) + input_ids = torch.ones(1, 2, dtype=torch.long) + out = Dummy()(pixel_values, input_ids) + assert out[1] is input_ids + + +@pytest.mark.unit +def test_groot_action_keeps_training_forward_without_context(): + from torch_tensorrt.hf.exporters.models.groot.patches import ( + _patch_groot_action_step_forward, + ) + + class Dummy(nn.Module): + def forward(self, backbone_output, action_input): + return backbone_output + + Dummy.forward = _patch_groot_action_step_forward(Dummy.forward) + assert Dummy()("backbone", "action") == "backbone" + + +@pytest.mark.unit +def test_groot_context_keeps_training_forward_without_hidden(): + from torch_tensorrt.hf.exporters.models.groot.patches import ( + _patch_groot_context_projection, + ) + + class Dummy(nn.Module): + def forward(self, backbone_inputs, action_inputs): + return backbone_inputs + + Dummy.forward = _patch_groot_context_projection(Dummy.forward) + assert Dummy()("backbone", "action") == "backbone" + + +@pytest.mark.unit +def test_nemotron_keeps_hf_forward_without_rope(): + from torch_tensorrt.hf.exporters.models.nemotron.patches import ( + _patch_nemotron_causal_lm, + ) + + class Dummy(nn.Module): + def forward(self, input_ids=None, inputs_embeds=None, **kwargs): + del input_ids, kwargs + return inputs_embeds * 2 + + Dummy.forward = _patch_nemotron_causal_lm(Dummy.forward) + hidden = torch.ones(1, 2, 4) + out = Dummy()(inputs_embeds=hidden) + torch.testing.assert_close(out, hidden * 2) + + +@pytest.mark.unit +def test_category_specific_linear_uses_index_select(): + from torch_tensorrt.hf.exporters.models.groot.patches import ( + _patch_category_specific_linear, + ) + + class Dummy(nn.Module): + def __init__(self): + super().__init__() + self.W = nn.Parameter( + torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4) + ) + self.b = nn.Parameter( + torch.arange(2 * 4, dtype=torch.float32).reshape(2, 4) + ) + + def forward(self, x, cat_ids): + raise AssertionError( + "original CategorySpecificLinear.forward should not run" + ) + + Dummy.forward = _patch_category_specific_linear(Dummy.forward) + layer = Dummy() + x = torch.ones(2, 5, 3) + cat_ids = torch.tensor([1, 0]) + out = layer(x, cat_ids) + expected = torch.bmm(x, layer.W[cat_ids]) + layer.b[cat_ids].unsqueeze(1) + torch.testing.assert_close(out, expected) From ac59e1c9cda906a79ea348fccb3c4527512566ed Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 16:33:20 -0700 Subject: [PATCH 05/11] Move the Nemotron mamba stub under the family package and document VLA export. Keep Hub load stubs next to the Nemotron spec, and mention PI05/GR00T run_vla.py in the LLM tools README. --- examples/dynamo/run_nemotron_export.py | 4 +++- .../{ => models/nemotron}/mamba_stub.py | 0 pyproject.toml | 1 - tools/llm/README.md | 24 +++++++++++++++++++ 4 files changed, 27 insertions(+), 2 deletions(-) rename py/torch_tensorrt/hf/exporters/{ => models/nemotron}/mamba_stub.py (100%) diff --git a/examples/dynamo/run_nemotron_export.py b/examples/dynamo/run_nemotron_export.py index 7c854ae0917..2df1af5139f 100644 --- a/examples/dynamo/run_nemotron_export.py +++ b/examples/dynamo/run_nemotron_export.py @@ -21,7 +21,9 @@ torch_tensorrt.__path__.append(_src_pkg) from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from torch_tensorrt.hf.exporters.mamba_stub import apply as apply_mamba_stub +from torch_tensorrt.hf.exporters.models.nemotron.mamba_stub import ( + apply as apply_mamba_stub, +) from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch from transformers import AutoModelForCausalLM, AutoTokenizer diff --git a/py/torch_tensorrt/hf/exporters/mamba_stub.py b/py/torch_tensorrt/hf/exporters/models/nemotron/mamba_stub.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/mamba_stub.py rename to py/torch_tensorrt/hf/exporters/models/nemotron/mamba_stub.py diff --git a/pyproject.toml b/pyproject.toml index d580cf19753..9dff51f61e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -458,7 +458,6 @@ module = [ "torch_tensorrt.hf.exporters.data", "torch_tensorrt.hf.exporters.rope", "torch_tensorrt.hf.exporters.prefix_cache", - "torch_tensorrt.hf.exporters.mamba_stub", ] ignore_errors = true diff --git a/tools/llm/README.md b/tools/llm/README.md index cb921bcadb5..bd6a0c6769b 100644 --- a/tools/llm/README.md +++ b/tools/llm/README.md @@ -6,6 +6,7 @@ This directory provides utilities and scripts for compiling, optimizing, and ben - **Model Support:** Works with popular LLMs such as Llama-3, Qwen2.5, etc. - **VLM Support:** Supports Visual Language Models like Qwen2.5-VL and Eagle2. +- **VLA Support:** Compiles PI05 and GR00T hybrid VLM graphs plus a separate action expert (`run_vla.py`). - **Precision Modes:** Supports FP16, BF16, and FP32. - **Quantization:** Supports FP8 and NVFP4 quantization formats for reduced memory usage and improved inference speed. - **KV Cache:** Supports static and dynamic KV cache for efficient autoregressive decoding. @@ -34,6 +35,13 @@ We have officially verified support for the following models: | Qwen 2.5 VL | Qwen/Qwen2.5-VL-3B-Instruct | FP16, FP32 | Yes | | Eagle2 | nvidia/Eagle2-2B | FP16, FP32 | Yes | +### Supported VLA Models + +| Policy | Model / alias | Precision | Notes | +|--------|---------------|-----------|-------| +| PI05 | `lerobot/pi05_libero_base` (`pi05`) | FP16, FP32 | Hybrid VLM + action expert | +| GR00T N1.5 | `nvidia/GR00T-N1.5-3B` (`groot`) | FP16, FP32 | Adds context projection island | + ### Usage #### Text-only LLMs: `run_llm.py` @@ -48,6 +56,19 @@ python run_llm.py --model meta-llama/Llama-3.2-1B-Instruct --prompt "What is par python run_vlm.py --model nvidia/Eagle2-2B --precision FP16 --num_tokens 128 --cache static_v1 --enable_pytorch_run --benchmark ``` +#### Vision Language Action policies: `run_vla.py` + +Compiles a hybrid VLM graph (vision + leftover fuse/scatter + language) and a +separate action expert. Chat template, tokenize, and Euler denoise stay on the +host. Requires the EdgeExporter prototype (`VLA_TEST_ROOT`, default +`/Test`) and `EDGE_LLM_PLUGIN_SO`. + +```bash +python run_vla.py --model pi05 --precision FP16 --dryrun +python run_vla.py --model lerobot/pi05_libero_base --precision FP16 --enable_pytorch_run --benchmark +python run_vla.py --model groot --precision FP16 --save /tmp/groot_engines +``` + #### Key Arguments - `--model`: Name or path of the HuggingFace LLM/VLM. @@ -62,6 +83,9 @@ python run_vlm.py --model nvidia/Eagle2-2B --precision FP16 --num_tokens 128 --c - `--cache`: KV cache type (`static_v1`, `static_v2`, or empty for no KV caching). - `--benchmark`: Enable benchmarking mode. - `--enable_pytorch_run`: Also run and compare PyTorch baseline. +- `--dryrun`: (VLA) Partition the hybrid graph without building TRT engines. +- `--save`: (VLA) Write named engines and `hybrid.pt2` to a directory. +- `--num_steps`: (VLA) Host Euler denoise steps. ### Quantization From cfe5aaa5f939786f0d01aaa054731bb496926fb9 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 3 Sep 2026 17:38:49 -0700 Subject: [PATCH 06/11] Add an Edge Exporter user guide. Document EdgeExporter.export, the returned ExportedProgram, family patches, and plugin converters, and link it from the user guide and HuggingFace tutorials. --- docsrc/tutorials/huggingface/index.rst | 7 +- docsrc/user_guide/edge_exporter.rst | 523 +++++++++++++++++++++++++ docsrc/user_guide/index.rst | 1 + 3 files changed, 528 insertions(+), 3 deletions(-) create mode 100644 docsrc/user_guide/edge_exporter.rst diff --git a/docsrc/tutorials/huggingface/index.rst b/docsrc/tutorials/huggingface/index.rst index 64e26f50d83..8041adae9e6 100644 --- a/docsrc/tutorials/huggingface/index.rst +++ b/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> diff --git a/docsrc/user_guide/edge_exporter.rst b/docsrc/user_guide/edge_exporter.rst new file mode 100644 index 00000000000..2ab1564bc00 --- /dev/null +++ b/docsrc/user_guide/edge_exporter.rst @@ -0,0 +1,523 @@ +.. _edge_exporter: + +Edge Exporter +============= + +Export a HuggingFace or LeRobot policy to TensorRT-Edge-LLM engines, then return +one ``torch.export`` graph that calls those engines. + +.. code-block:: python + + from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig + + exporter = EdgeExporter() + config = EdgeConfig(dryrun=True, 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 +shape as Transformers export: ``exporter.export(model, inputs, config)``. The +difference is what happens inside. Instead of tracing the whole policy as one +graph, Edge compiles one TensorRT engine per component, then records a small +outer graph that only *calls* those engines. + +.. note:: + + The Edge exporter is experimental. Family patches target a specific modeling + surface (LeRobot PI05 / GR00T, HuggingFace Nemotron-H) and TensorRT-Edge-LLM + plugins. Treat patches as tied to the versions you test against. + +.. list-table:: + :header-rows: 1 + :widths: 18 18 40 24 + + * - Family + - Spec key + - Engines + - Typical checkpoint + * - PI05 + - ``pi05`` + - ``vision``, ``language``, ``action`` + - ``lerobot/pi05_libero_base`` + * - GR00T + - ``groot`` + - ``vision``, ``language``, ``context_projection``, ``action`` + - ``nvidia/GR00T-N1.5-3B`` + * - Nemotron-H + - ``nemotron_h`` + - ``language`` + - ``nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16`` + +Installation +------------ + +You need Torch-TensorRT, TensorRT-Edge-LLM plugins, and the modeling stacks for +the families you export. + +.. code-block:: bash + + pip install transformers + +VLA families (PI05, GR00T) also need LeRobot. Nemotron-H uses Transformers only. + +For a real compile, set ``EDGE_LLM_PLUGIN_SO`` (or ``EDGELLM_PLUGIN_PATH`` / +``EDGELLM_TRT_PLUGIN_SO``) to ``libNvInfer_edgellm_plugin.so``. + +Before ``export()``, load the Edge-LLM plugins and force HuggingFace attention to +``eager`` (FlashAttention / SDPA are not the plugin path): + +.. code-block:: python + + from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt + from torch_tensorrt.hf.exporters.utils import force_hf_attention + + load_plugins_for_trt() + force_hf_attention(policy.model.paligemma_with_expert.paligemma.model.vision_tower, "eager") + force_hf_attention(policy.model.paligemma_with_expert.paligemma.model.language_model, "eager") + +Export a model +-------------- + +All families share one interface. Create an exporter, pass a policy or causal LM +plus sample inputs, and call ``export()``. + +.. code-block:: python + + from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig + from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt + + load_plugins_for_trt() + + exporter = EdgeExporter() + config = EdgeConfig( + 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) + + # engines: {"vision": ".../vision", "language": ".../language", "action": ".../action"} + print(exporter.engines) + + # run the exported program + outputs = exported.module()(**exporter.sample) + +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``). + +Export Program +-------------- + +``EdgeExporter.export`` returns an ``ExportedProgram``. Each policy +component (vision, language, action, …) is compiled into its own TensorRT +engine under ``engine_dir//``. The exported graph is the runtime that +calls those engines in order through ``torch.ops.edge_llm.execute_engine``. +A packing op sits between vision and language so image tokens land in the +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`` → +``language`` → ``context_projection`` → ``action``): + +.. code-block:: text + + graph(): + %pixel_values : [num_users=1] = placeholder[target=pixel_values] + %lang_embeds : [num_users=1] = placeholder[target=lang_embeds] + %image_token_mask : [num_users=1] = placeholder[target=image_token_mask] + %rope_rotary_cos_sin : [num_users=1] = placeholder[target=rope_rotary_cos_sin] + %context_lengths : [num_users=1] = placeholder[target=context_lengths] + %kvcache_start_index : [num_users=1] = placeholder[target=kvcache_start_index] + %last_token_ids : [num_users=1] = placeholder[target=last_token_ids] + %ds_stack : [num_users=1] = placeholder[target=ds_stack] + %past_key_values_0 : [num_users=1] = placeholder[target=past_key_values_0] + %step_actions : [num_users=1] = placeholder[target=step_actions] + %step_timestep : [num_users=1] = placeholder[target=step_timestep] + %state : [num_users=1] = placeholder[target=state] + %embodiment_id : [num_users=1] = placeholder[target=embodiment_id] + %execute_engine : [num_users=1] = call_function[target=torch.ops.edge_llm.execute_engine.default](args = (edge_engines/vision, vision, [%pixel_values]), kwargs = {}) + %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine, 0), kwargs = {}) + %scatter_image_tokens : [num_users=1] = call_function[target=torch.ops.edge_llm.scatter_image_tokens.default](args = (%getitem, %lang_embeds, %image_token_mask), kwargs = {}) + %execute_engine_1 : [num_users=4] = call_function[target=torch.ops.edge_llm.execute_engine.default](args = (edge_engines/language, language, [%scatter_image_tokens, %rope_rotary_cos_sin, %context_lengths, %kvcache_start_index, %last_token_ids, %ds_stack, %past_key_values_0]), kwargs = {}) + %getitem_1 : [num_users=0] = call_function[target=operator.getitem](args = (%execute_engine_1, 0), kwargs = {}) + %getitem_2 : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_1, 1), kwargs = {}) + %getitem_3 : [num_users=0] = call_function[target=operator.getitem](args = (%execute_engine_1, 2), kwargs = {}) + %getitem_4 : [num_users=0] = call_function[target=operator.getitem](args = (%execute_engine_1, 3), kwargs = {}) + %execute_engine_2 : [num_users=1] = call_function[target=torch.ops.edge_llm.execute_engine.default](args = (edge_engines/context_projection, context_projection, [%getitem_2]), kwargs = {}) + %getitem_5 : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_2, 0), kwargs = {}) + %execute_engine_3 : [num_users=1] = call_function[target=torch.ops.edge_llm.execute_engine.default](args = (edge_engines/action, action, [%step_actions, %step_timestep, %getitem_5, %state, %embodiment_id]), kwargs = {}) + %getitem_6 : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_3, 0), kwargs = {}) + return (getitem_6,) + +PI05 is the same shape with ``fuse_prefix`` instead of ``scatter_image_tokens`` +and no ``context_projection`` engine. Nemotron is a single +``execute_engine`` on ``language``. + +``execute_engine`` is a Torch custom op, not a normal ``nn.Module`` call. Family +``spec.run()`` calls ``call_engine(...)``, so ``torch.export`` records **one +node per engine**. Matching ``register_fake`` kernels give Dynamo the output +shapes. Two packing ops live in the same file +(``py/torch_tensorrt/hf/exporters/ops.py``): + +* ``edge_llm::fuse_prefix`` — PI05: concat vision tokens with language + embeddings and gather the compact prefix. +* ``edge_llm::scatter_image_tokens`` — GR00T: write vision tokens into the + ```` slots of the language embeddings. + +These appear in the **outer** ExportedProgram. They are not TensorRT plugins. + +Patches +------- + +HuggingFace ``DynamoExporter`` uses ``@register_patch`` plus a temporary class +``setattr``. Edge uses the same contract. + +Each family has a ``patches.py`` that registers factories on a backend name +(``"pi05"``, ``"groot"``, ``"nemotron"``). A factory receives the **original** +``Class.forward`` and returns a replacement. ``apply_patches(backend)`` resolves +the dotted class path and does ``setattr(Cls, "forward", factory(original))`` +for the duration of ``export()``. After a real compile the original methods are +restored. Dryrun leaves the replacements installed so ``execute_engine`` still +hits the patched Python modules. + +.. code-block:: python + + from torch_tensorrt.hf.exporters.plugin.attn_patches import register_patch + + PI05 = "pi05" + + @register_patch( + PI05, + "transformers.models.paligemma.modeling_paligemma.PaliGemmaModel.forward", + ) + def _patch_paligemma_image_features(_original): + def forward(self, pixel_values, **kwargs): + image_outputs = self.vision_tower(pixel_values, **kwargs) + hidden = image_outputs.last_hidden_state + return self.multi_modal_projector(hidden) + + return forward + +The replacement is the thing TensorRT traces. You compile the **original +submodule** (``PaliGemmaModel``, ``PiGemmaModel``, ``FlowmatchingActionHead``, …), +not a wrapper ``nn.Module``. The patched ``forward`` is what makes that submodule +look like an Edge engine: a tensor in, a tensor out, plugin attention inside. + +When the same class is used in two roles, the patched ``forward`` dispatches. +PI05 language is ``PiGemmaModel`` for both the language tower and the action +expert. Edge prefill passes ``rope_rotary_cos_sin``; the action expert does not. +If that argument is missing, the original HuggingFace forward runs: + +.. code-block:: python + + def forward(self, inputs_embeds=None, rope_rotary_cos_sin=None, **kwargs): + if rope_rotary_cos_sin is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + return causal_lm_plugin_forward(self, inputs_embeds, rope_rotary_cos_sin, ...) + +Attention patches follow the same rule: ``GemmaAttention.forward`` uses the +language plugin when ``rope_rotary_cos_sin`` is present, otherwise eager HF +attention. + +The spec installs the whole family backend once around the component loop: + +.. code-block:: python + + class Pi05Spec(EdgeSpec): + def apply_patches(self, model=None): + from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + return apply_patches("pi05") + +Nemotron also wraps hybrid mixers on the live model inside ``apply_patches`` +(MoE packing needs the instance). That is still not a separate export wrapper; +the compiled module is the original ``NemotronHForCausalLM``. + +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. + +Create ``py/torch_tensorrt/hf/exporters/models//``: + +.. code-block:: text + + / + spec.py # EdgeSpec: components, sample inputs, prepare, run + patches.py # @register_patch factories on this family's backend + helpers.py # optional packing / submodule lookup + +Import the spec from ``models/__init__.py`` so registration happens when the +exporter package loads: + +.. code-block:: python + + from torch_tensorrt.hf.exporters.models.my_vla import spec as _my_vla # noqa: F401 + +1. Register the spec +^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: python + + from torch_tensorrt.hf.exporters.spec import EdgeSpec, register_edge_spec + + @register_edge_spec("my_vla") + class MyVlaSpec(EdgeSpec): + components = ("vision", "language", "action") + + def apply_patches(self, model=None): + from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + from .patches import MY_VLA + return apply_patches(MY_VLA) + +Add a heuristic in ``infer_model_type`` (in ``exporters/spec.py``) if you want +``model_type`` to be optional, or always pass ``EdgeConfig(model_type="my_vla")``. + +2. Patch original ``forward`` methods +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +In ``patches.py``, register one factory per class you need to change. Put the +Edge I/O in that ``forward``. Compile the original module; do not introduce a +wrapper class. + +Reuse the shared plugin attention factories when the layout matches +(``_patch_vision_attention``, ``_patch_language_attention``). Reuse +``causal_lm_plugin_forward`` for decoder-only prefill. + +3. Select submodules and flatten I/O +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``prepare(name, model, sample, upstream, config)`` returns a ``ComponentBundle``: + +* ``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). + +4. Call engines in ``run()`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``run(engines, sample)`` is the outer graph. Call ``call_engine`` for each +component and keep packing in Python (``fuse_prefix``, ``scatter_image_tokens``, +or your own custom op). + +.. code-block:: python + + def run(self, engines, sample): + vis = call_engine(engines["vision"], "vision", sample["pixel_values"])[0] + prefix = fuse_prefix(vis, sample["lang_embeds"], sample["compact_index"]) + lm = call_engine(engines["language"], "language", prefix, ...) + return call_engine( + engines["action"], "action", sample["step_actions"], ..., lm[2], lm[3] + ) + +5. Collate sample inputs +^^^^^^^^^^^^^^^^^^^^^^^^ + +``prepare_sample_inputs`` turns the caller payload into the stem dict ``prepare`` +and ``run`` share. If the caller already passed tensors (``pixel_values``, +``input_ids``, …), return them. Otherwise load a preprocessor / tokenizer here +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. + +.. code-block:: bash + + cd TensorRT/examples/dynamo + + 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 + +.. list-table:: + :header-rows: 1 + :widths: 28 32 40 + + * - Script + - What to pass + - Engines + * - ``run_pi05_export.py`` + - LeRobot ``PI05Policy`` + - vision, language, action + * - ``run_groot_export.py`` + - LeRobot ``GrootPolicy`` + - vision, language, context_projection, action + * - ``run_nemotron_export.py`` + - HuggingFace ``NemotronHForCausalLM`` + - language + +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 +------------- + +``EdgeConfig`` knobs that matter for this path: + +.. list-table:: + :header-rows: 1 + :widths: 24 16 60 + + * - Field + - Default + - Role + * - ``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 + * - ``trt_settings`` + - ``{}`` + - Forwarded into ``torch_tensorrt.dynamo.compile`` + +``strict``, ``dynamic``, and ``dynamic_shapes`` match HuggingFace ``DynamoConfig`` +and apply to the **outer** ``torch.export`` of the runtime, not to the +per-component TensorRT compiles. + +Plugins vs ``execute_engine`` +----------------------------- + +There are two different custom-op namespaces. Do not mix them up. + +.. list-table:: + :header-rows: 1 + :widths: 28 38 34 + + * - Where + - Ops + - What you see + * - Outer ``ExportedProgram`` + - ``edge_llm::execute_engine``, ``fuse_prefix``, ``scatter_image_tokens`` + - ``print(program.graph)`` after ``EdgeExporter.export`` + * - Inside one component engine + - ``trt::attention_plugin``, ``trt::vit_attention_plugin``, + ``trt::causal_conv1d``, ``trt::update_ssm_state``, + ``trt::nvfp4_moe_plugin`` + - The FX graph of that component during ``dynamo.compile`` + +A patched attention ``forward`` calls ``torch.ops.trt.attention_plugin``. +When that component is ``torch.export`` + ``torch_tensorrt.dynamo.compile``'d, +the plugin converter turns that op into an ``AttentionPlugin`` layer inside +that engine. The outer ``ExportedProgram`` never sees it; it only sees +``execute_engine("language")``. + + +Plugin Converters +----------------- + +A **plugin converter** is the Torch-TensorRT dynamo hook that maps a +``torch.ops.trt.*`` node onto a TensorRT ``IPluginV3`` (from +``libNvInfer_edgellm_plugin.so``). + +Without a converter, ``dynamo.compile`` cannot lower the custom op and +either graph-breaks or fails. With a converter, the op becomes one TensorRT +plugin layer. + +Converters live in +``py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py`` and are +registered with ``@dynamo_tensorrt_converter``. Example for ViT attention: + +.. code-block:: python + + @dynamo_tensorrt_converter( + torch.ops.trt.vit_attention_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, + ) + def convert_vit_attention_plugin(ctx, target, args, kwargs, name): + creator = get_trt_plugin_creator("ViTAttentionPlugin", "1", "") + plugin = creator.create_plugin(name, fields, trt.TensorRTPhase.BUILD) + layer = ctx.net.add_plugin_v3(inputs, [], plugin) + return layer.get_output(0) + +That is: look up the plugin by TensorRT name, fill ``PluginField``s, add the +layer, return its outputs. ``load_plugins_for_trt()`` imports this module so +the converters are registered before compile. + + +How to add a plugin +------------------- + +You need four pieces. The C++ plugin in Edge-LLM is assumed to already +exist and be in ``libNvInfer_edgellm_plugin.so``. + +**1. Torch custom op (eager + fake)** so Dynamo can trace. + +Register in ``plugin_utils.py``, ``mamba.py``, or ``moe.py``: + +.. code-block:: python + + @torch.library.custom_op("trt::my_plugin", mutates_args=()) + def my_plugin(x: torch.Tensor, ...) -> torch.Tensor: + return _my_plugin_eager(x, ...) + + @my_plugin.register_fake + def _(x, ...): + return torch.empty_like(x) + +Call ``load_plugins_for_trt()`` so this op exists before export. + +**2. Call it from a patched ``forward``.** + +In ``models//patches.py``, ``@register_patch`` a class ``forward`` +that emits ``torch.ops.trt.my_plugin`` when Edge I/O is present. That is +how attention already works: the plugin call is in the patched +``GemmaAttention.forward``, not in a wrapper module. + +**3. Plugin converter.** + +Add ``@dynamo_tensorrt_converter(torch.ops.trt.my_plugin.default)`` in +``plugin_converter.py``. It must use the same TensorRT plugin name / version +the ``.so`` registered (see ``get_trt_plugin_creator``). + +**4. Load the ``.so``.** + +``load_plugins_for_trt()`` already calls ``load_plugin()``, which +``ctypes.CDLL``s ``EDGE_LLM_PLUGIN_SO``. After that, TensorRT can create +the plugin during ``dynamo.compile``. + +Checklist for a new kernel: + +* Eager op + fake kernel (traceable). +* Patched ``forward`` that actually calls the op. +* Converter that adds the TRT plugin layer. +* Plugin present in ``libNvInfer_edgellm_plugin.so``. +* ``load_plugins_for_trt()`` before ``EdgeExporter().export``. diff --git a/docsrc/user_guide/index.rst b/docsrc/user_guide/index.rst index e46c61b5cfd..09be207dcf6 100644 --- a/docsrc/user_guide/index.rst +++ b/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 From bcf1fac963d996142d2aa0815342d28a8f3c3d2e Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Fri, 4 Sep 2026 23:23:24 -0700 Subject: [PATCH 07/11] Move Edge exporter to tools/hf and add compile-time parity. --- docsrc/user_guide/edge_exporter.rst | 111 ++++++++++++------ py/torch_tensorrt/hf/__init__.py | 4 - py/torch_tensorrt/hf/exporters/__init__.py | 26 ---- .../hf/exporters/models/__init__.py | 5 - pyproject.toml | 10 -- setup.py | 16 --- tools/hf/__init__.py | 4 + tools/hf/exporters/__init__.py | 11 ++ .../hf/exporters/compile.py | 28 ++++- .../hf/exporters/config.py | 2 +- .../hf/exporters/data.py | 0 .../hf/exporters/exporter.py | 14 ++- tools/hf/exporters/measure.py | 62 ++++++++++ tools/hf/exporters/models/__init__.py | 5 + .../hf/exporters/models/common/__init__.py | 0 .../hf/exporters/models/common/helpers.py | 2 +- .../hf/exporters/models/common/patches.py | 0 .../hf/exporters/models/groot/__init__.py | 0 .../hf/exporters/models/groot/helpers.py | 0 .../hf/exporters/models/groot/patches.py | 30 ++++- .../hf/exporters/models/groot/spec.py | 29 +++-- .../hf/exporters/models/nemotron/__init__.py | 0 .../hf/exporters/models/nemotron/helpers.py | 0 .../exporters/models/nemotron/mamba_stub.py | 0 .../hf/exporters/models/nemotron/patches.py | 8 +- .../hf/exporters/models/nemotron/spec.py | 21 ++-- .../hf/exporters/models/pi05/__init__.py | 0 .../hf/exporters/models/pi05/helpers.py | 0 .../hf/exporters/models/pi05/patches.py | 8 +- .../hf/exporters/models/pi05/spec.py | 28 +++-- .../hf/exporters/ops.py | 0 .../hf/exporters/plugin/__init__.py | 0 .../hf/exporters/plugin/attention.py | 0 .../hf/exporters/plugin/attn_patches.py | 3 +- .../hf/exporters/plugin/mamba.py | 0 .../hf/exporters/plugin/moe.py | 0 .../hf/exporters/plugin/plugin_converter.py | 0 .../hf/exporters/plugin/plugin_utils.py | 0 .../hf/exporters/prefix_cache.py | 0 .../hf/exporters/rope.py | 0 .../hf/exporters/runtime.py | 3 +- .../hf/exporters/spec.py | 0 .../hf/exporters/tests}/test_edge_exporter.py | 95 +++++++++++---- .../hf/exporters/utils.py | 11 -- .../dynamo => tools/hf}/run_groot_export.py | 26 ++-- .../hf}/run_nemotron_export.py | 25 ++-- .../dynamo => tools/hf}/run_pi05_export.py | 31 +++-- 47 files changed, 376 insertions(+), 242 deletions(-) delete mode 100644 py/torch_tensorrt/hf/__init__.py delete mode 100644 py/torch_tensorrt/hf/exporters/__init__.py delete mode 100644 py/torch_tensorrt/hf/exporters/models/__init__.py create mode 100644 tools/hf/__init__.py create mode 100644 tools/hf/exporters/__init__.py rename {py/torch_tensorrt => tools}/hf/exporters/compile.py (81%) rename {py/torch_tensorrt => tools}/hf/exporters/config.py (92%) rename {py/torch_tensorrt => tools}/hf/exporters/data.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/exporter.py (94%) create mode 100644 tools/hf/exporters/measure.py create mode 100644 tools/hf/exporters/models/__init__.py rename {py/torch_tensorrt => tools}/hf/exporters/models/common/__init__.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/common/helpers.py (96%) rename {py/torch_tensorrt => tools}/hf/exporters/models/common/patches.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/groot/__init__.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/groot/helpers.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/groot/patches.py (85%) rename {py/torch_tensorrt => tools}/hf/exporters/models/groot/spec.py (94%) rename {py/torch_tensorrt => tools}/hf/exporters/models/nemotron/__init__.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/nemotron/helpers.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/nemotron/mamba_stub.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/nemotron/patches.py (91%) rename {py/torch_tensorrt => tools}/hf/exporters/models/nemotron/spec.py (91%) rename {py/torch_tensorrt => tools}/hf/exporters/models/pi05/__init__.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/pi05/helpers.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/models/pi05/patches.py (95%) rename {py/torch_tensorrt => tools}/hf/exporters/models/pi05/spec.py (93%) rename {py/torch_tensorrt => tools}/hf/exporters/ops.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/__init__.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/attention.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/attn_patches.py (99%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/mamba.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/moe.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/plugin_converter.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/plugin/plugin_utils.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/prefix_cache.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/rope.py (100%) rename {py/torch_tensorrt => tools}/hf/exporters/runtime.py (90%) rename {py/torch_tensorrt => tools}/hf/exporters/spec.py (100%) rename {tests/py/dynamo/hf => tools/hf/exporters/tests}/test_edge_exporter.py (82%) rename {py/torch_tensorrt => tools}/hf/exporters/utils.py (75%) rename {examples/dynamo => tools/hf}/run_groot_export.py (81%) rename {examples/dynamo => tools/hf}/run_nemotron_export.py (80%) rename {examples/dynamo => tools/hf}/run_pi05_export.py (79%) diff --git a/docsrc/user_guide/edge_exporter.rst b/docsrc/user_guide/edge_exporter.rst index 2ab1564bc00..7263ebeb39a 100644 --- a/docsrc/user_guide/edge_exporter.rst +++ b/docsrc/user_guide/edge_exporter.rst @@ -8,7 +8,7 @@ one ``torch.export`` graph that calls those engines. .. code-block:: python - from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig + from hf.exporters import EdgeExporter, EdgeConfig exporter = EdgeExporter() config = EdgeConfig(dryrun=True, engine_dir="/tmp/pi05_edge") @@ -20,6 +20,9 @@ difference is what happens inside. Instead of tracing the whole policy as one graph, Edge compiles one TensorRT engine per component, then records a small outer graph that only *calls* those engines. +The code is in ``tools/hf``. The entry points are ``tools/hf/run_pi05_export.py``, +``run_groot_export.py``, and ``run_nemotron_export.py``. + .. note:: The Edge exporter is experimental. Family patches target a specific modeling @@ -67,8 +70,8 @@ Before ``export()``, load the Edge-LLM plugins and force HuggingFace attention t .. code-block:: python - from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt - from torch_tensorrt.hf.exporters.utils import force_hf_attention + from hf.exporters.plugin.plugin_utils import load_plugins_for_trt + from hf.exporters.utils import force_hf_attention load_plugins_for_trt() force_hf_attention(policy.model.paligemma_with_expert.paligemma.model.vision_tower, "eager") @@ -82,8 +85,8 @@ plus sample inputs, and call ``export()``. .. code-block:: python - from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig - from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt + from hf.exporters import EdgeExporter, EdgeConfig + from hf.exporters.plugin.plugin_utils import load_plugins_for_trt load_plugins_for_trt() @@ -166,7 +169,7 @@ and no ``context_projection`` engine. Nemotron is a single ``spec.run()`` calls ``call_engine(...)``, so ``torch.export`` records **one node per engine**. Matching ``register_fake`` kernels give Dynamo the output shapes. Two packing ops live in the same file -(``py/torch_tensorrt/hf/exporters/ops.py``): +(``tools/hf/exporters/ops.py``): * ``edge_llm::fuse_prefix`` — PI05: concat vision tokens with language embeddings and gather the compact prefix. @@ -178,20 +181,26 @@ These appear in the **outer** ExportedProgram. They are not TensorRT plugins. Patches ------- -HuggingFace ``DynamoExporter`` uses ``@register_patch`` plus a temporary class -``setattr``. Edge uses the same contract. +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). + +``@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()``. -Each family has a ``patches.py`` that registers factories on a backend name -(``"pi05"``, ``"groot"``, ``"nemotron"``). A factory receives the **original** -``Class.forward`` and returns a replacement. ``apply_patches(backend)`` resolves -the dotted class path and does ``setattr(Cls, "forward", factory(original))`` -for the duration of ``export()``. After a real compile the original methods are -restored. Dryrun leaves the replacements installed so ``execute_engine`` still -hits the patched Python modules. +HuggingFace ``DynamoExporter`` uses the same two steps. The purpose is +different. HF patches make the original modeling ``forward`` traceable. Edge +patches change ``forward`` first so TensorRT traces plugin ops +(``torch.ops.trt.*``), not HuggingFace attention. After compile, the outer +graph is ``spec.run()`` → ``execute_engine``. There is no HF attention left +to patch, so the HuggingFace ``"dynamo"`` registry does not apply. .. code-block:: python - from torch_tensorrt.hf.exporters.plugin.attn_patches import register_patch + from exporters.plugin.attn_patches import register_patch PI05 = "pi05" @@ -207,15 +216,15 @@ hits the patched Python modules. return forward -The replacement is the thing TensorRT traces. You compile the **original -submodule** (``PaliGemmaModel``, ``PiGemmaModel``, ``FlowmatchingActionHead``, …), -not a wrapper ``nn.Module``. The patched ``forward`` is what makes that submodule -look like an Edge engine: a tensor in, a tensor out, plugin attention inside. +You compile the original submodule (``PaliGemmaModel``, ``PiGemmaModel``, +``FlowmatchingActionHead``, …), not a wrapper ``nn.Module``. The patched +``forward`` is what TensorRT traces: a tensor in, a tensor out, plugin +attention inside. -When the same class is used in two roles, the patched ``forward`` dispatches. -PI05 language is ``PiGemmaModel`` for both the language tower and the action -expert. Edge prefill passes ``rope_rotary_cos_sin``; the action expert does not. -If that argument is missing, the original HuggingFace forward runs: +When the same class is used twice (PI05 language vs action expert), the +patched ``forward`` checks for Edge arguments such as +``rope_rotary_cos_sin``. If they are missing, the original HuggingFace +``forward`` runs: .. code-block:: python @@ -228,18 +237,52 @@ Attention patches follow the same rule: ``GemmaAttention.forward`` uses the language plugin when ``rope_rotary_cos_sin`` is present, otherwise eager HF attention. -The spec installs the whole family backend once around the component loop: +The spec installs the family once around the component loop: .. code-block:: python class Pi05Spec(EdgeSpec): def apply_patches(self, model=None): - from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + from exporters.plugin.attn_patches import apply_patches return apply_patches("pi05") -Nemotron also wraps hybrid mixers on the live model inside ``apply_patches`` -(MoE packing needs the instance). That is still not a separate export wrapper; -the compiled module is the original ``NemotronHForCausalLM``. +When the decorator is enough +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +If ``type(module)`` is the class named in the dotted path, +``@register_patch`` plus ``apply_patches`` is all you need. That is Siglip, +Qwen3, Llama, ``GR00TN15``, the action head, and so on. + +When it is not +^^^^^^^^^^^^^^ + +The path must be the **same class object** as the live module. A look-alike +file under another import is a different class. ``setattr`` on one does not +change the other. + +``trust_remote_code=True`` downloads Hub ``.py`` files into HuggingFace's +module cache and imports them. ``AutoModel.from_config`` then builds an +instance in memory. That class's module path looks like +``transformers_modules...modeling_...``. It is not stable and +does not exist until load, so the decorator cannot name it. The cache stores +source, not the ``nn.Module``. + +GR00T's Eagle is this case. The LeRobot path +``lerobot.policies.groot.eagle2_hg_model....Eagle25VLForConditionalGeneration`` +is a different class from the HuggingFace cache copy that ``from_config`` +actually constructs. + +Live-object patches +^^^^^^^^^^^^^^^^^^^ + +``apply_groot_patches(model)`` is not a second decorator. It is the place +that has the instance, so it can patch ``type(eagle_model)``. It still runs +``apply_patches("groot")`` for every class that has a stable path. + +Nemotron's ``apply_nemotron_patches(model)`` is the same idea for mixers: +the registry is string paths; anything that only exists on the live object +needs ``model``. The compiled module is still the original +``NemotronHForCausalLM``, not a wrapper. Add a new model --------------- @@ -248,7 +291,7 @@ Add a new model and loops ``spec.components``. A new architecture is a new spec plus a patch backend. -Create ``py/torch_tensorrt/hf/exporters/models//``: +Create ``tools/hf/exporters/models//``: .. code-block:: text @@ -262,21 +305,21 @@ exporter package loads: .. code-block:: python - from torch_tensorrt.hf.exporters.models.my_vla import spec as _my_vla # noqa: F401 + from hf.exporters.models.my_vla import spec as _my_vla # noqa: F401 1. Register the spec ^^^^^^^^^^^^^^^^^^^^ .. code-block:: python - from torch_tensorrt.hf.exporters.spec import EdgeSpec, register_edge_spec + from hf.exporters.spec import EdgeSpec, register_edge_spec @register_edge_spec("my_vla") class MyVlaSpec(EdgeSpec): components = ("vision", "language", "action") def apply_patches(self, model=None): - from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + from hf.exporters.plugin.attn_patches import apply_patches from .patches import MY_VLA return apply_patches(MY_VLA) @@ -452,7 +495,7 @@ either graph-breaks or fails. With a converter, the op becomes one TensorRT plugin layer. Converters live in -``py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py`` and are +``tools/hf/exporters/plugin/plugin_converter.py`` and are registered with ``@dynamo_tensorrt_converter``. Example for ViT attention: .. code-block:: python diff --git a/py/torch_tensorrt/hf/__init__.py b/py/torch_tensorrt/hf/__init__.py deleted file mode 100644 index 1a52009e770..00000000000 --- a/py/torch_tensorrt/hf/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""HuggingFace-facing export helpers for Torch-TensorRT. - -Use ``from torch_tensorrt.hf.exporters import EdgeExporter, EdgeConfig``. -""" diff --git a/py/torch_tensorrt/hf/exporters/__init__.py b/py/torch_tensorrt/hf/exporters/__init__.py deleted file mode 100644 index 01dd627defa..00000000000 --- a/py/torch_tensorrt/hf/exporters/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -from torch_tensorrt.hf.exporters.config import EdgeConfig -from torch_tensorrt.hf.exporters.exporter import EdgeExporter -from torch_tensorrt.hf.exporters.models.groot.spec import ( # noqa: F401 - GrootSpec as _GrootSpec, -) -from torch_tensorrt.hf.exporters.models.nemotron.spec import ( # noqa: F401 - NemotronSpec as _NemotronSpec, -) -from torch_tensorrt.hf.exporters.models.pi05.spec import ( # noqa: F401 - Pi05Spec as _Pi05Spec, -) -from torch_tensorrt.hf.exporters.spec import ( - ComponentBundle, - EdgeSpec, - get_edge_spec, - register_edge_spec, -) - -__all__ = [ - "ComponentBundle", - "EdgeConfig", - "EdgeExporter", - "EdgeSpec", - "get_edge_spec", - "register_edge_spec", -] diff --git a/py/torch_tensorrt/hf/exporters/models/__init__.py b/py/torch_tensorrt/hf/exporters/models/__init__.py deleted file mode 100644 index bfb17db20d4..00000000000 --- a/py/torch_tensorrt/hf/exporters/models/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Model families. Importing this package registers EdgeSpecs.""" - -from torch_tensorrt.hf.exporters.models.groot import spec as _groot # noqa: F401 -from torch_tensorrt.hf.exporters.models.nemotron import spec as _nemotron # noqa: F401 -from torch_tensorrt.hf.exporters.models.pi05 import spec as _pi05 # noqa: F401 diff --git a/pyproject.toml b/pyproject.toml index 9dff51f61e2..f48d7abdc4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -451,16 +451,6 @@ module = "torch_tensorrt.fx.*" ignore_errors = true follow_imports = "skip" -[[tool.mypy.overrides]] -module = [ - "torch_tensorrt.hf.exporters.plugin.*", - "torch_tensorrt.hf.exporters.models.*", - "torch_tensorrt.hf.exporters.data", - "torch_tensorrt.hf.exporters.rope", - "torch_tensorrt.hf.exporters.prefix_cache", -] -ignore_errors = true - [tool.typos] files.extend-exclude = [ "docs/**/*", diff --git a/setup.py b/setup.py index 64a039f61fa..1c994f6f39b 100644 --- a/setup.py +++ b/setup.py @@ -616,14 +616,6 @@ def run(self): "torch_tensorrt.dynamo.runtime", "torch_tensorrt.dynamo.tools", "torch_tensorrt.executorch", - "torch_tensorrt.hf", - "torch_tensorrt.hf.exporters", - "torch_tensorrt.hf.exporters.models", - "torch_tensorrt.hf.exporters.models.common", - "torch_tensorrt.hf.exporters.models.groot", - "torch_tensorrt.hf.exporters.models.nemotron", - "torch_tensorrt.hf.exporters.models.pi05", - "torch_tensorrt.hf.exporters.plugin", "torch_tensorrt.runtime", ] @@ -663,14 +655,6 @@ def run(self): "torch_tensorrt.dynamo.runtime": "py/torch_tensorrt/dynamo/runtime", "torch_tensorrt.dynamo.tools": "py/torch_tensorrt/dynamo/tools", "torch_tensorrt.executorch": "py/torch_tensorrt/executorch", - "torch_tensorrt.hf": "py/torch_tensorrt/hf", - "torch_tensorrt.hf.exporters": "py/torch_tensorrt/hf/exporters", - "torch_tensorrt.hf.exporters.models": "py/torch_tensorrt/hf/exporters/models", - "torch_tensorrt.hf.exporters.models.common": "py/torch_tensorrt/hf/exporters/models/common", - "torch_tensorrt.hf.exporters.models.groot": "py/torch_tensorrt/hf/exporters/models/groot", - "torch_tensorrt.hf.exporters.models.nemotron": "py/torch_tensorrt/hf/exporters/models/nemotron", - "torch_tensorrt.hf.exporters.models.pi05": "py/torch_tensorrt/hf/exporters/models/pi05", - "torch_tensorrt.hf.exporters.plugin": "py/torch_tensorrt/hf/exporters/plugin", "torch_tensorrt.runtime": "py/torch_tensorrt/runtime", } diff --git a/tools/hf/__init__.py b/tools/hf/__init__.py new file mode 100644 index 00000000000..a58922e44e3 --- /dev/null +++ b/tools/hf/__init__.py @@ -0,0 +1,4 @@ +"""HuggingFace-facing export helpers for Torch-TensorRT. + +Use ``from exporters import EdgeExporter, EdgeConfig``. +""" diff --git a/tools/hf/exporters/__init__.py b/tools/hf/exporters/__init__.py new file mode 100644 index 00000000000..2ed7d30ed14 --- /dev/null +++ b/tools/hf/exporters/__init__.py @@ -0,0 +1,11 @@ +from .config import EdgeConfig +from .exporter import EdgeExporter +from .models.groot.spec import GrootSpec as _GrootSpec # 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, +) diff --git a/py/torch_tensorrt/hf/exporters/compile.py b/tools/hf/exporters/compile.py similarity index 81% rename from py/torch_tensorrt/hf/exporters/compile.py rename to tools/hf/exporters/compile.py index a8b9dee20d6..81777e5c66c 100644 --- a/py/torch_tensorrt/hf/exporters/compile.py +++ b/tools/hf/exporters/compile.py @@ -6,14 +6,17 @@ import torch import torch_tensorrt -from torch_tensorrt.hf.exporters.ops import _as_tuple, record_engine -from torch_tensorrt.hf.exporters.spec import ComponentBundle + +from .measure import cuda_ms, parity +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) | { @@ -33,6 +36,7 @@ def compile_component( 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, ...]]: """Export one component, compile it, write ``engine_dir//``. @@ -42,7 +46,7 @@ def compile_component( Family setattr is owned by ``EdgeSpec.apply_patches``, not this helper. ``dryrun`` records the patched eager module for ``execute_engine``. """ - from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + from .plugin.attn_patches import ( set_language_mask_type, ) @@ -89,6 +93,22 @@ def compile_component( 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 + ): + 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, @@ -111,7 +131,7 @@ def compile_component( return engine_path, outputs finally: if not dryrun and patched is not None: - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + from .plugin.plugin_utils import ( restore_attention, ) diff --git a/py/torch_tensorrt/hf/exporters/config.py b/tools/hf/exporters/config.py similarity index 92% rename from py/torch_tensorrt/hf/exporters/config.py rename to tools/hf/exporters/config.py index ae9f140366e..d81055589b8 100644 --- a/py/torch_tensorrt/hf/exporters/config.py +++ b/tools/hf/exporters/config.py @@ -7,7 +7,7 @@ @dataclass class EdgeConfig: - """Knobs for :class:`~torch_tensorrt.hf.exporters.EdgeExporter`. + """Knobs for :class:`~exporters.EdgeExporter`. ``strict`` / ``dynamic`` / ``dynamic_shapes`` match HuggingFace ``DynamoConfig`` so this can subclass it later without an API break. diff --git a/py/torch_tensorrt/hf/exporters/data.py b/tools/hf/exporters/data.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/data.py rename to tools/hf/exporters/data.py diff --git a/py/torch_tensorrt/hf/exporters/exporter.py b/tools/hf/exporters/exporter.py similarity index 94% rename from py/torch_tensorrt/hf/exporters/exporter.py rename to tools/hf/exporters/exporter.py index f619ac01b23..f93ca57bc08 100644 --- a/py/torch_tensorrt/hf/exporters/exporter.py +++ b/tools/hf/exporters/exporter.py @@ -11,13 +11,14 @@ import torch import torch.nn as nn from torch.export import ExportedProgram -from torch_tensorrt.hf.exporters import ops as _ops # noqa: F401 -from torch_tensorrt.hf.exporters.compile import compile_component -from torch_tensorrt.hf.exporters.config import EdgeConfig -from torch_tensorrt.hf.exporters.runtime import EdgeRuntimeModule -from torch_tensorrt.hf.exporters.spec import get_edge_spec from transformers.exporters.exporter_dynamo import DynamoExporter +from . import ops as _ops # noqa: F401 +from .compile import compile_component +from .config import EdgeConfig +from .runtime import EdgeRuntimeModule +from .spec import get_edge_spec + logger = logging.getLogger(__name__) @@ -42,6 +43,7 @@ def __init__(self) -> None: 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( @@ -71,6 +73,7 @@ def export( engines: dict[str, str] = {} upstream: dict[str, Any] = {} + self.bench = {} def _compile_components() -> None: for name in names: @@ -81,6 +84,7 @@ def _compile_components() -> None: engine_dir=engine_dir, dryrun=config.dryrun, trt_settings=config.trt_settings, + bench=self.bench, ) upstream.update(spec.capture_upstream(name, outs, sample, bundle)) diff --git a/tools/hf/exporters/measure.py b/tools/hf/exporters/measure.py new file mode 100644 index 00000000000..cdd43b3f878 --- /dev/null +++ b/tools/hf/exporters/measure.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping + +import torch + + +def parity(name: str, eager: torch.Tensor, trt: torch.Tensor) -> None: + """Print eager-vs-TRT numeric agreement for one tensor.""" + a, b = eager.float(), trt.float() + delta = a - b + diff = delta.abs() + rel_l2 = delta.norm() / b.norm().clamp_min(1e-8) + close = torch.isclose(a, b, rtol=1e-2, atol=1e-2).float().mean() * 100 + print( + f"{name:<36} mean_abs={float(diff.mean()):.6f} " + f"max_abs={float(diff.max()):.6f} rel_l2={float(rel_l2):.4f} " + f"close%={float(close):.1f}" + ) + + +def cuda_ms(fn: Callable[[], object], *, warmup: int = 10, iters: int = 100) -> float: + """Average runtime of ``fn`` in milliseconds (CUDA events, else wall time).""" + with torch.no_grad(): + for _ in range(warmup): + fn() + if torch.cuda.is_available(): + torch.cuda.synchronize() + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + t0 = time.perf_counter() + for _ in range(iters): + fn() + return (time.perf_counter() - t0) * 1000.0 / iters + + +def speedup(eager_ms: float, trt_ms: float) -> str: + if eager_ms <= 0.0 or trt_ms <= 0.0: + return "n/a" + return f"{eager_ms / trt_ms:.3f}x" + + +def print_bench(bench: Mapping[str, tuple[float, float]]) -> None: + """Print per-component CUDA timings collected during ``export()``.""" + if not bench: + return + eager_total = trt_total = 0.0 + for name, (eager_ms, trt_ms) in bench.items(): + print(f"{name} eager execute: {eager_ms:.3f} ms") + print(f"{name} trt execute: {trt_ms:.3f} ms") + print(f"{name} speedup: {speedup(eager_ms, trt_ms)}") + eager_total += eager_ms + trt_total += trt_ms + print(f"total eager execute: {eager_total:.3f} ms") + print(f"total trt execute: {trt_total:.3f} ms") + print(f"total speedup: {speedup(eager_total, trt_total)}") diff --git a/tools/hf/exporters/models/__init__.py b/tools/hf/exporters/models/__init__.py new file mode 100644 index 00000000000..eaebd1daeb8 --- /dev/null +++ b/tools/hf/exporters/models/__init__.py @@ -0,0 +1,5 @@ +"""Model families. Importing this package registers EdgeSpecs.""" + +from .groot import spec as _groot # noqa: F401 +from .nemotron import spec as _nemotron # noqa: F401 +from .pi05 import spec as _pi05 # noqa: F401 diff --git a/py/torch_tensorrt/hf/exporters/models/common/__init__.py b/tools/hf/exporters/models/common/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/common/__init__.py rename to tools/hf/exporters/models/common/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/models/common/helpers.py b/tools/hf/exporters/models/common/helpers.py similarity index 96% rename from py/torch_tensorrt/hf/exporters/models/common/helpers.py rename to tools/hf/exporters/models/common/helpers.py index a50c6bf0a6b..5f4951e07c7 100644 --- a/py/torch_tensorrt/hf/exporters/models/common/helpers.py +++ b/tools/hf/exporters/models/common/helpers.py @@ -25,7 +25,7 @@ def causal_lm_flat( num_kv = int(cfg.num_key_value_heads) head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) try: - from torch_tensorrt.hf.exporters.rope import make_rope_rotary_cos_sin + from ...rope import make_rope_rotary_cos_sin rope = make_rope_rotary_cos_sin( cfg, int(max_seq_len), device, language_model=language diff --git a/py/torch_tensorrt/hf/exporters/models/common/patches.py b/tools/hf/exporters/models/common/patches.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/common/patches.py rename to tools/hf/exporters/models/common/patches.py diff --git a/py/torch_tensorrt/hf/exporters/models/groot/__init__.py b/tools/hf/exporters/models/groot/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/groot/__init__.py rename to tools/hf/exporters/models/groot/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/models/groot/helpers.py b/tools/hf/exporters/models/groot/helpers.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/groot/helpers.py rename to tools/hf/exporters/models/groot/helpers.py diff --git a/py/torch_tensorrt/hf/exporters/models/groot/patches.py b/tools/hf/exporters/models/groot/patches.py similarity index 85% rename from py/torch_tensorrt/hf/exporters/models/groot/patches.py rename to tools/hf/exporters/models/groot/patches.py index b0eb154d095..4dbc9feb1e2 100644 --- a/py/torch_tensorrt/hf/exporters/models/groot/patches.py +++ b/tools/hf/exporters/models/groot/patches.py @@ -2,15 +2,17 @@ from __future__ import annotations -from typing import Any, Callable +from contextlib import contextmanager +from typing import Any, Callable, Iterator import torch -from torch_tensorrt.hf.exporters.models.common.patches import causal_lm_plugin_forward -from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + +from ...plugin.attn_patches import ( _patch_language_attention, _patch_vision_attention, register_patch, ) +from ..common.patches import causal_lm_plugin_forward GROOT = "groot" @@ -45,6 +47,28 @@ def forward(self, pixel_values, input_ids=None, **kwargs: Any): return forward +@contextmanager +def apply_groot_patches(model: Any | None = None) -> Iterator[None]: + """Family setattr, plus the live Eagle class. + + LeRobot builds Eagle with ``AutoModel.from_config(..., trust_remote_code=True)``, + so the running class is HuggingFace ``transformers_modules`` code, not + ``lerobot.policies.groot.eagle2_hg_model``. The dotted path still covers the + in-tree copy; this patches ``type(eagle_model)`` so vision + ``eagle(pixel_values)`` hits ``extract_feature``. + """ + from ...plugin.attn_patches import apply_patches, patch_attribute + from .helpers import _groot + + with apply_patches(GROOT): + if model is None: + yield + return + eagle_cls = type(_groot(model).backbone.eagle_model) + with patch_attribute(eagle_cls, "forward", _patch_eagle_image_features): + yield + + @register_patch( GROOT, "transformers.models.llama.modeling_llama.LlamaForCausalLM.forward", diff --git a/py/torch_tensorrt/hf/exporters/models/groot/spec.py b/tools/hf/exporters/models/groot/spec.py similarity index 94% rename from py/torch_tensorrt/hf/exporters/models/groot/spec.py rename to tools/hf/exporters/models/groot/spec.py index 59b8a97e139..82a542b57b7 100644 --- a/py/torch_tensorrt/hf/exporters/models/groot/spec.py +++ b/tools/hf/exporters/models/groot/spec.py @@ -5,22 +5,23 @@ import torch import torch.nn as nn -from torch_tensorrt.hf.exporters.models.common.helpers import ( + +from ...ops import call_engine, scatter_image_tokens +from ...spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from ..common.helpers import ( causal_lm_flat, kv_kwargs, split_flat_to_kwargs, ) -from torch_tensorrt.hf.exporters.models.groot.helpers import ( +from .helpers import ( _groot, make_embodiment_id, ) -from torch_tensorrt.hf.exporters.models.groot.patches import GROOT -from torch_tensorrt.hf.exporters.ops import call_engine, scatter_image_tokens -from torch_tensorrt.hf.exporters.spec import ( - ComponentBundle, - EdgeSpec, - register_edge_spec, -) +from .patches import apply_groot_patches def _export_module(module: nn.Module, sample: Mapping[str, Any]) -> nn.Module: @@ -44,10 +45,7 @@ class GrootSpec(EdgeSpec): # type: ignore[misc] components = ("vision", "language", "context_projection", "action") def apply_patches(self, model=None): - del model - from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches - - return apply_patches(GROOT) + return apply_groot_patches(model) def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any @@ -57,7 +55,8 @@ def prepare_sample_inputs( from lerobot.policies.factory import make_pre_post_processors from lerobot.policies.groot.processor_groot import GrootEagleEncodeStep - from torch_tensorrt.hf.exporters.data import ( + + from ...data import ( create_pil_messages, load_test_data, pack_state, @@ -127,7 +126,7 @@ def prepare( upstream: Mapping[str, Any], config: Any, ) -> ComponentBundle: - from torch_tensorrt.hf.exporters.plugin.attention import ( + from ...plugin.attention import ( ContextAttentionMaskType, ) diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py b/tools/hf/exporters/models/nemotron/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/nemotron/__init__.py rename to tools/hf/exporters/models/nemotron/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py b/tools/hf/exporters/models/nemotron/helpers.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/nemotron/helpers.py rename to tools/hf/exporters/models/nemotron/helpers.py diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/mamba_stub.py b/tools/hf/exporters/models/nemotron/mamba_stub.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/nemotron/mamba_stub.py rename to tools/hf/exporters/models/nemotron/mamba_stub.py diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py b/tools/hf/exporters/models/nemotron/patches.py similarity index 91% rename from py/torch_tensorrt/hf/exporters/models/nemotron/patches.py rename to tools/hf/exporters/models/nemotron/patches.py index e9856dc285e..7b772500d9c 100644 --- a/py/torch_tensorrt/hf/exporters/models/nemotron/patches.py +++ b/tools/hf/exporters/models/nemotron/patches.py @@ -5,12 +5,12 @@ from contextlib import contextmanager from typing import Any, Callable, Iterator -from torch_tensorrt.hf.exporters.models.common.patches import gather_last_token_hidden -from torch_tensorrt.hf.exporters.models.nemotron.helpers import _decoder, _kind -from torch_tensorrt.hf.exporters.plugin.attn_patches import ( +from ...plugin.attn_patches import ( apply_patches, register_patch, ) +from ..common.patches import gather_last_token_hidden +from .helpers import _decoder, _kind NEMOTRON = "nemotron" @@ -78,7 +78,7 @@ def forward( @contextmanager def apply_nemotron_patches(model: Any | None = None) -> Iterator[None]: """Class setattr plus mixer plugin wrappers (MoE packing needs the instance).""" - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + from ...plugin.plugin_utils import ( patch_nemotron_mixers, restore_attention, ) diff --git a/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py b/tools/hf/exporters/models/nemotron/spec.py similarity index 91% rename from py/torch_tensorrt/hf/exporters/models/nemotron/spec.py rename to tools/hf/exporters/models/nemotron/spec.py index 413a5b21e85..b4a283f410b 100644 --- a/py/torch_tensorrt/hf/exporters/models/nemotron/spec.py +++ b/tools/hf/exporters/models/nemotron/spec.py @@ -5,24 +5,25 @@ import torch import torch.nn as nn -from torch_tensorrt.hf.exporters.models.common.helpers import ( + +from ...ops import call_engine +from ...spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from ..common.helpers import ( kv_kwargs, split_flat_to_kwargs, ) -from torch_tensorrt.hf.exporters.models.nemotron.helpers import ( +from .helpers import ( _decoder, _kind, allocate_plugin_states, ) -from torch_tensorrt.hf.exporters.models.nemotron.patches import ( +from .patches import ( apply_nemotron_patches, ) -from torch_tensorrt.hf.exporters.ops import call_engine -from torch_tensorrt.hf.exporters.spec import ( - ComponentBundle, - EdgeSpec, - register_edge_spec, -) @register_edge_spec("nemotron_h", "nemotron") @@ -72,7 +73,7 @@ def prepare( upstream: Mapping[str, Any], config: Any, ) -> ComponentBundle: - from torch_tensorrt.hf.exporters.rope import make_rope_rotary_cos_sin + from ...rope import make_rope_rotary_cos_sin del name, upstream embeds = sample["inputs_embeds"] diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/__init__.py b/tools/hf/exporters/models/pi05/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/pi05/__init__.py rename to tools/hf/exporters/models/pi05/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/helpers.py b/tools/hf/exporters/models/pi05/helpers.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/models/pi05/helpers.py rename to tools/hf/exporters/models/pi05/helpers.py diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/patches.py b/tools/hf/exporters/models/pi05/patches.py similarity index 95% rename from py/torch_tensorrt/hf/exporters/models/pi05/patches.py rename to tools/hf/exporters/models/pi05/patches.py index 1ef966509a2..5291c96ebe3 100644 --- a/py/torch_tensorrt/hf/exporters/models/pi05/patches.py +++ b/tools/hf/exporters/models/pi05/patches.py @@ -6,12 +6,13 @@ import torch import torch.nn.functional as F -from torch_tensorrt.hf.exporters.models.common.patches import causal_lm_plugin_forward -from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + +from ...plugin.attn_patches import ( _patch_language_attention, _patch_vision_attention, register_patch, ) +from ..common.patches import causal_lm_plugin_forward PI05 = "pi05" @@ -97,7 +98,8 @@ def forward( **kwargs, ) from lerobot.policies.pi05.modeling_pi05 import create_sinusoidal_pos_embedding - from torch_tensorrt.hf.exporters.prefix_cache import PrefixKVCache + + from ...prefix_cache import PrefixKVCache suffix_embs = self.action_in_proj(x_t) time_emb = create_sinusoidal_pos_embedding( diff --git a/py/torch_tensorrt/hf/exporters/models/pi05/spec.py b/tools/hf/exporters/models/pi05/spec.py similarity index 93% rename from py/torch_tensorrt/hf/exporters/models/pi05/spec.py rename to tools/hf/exporters/models/pi05/spec.py index 7fb9381f449..d53dd05dcc7 100644 --- a/py/torch_tensorrt/hf/exporters/models/pi05/spec.py +++ b/tools/hf/exporters/models/pi05/spec.py @@ -5,25 +5,26 @@ import torch import torch.nn as nn -from torch_tensorrt.hf.exporters.models.common.helpers import ( + +from ...ops import call_engine, fuse_prefix +from ...spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from ..common.helpers import ( causal_lm_flat, kv_kwargs, split_flat_to_kwargs, ) -from torch_tensorrt.hf.exporters.models.common.patches import language_decoder -from torch_tensorrt.hf.exporters.models.pi05.helpers import ( +from ..common.patches import language_decoder +from .helpers import ( _core, build_pi05_prefix_embs, make_pi05_suffix_position_and_mask, pi05_compact_index, ) -from torch_tensorrt.hf.exporters.models.pi05.patches import PI05 -from torch_tensorrt.hf.exporters.ops import call_engine, fuse_prefix -from torch_tensorrt.hf.exporters.spec import ( - ComponentBundle, - EdgeSpec, - register_edge_spec, -) +from .patches import PI05 @register_edge_spec("pi05") @@ -33,7 +34,7 @@ class Pi05Spec(EdgeSpec): # type: ignore[misc] def apply_patches(self, model=None): """Install vision, language, and action setattr replacements.""" del model - from torch_tensorrt.hf.exporters.plugin.attn_patches import apply_patches + from ...plugin.attn_patches import apply_patches return apply_patches(PI05) @@ -48,7 +49,8 @@ def prepare_sample_inputs( OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, ) - from torch_tensorrt.hf.exporters.data import ( + + from ...data import ( frame_from_test_data, load_test_data, ) @@ -96,7 +98,7 @@ def prepare( upstream: Mapping[str, Any], config: Any, ) -> ComponentBundle: - from torch_tensorrt.hf.exporters.plugin.attention import ( + from ...plugin.attention import ( ContextAttentionMaskType, ) diff --git a/py/torch_tensorrt/hf/exporters/ops.py b/tools/hf/exporters/ops.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/ops.py rename to tools/hf/exporters/ops.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/__init__.py b/tools/hf/exporters/plugin/__init__.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/__init__.py rename to tools/hf/exporters/plugin/__init__.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/attention.py b/tools/hf/exporters/plugin/attention.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/attention.py rename to tools/hf/exporters/plugin/attention.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/attn_patches.py b/tools/hf/exporters/plugin/attn_patches.py similarity index 99% rename from py/torch_tensorrt/hf/exporters/plugin/attn_patches.py rename to tools/hf/exporters/plugin/attn_patches.py index cea2ea5b6c5..4b30576132b 100644 --- a/py/torch_tensorrt/hf/exporters/plugin/attn_patches.py +++ b/tools/hf/exporters/plugin/attn_patches.py @@ -20,7 +20,8 @@ import torch import torch.nn as nn -from torch_tensorrt.hf.exporters.plugin.attention import ContextAttentionMaskType + +from .attention import ContextAttentionMaskType _PATCHES: dict[str, list[tuple[str, Callable]]] = {} _LANGUAGE_MASK_TYPE = int(ContextAttentionMaskType.PADDING) diff --git a/py/torch_tensorrt/hf/exporters/plugin/mamba.py b/tools/hf/exporters/plugin/mamba.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/mamba.py rename to tools/hf/exporters/plugin/mamba.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/moe.py b/tools/hf/exporters/plugin/moe.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/moe.py rename to tools/hf/exporters/plugin/moe.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py b/tools/hf/exporters/plugin/plugin_converter.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/plugin_converter.py rename to tools/hf/exporters/plugin/plugin_converter.py diff --git a/py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py b/tools/hf/exporters/plugin/plugin_utils.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/plugin/plugin_utils.py rename to tools/hf/exporters/plugin/plugin_utils.py diff --git a/py/torch_tensorrt/hf/exporters/prefix_cache.py b/tools/hf/exporters/prefix_cache.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/prefix_cache.py rename to tools/hf/exporters/prefix_cache.py diff --git a/py/torch_tensorrt/hf/exporters/rope.py b/tools/hf/exporters/rope.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/rope.py rename to tools/hf/exporters/rope.py diff --git a/py/torch_tensorrt/hf/exporters/runtime.py b/tools/hf/exporters/runtime.py similarity index 90% rename from py/torch_tensorrt/hf/exporters/runtime.py rename to tools/hf/exporters/runtime.py index 8e98947140a..e55547ea222 100644 --- a/py/torch_tensorrt/hf/exporters/runtime.py +++ b/tools/hf/exporters/runtime.py @@ -4,7 +4,8 @@ from typing import Any import torch.nn as nn -from torch_tensorrt.hf.exporters.spec import EdgeSpec + +from .spec import EdgeSpec class EdgeRuntimeModule(nn.Module): # type: ignore[misc] diff --git a/py/torch_tensorrt/hf/exporters/spec.py b/tools/hf/exporters/spec.py similarity index 100% rename from py/torch_tensorrt/hf/exporters/spec.py rename to tools/hf/exporters/spec.py diff --git a/tests/py/dynamo/hf/test_edge_exporter.py b/tools/hf/exporters/tests/test_edge_exporter.py similarity index 82% rename from tests/py/dynamo/hf/test_edge_exporter.py rename to tools/hf/exporters/tests/test_edge_exporter.py index 43bd6fec71f..23437408fd9 100644 --- a/tests/py/dynamo/hf/test_edge_exporter.py +++ b/tools/hf/exporters/tests/test_edge_exporter.py @@ -6,9 +6,9 @@ import pytest import torch import torch.nn as nn -from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter, register_edge_spec -from torch_tensorrt.hf.exporters.ops import call_engine -from torch_tensorrt.hf.exporters.spec import ComponentBundle, EdgeSpec, registered_specs +from exporters import EdgeConfig, EdgeExporter, register_edge_spec +from exporters.ops import call_engine +from exporters.spec import ComponentBundle, EdgeSpec, registered_specs @register_edge_spec("dummy_edge") @@ -177,7 +177,7 @@ def test_edge_exporter_dryrun_keeps_attention_patch(tmp_path): @pytest.mark.unit def test_attn_patch_attribute_restores(): - from torch_tensorrt.hf.exporters.plugin.attn_patches import patch_attribute + from exporters.plugin.attn_patches import patch_attribute class Owner: def go(self): @@ -196,7 +196,7 @@ def go(self): @pytest.mark.unit def test_language_attn_keeps_hf_forward_without_rope(): - from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + from exporters.plugin.attn_patches import ( _patch_language_attention, ) @@ -214,8 +214,8 @@ def forward(self, hidden_states, past_key_values=None, **kwargs): @pytest.mark.unit def test_pi05_backend_registers_vision_and_language(): - from torch_tensorrt.hf.exporters.models.pi05.patches import PI05 - from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + from exporters.models.pi05.patches import PI05 + from exporters.plugin.attn_patches import _PATCHES paths = [p for p, _ in _PATCHES[PI05]] assert any("SiglipAttention.forward" in p for p in paths) @@ -227,7 +227,7 @@ def test_pi05_backend_registers_vision_and_language(): @pytest.mark.unit def test_paligemma_image_features_patch_returns_tensor(): - from torch_tensorrt.hf.exporters.models.pi05.patches import ( + from exporters.models.pi05.patches import ( _patch_paligemma_image_features, ) @@ -262,7 +262,7 @@ def forward(self, *args, **kwargs): @pytest.mark.unit def test_pi05_language_model_keeps_hf_forward_without_rope(): - from torch_tensorrt.hf.exporters.models.pi05.patches import ( + from exporters.models.pi05.patches import ( _patch_pi05_language_model, ) @@ -279,7 +279,7 @@ def forward(self, inputs_embeds=None, past_key_values=None, **kwargs): @pytest.mark.unit def test_pi05_action_keeps_training_forward_without_prefix_kv(): - from torch_tensorrt.hf.exporters.models.pi05.patches import ( + from exporters.models.pi05.patches import ( _patch_pi05_action_step_forward, ) @@ -294,10 +294,10 @@ def forward(self, images, img_masks, tokens, masks, actions, noise, time): @pytest.mark.unit def test_language_attn_plugin_when_rope_present(): - from torch_tensorrt.hf.exporters.plugin.attn_patches import ( + from exporters.plugin.attn_patches import ( _patch_language_attention, ) - from torch_tensorrt.hf.exporters.plugin.plugin_utils import ( + from exporters.plugin.plugin_utils import ( _register_attention_plugin_op, ) @@ -336,8 +336,8 @@ def forward(self, hidden_states, **kwargs): @pytest.mark.unit def test_groot_backend_registers_components(): - from torch_tensorrt.hf.exporters.models.groot.patches import GROOT - from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + from exporters.models.groot.patches import GROOT + from exporters.plugin.attn_patches import _PATCHES paths = [p for p, _ in _PATCHES[GROOT]] assert any("SiglipAttention.forward" in p for p in paths) @@ -351,8 +351,8 @@ def test_groot_backend_registers_components(): @pytest.mark.unit def test_nemotron_backend_registers_causal_lm(): - from torch_tensorrt.hf.exporters.models.nemotron.patches import NEMOTRON - from torch_tensorrt.hf.exporters.plugin.attn_patches import _PATCHES + from exporters.models.nemotron.patches import NEMOTRON + from exporters.plugin.attn_patches import _PATCHES paths = [p for p, _ in _PATCHES[NEMOTRON]] assert any("NemotronHForCausalLM.forward" in p for p in paths) @@ -360,7 +360,7 @@ def test_nemotron_backend_registers_causal_lm(): @pytest.mark.unit def test_eagle_vision_patch_extracts_features(): - from torch_tensorrt.hf.exporters.models.groot.patches import ( + from exporters.models.groot.patches import ( _patch_eagle_image_features, ) @@ -376,9 +376,36 @@ def forward(self, *args, **kwargs): torch.testing.assert_close(Dummy()(pixel_values), pixel_values + 1) +@pytest.mark.unit +def test_groot_patches_live_eagle_class(): + from exporters.models.groot.patches import apply_groot_patches + + class Eagle: + def extract_feature(self, pixel_values): + return pixel_values + 1 + + def forward(self, pixel_values, input_ids=None, **kwargs): + raise AssertionError("unpatched Eagle.forward should not run") + + class Groot: + def __init__(self): + self.backbone = type("Backbone", (), {})() + self.backbone.eagle_model = Eagle() + + class Policy: + def __init__(self): + self._groot_model = Groot() + + policy = Policy() + eagle = policy._groot_model.backbone.eagle_model + pixel_values = torch.zeros(1, 3, 4, 4) + with apply_groot_patches(policy): + torch.testing.assert_close(eagle(pixel_values), pixel_values + 1) + + @pytest.mark.unit def test_eagle_vision_keeps_vlm_forward_with_input_ids(): - from torch_tensorrt.hf.exporters.models.groot.patches import ( + from exporters.models.groot.patches import ( _patch_eagle_image_features, ) @@ -399,7 +426,7 @@ def forward(self, pixel_values, input_ids=None, **kwargs): @pytest.mark.unit def test_groot_action_keeps_training_forward_without_context(): - from torch_tensorrt.hf.exporters.models.groot.patches import ( + from exporters.models.groot.patches import ( _patch_groot_action_step_forward, ) @@ -413,7 +440,7 @@ def forward(self, backbone_output, action_input): @pytest.mark.unit def test_groot_context_keeps_training_forward_without_hidden(): - from torch_tensorrt.hf.exporters.models.groot.patches import ( + from exporters.models.groot.patches import ( _patch_groot_context_projection, ) @@ -427,7 +454,7 @@ def forward(self, backbone_inputs, action_inputs): @pytest.mark.unit def test_nemotron_keeps_hf_forward_without_rope(): - from torch_tensorrt.hf.exporters.models.nemotron.patches import ( + from exporters.models.nemotron.patches import ( _patch_nemotron_causal_lm, ) @@ -444,7 +471,7 @@ def forward(self, input_ids=None, inputs_embeds=None, **kwargs): @pytest.mark.unit def test_category_specific_linear_uses_index_select(): - from torch_tensorrt.hf.exporters.models.groot.patches import ( + from exporters.models.groot.patches import ( _patch_category_specific_linear, ) @@ -470,3 +497,27 @@ def forward(self, x, cat_ids): out = layer(x, cat_ids) expected = torch.bmm(x, layer.W[cat_ids]) + layer.b[cat_ids].unsqueeze(1) torch.testing.assert_close(out, expected) + + +@pytest.mark.unit +def test_measure_parity_and_bench(capsys): + from exporters.measure import cuda_ms, parity, print_bench, speedup + + a = torch.ones(2, 2) + parity("dummy A vs C (TRT)", a, a) + log = capsys.readouterr().out + assert "dummy A vs C (TRT)" in log + assert "close%=100.0" in log + assert speedup(10.0, 5.0) == "2.000x" + assert speedup(0.0, 5.0) == "n/a" + + elapsed = cuda_ms(lambda: torch.ones(2, 2).sum(), warmup=1, iters=3) + assert elapsed >= 0.0 + + print_bench({"vision": (10.0, 5.0), "language": (4.0, 2.0)}) + log = capsys.readouterr().out + assert "vision eager execute: 10.000 ms" in log + assert "vision trt execute: 5.000 ms" in log + assert "total speedup: 2.000x" in log + print_bench({}) + assert capsys.readouterr().out == "" diff --git a/py/torch_tensorrt/hf/exporters/utils.py b/tools/hf/exporters/utils.py similarity index 75% rename from py/torch_tensorrt/hf/exporters/utils.py rename to tools/hf/exporters/utils.py index fb6c04b1779..b8d21d3b40b 100644 --- a/py/torch_tensorrt/hf/exporters/utils.py +++ b/tools/hf/exporters/utils.py @@ -8,17 +8,6 @@ import torch -_THOR_CUDA_LIB = Path("/usr/local/cuda-13.0/thor/targets/aarch64-linux/lib") - - -def configure_thor_pytorch() -> None: - """Use PyTorch fallbacks for ops whose pip CUDA wheels mismatch DriveOS Thor.""" - on_thor = os.environ.get("TRT_VLA_THOR", "auto") - if on_thor == "auto": - on_thor = "1" if _THOR_CUDA_LIB.is_dir() else "0" - if on_thor == "1": - torch.backends.cudnn.enabled = False - def force_hf_attention(module: Any, attn: str, use_cache: bool | None = False) -> None: """Force HuggingFace attention implementation on a module tree.""" diff --git a/examples/dynamo/run_groot_export.py b/tools/hf/run_groot_export.py similarity index 81% rename from examples/dynamo/run_groot_export.py rename to tools/hf/run_groot_export.py index 19baf8558f9..bf918c9add6 100644 --- a/examples/dynamo/run_groot_export.py +++ b/tools/hf/run_groot_export.py @@ -8,25 +8,19 @@ from __future__ import annotations import argparse +import sys from pathlib import Path -_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ -_TRT_PY = _REPO_ROOT / "py" - import torch # noqa: E402 import torch_tensorrt # noqa: E402 - -_src_pkg = str(_TRT_PY / "torch_tensorrt") -if _src_pkg not in list(torch_tensorrt.__path__): - torch_tensorrt.__path__.append(_src_pkg) - +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt +from exporters.utils import force_hf_attention from lerobot.configs import FeatureType, PolicyFeature from lerobot.policies.groot import GrootPolicy from lerobot.policies.groot.configuration_groot import GrootConfig from lerobot.utils.constants import ACTION, OBS_STATE -from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt -from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch, force_hf_attention def load_groot(device: torch.device) -> GrootPolicy: @@ -62,7 +56,6 @@ def main() -> None: parser.add_argument("--engine-dir", default="/tmp/groot_edge_exporter") args = parser.parse_args() - configure_thor_pytorch() load_plugins_for_trt() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -75,13 +68,7 @@ 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, - dryrun=not args.compile, - skip_runtime_export=False, - ) + config = EdgeConfig(model_type="groot", engine_dir=args.engine_dir, max_seq_len=968) # Spec tokenizes libero via Eagle chat template because we pass the policy. sample_inputs = {"device": device, "dtype": dtype} @@ -98,6 +85,7 @@ def main() -> None: out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity print("velocity", tuple(out.shape), "mean", float(out.float().mean())) + print_bench(exporter.bench) if __name__ == "__main__": diff --git a/examples/dynamo/run_nemotron_export.py b/tools/hf/run_nemotron_export.py similarity index 80% rename from examples/dynamo/run_nemotron_export.py rename to tools/hf/run_nemotron_export.py index 2df1af5139f..977cfe0056b 100644 --- a/examples/dynamo/run_nemotron_export.py +++ b/tools/hf/run_nemotron_export.py @@ -8,24 +8,15 @@ from __future__ import annotations import argparse +import sys from pathlib import Path -_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ -_TRT_PY = _REPO_ROOT / "py" - -import torch # noqa: E402 -import torch_tensorrt # noqa: E402 - -_src_pkg = str(_TRT_PY / "torch_tensorrt") -if _src_pkg not in list(torch_tensorrt.__path__): - torch_tensorrt.__path__.append(_src_pkg) - -from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from torch_tensorrt.hf.exporters.models.nemotron.mamba_stub import ( - apply as apply_mamba_stub, -) -from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt -from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch +import torch +import torch_tensorrt +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.models.nemotron.mamba_stub import apply as apply_mamba_stub +from exporters.plugin.plugin_utils import load_plugins_for_trt from transformers import AutoModelForCausalLM, AutoTokenizer @@ -60,7 +51,6 @@ def main() -> None: parser.add_argument("--max-seq-len", type=int, default=128) args = parser.parse_args() - configure_thor_pytorch() load_plugins_for_trt() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -94,6 +84,7 @@ def main() -> None: logits = out[0] if isinstance(out, (tuple, list)) else out print("logits", tuple(logits.shape), "mean", float(logits.float().mean())) + print_bench(exporter.bench) if __name__ == "__main__": diff --git a/examples/dynamo/run_pi05_export.py b/tools/hf/run_pi05_export.py similarity index 79% rename from examples/dynamo/run_pi05_export.py rename to tools/hf/run_pi05_export.py index fc84c8ce07f..2c2017c0e06 100644 --- a/examples/dynamo/run_pi05_export.py +++ b/tools/hf/run_pi05_export.py @@ -1,25 +1,25 @@ #!/usr/bin/env python3 +"""Smoke EdgeExporter on pi05 model. + +Pass the LeRobot PI05Policy, not policy.model — prepare_sample_inputs +needs the preprocessor on the policy wrapper. +""" + from __future__ import annotations import argparse +import sys from pathlib import Path -_REPO_ROOT = Path(__file__).resolve().parents[2] # TensorRT/ -_TRT_PY = _REPO_ROOT / "py" - -import torch # noqa: E402 -import torch_tensorrt # noqa: E402 - -_src_pkg = str(_TRT_PY / "torch_tensorrt") -if _src_pkg not in list(torch_tensorrt.__path__): - torch_tensorrt.__path__.append(_src_pkg) - +import torch +import torch_tensorrt +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt +from exporters.utils import force_hf_attention from lerobot.configs import FeatureType, PolicyFeature from lerobot.policies.pi05 import PI05Policy from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE -from torch_tensorrt.hf.exporters import EdgeConfig, EdgeExporter -from torch_tensorrt.hf.exporters.plugin.plugin_utils import load_plugins_for_trt -from torch_tensorrt.hf.exporters.utils import configure_thor_pytorch, force_hf_attention def load_pi05(device: torch.device) -> PI05Policy: @@ -59,7 +59,6 @@ def main() -> None: parser.add_argument("--engine-dir", default="/tmp/pi05_edge_exporter") args = parser.parse_args() - configure_thor_pytorch() load_plugins_for_trt() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -78,9 +77,6 @@ def main() -> None: model_type="pi05", # optional; inferred from paligemma_with_expert engine_dir=args.engine_dir, max_seq_len=968, - dryrun=not args.compile, # True = no TRT, still writes config.json + runtime graph - skip_runtime_export=False, # False = also torch.export the stitched execute_engine graph - # components=("vision",), # uncomment to export only vision ) # Spec loads libero + preprocessor because we pass the policy, not a tensor dict. @@ -102,6 +98,7 @@ def main() -> None: out = velocity[0] if isinstance(velocity, (tuple, list)) else velocity print("velocity", tuple(out.shape), "mean", float(out.float().mean())) + print_bench(exporter.bench) if __name__ == "__main__": From c99e14061030ec5b46f4bf3c0f6517151f28faa7 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 8 Sep 2026 00:06:46 -0700 Subject: [PATCH 08/11] 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 7263ebeb39a..74009efe52e 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 81777e5c66c..8030f3319a3 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 d81055589b8..3a822da16be 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 f93ca57bc08..a46b5ad9645 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 5f4951e07c7..2eb153e7bfd 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 c9434f0a278..c2d916325ca 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 82a542b57b7..c019e5851bf 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 7b772500d9c..568a5f436c7 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 b4a283f410b..2f41a99dc7a 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 5a1608837c6..0808afac775 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 d53dd05dcc7..c5f0bb74b69 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 9755d451aff..a71f791827f 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 4b30576132b..2693dbb853e 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 0186cdf2600..3dbea351df1 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 23437408fd9..590ae602b29 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 bf918c9add6..d21007e8250 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 977cfe0056b..14da09de3e8 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 2c2017c0e06..f7abd318bc1 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())) From 42f0216032fe5c7b09c2ecf0cbb9b010878e2777 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Thu, 10 Sep 2026 14:26:40 -0700 Subject: [PATCH 09/11] Add Nanbeige4.2-3B export support --- tools/hf/exporters/__init__.py | 1 + tools/hf/exporters/models/groot/patches.py | 22 -- tools/hf/exporters/models/groot/spec.py | 12 +- .../hf/exporters/models/nanbeige/__init__.py | 1 + tools/hf/exporters/models/nanbeige/helpers.py | 53 ++++ tools/hf/exporters/models/nanbeige/patches.py | 101 ++++++++ tools/hf/exporters/models/nanbeige/spec.py | 239 ++++++++++++++++++ tools/hf/run_nanbeige_export.py | 71 ++++++ tools/hf/run_pi05_export.py | 4 +- 9 files changed, 478 insertions(+), 26 deletions(-) create mode 100644 tools/hf/exporters/models/nanbeige/__init__.py create mode 100644 tools/hf/exporters/models/nanbeige/helpers.py create mode 100644 tools/hf/exporters/models/nanbeige/patches.py create mode 100644 tools/hf/exporters/models/nanbeige/spec.py create mode 100644 tools/hf/run_nanbeige_export.py diff --git a/tools/hf/exporters/__init__.py b/tools/hf/exporters/__init__.py index 2ed7d30ed14..1e2286ba78a 100644 --- a/tools/hf/exporters/__init__.py +++ b/tools/hf/exporters/__init__.py @@ -1,6 +1,7 @@ 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 ( diff --git a/tools/hf/exporters/models/groot/patches.py b/tools/hf/exporters/models/groot/patches.py index 4dbc9feb1e2..f56f1c4f71f 100644 --- a/tools/hf/exporters/models/groot/patches.py +++ b/tools/hf/exporters/models/groot/patches.py @@ -47,28 +47,6 @@ def forward(self, pixel_values, input_ids=None, **kwargs: Any): return forward -@contextmanager -def apply_groot_patches(model: Any | None = None) -> Iterator[None]: - """Family setattr, plus the live Eagle class. - - LeRobot builds Eagle with ``AutoModel.from_config(..., trust_remote_code=True)``, - so the running class is HuggingFace ``transformers_modules`` code, not - ``lerobot.policies.groot.eagle2_hg_model``. The dotted path still covers the - in-tree copy; this patches ``type(eagle_model)`` so vision - ``eagle(pixel_values)`` hits ``extract_feature``. - """ - from ...plugin.attn_patches import apply_patches, patch_attribute - from .helpers import _groot - - with apply_patches(GROOT): - if model is None: - yield - return - eagle_cls = type(_groot(model).backbone.eagle_model) - with patch_attribute(eagle_cls, "forward", _patch_eagle_image_features): - yield - - @register_patch( GROOT, "transformers.models.llama.modeling_llama.LlamaForCausalLM.forward", diff --git a/tools/hf/exporters/models/groot/spec.py b/tools/hf/exporters/models/groot/spec.py index c019e5851bf..33f7599a34c 100644 --- a/tools/hf/exporters/models/groot/spec.py +++ b/tools/hf/exporters/models/groot/spec.py @@ -21,7 +21,6 @@ _groot, make_embodiment_id, ) -from .patches import apply_groot_patches def _export_module(module: nn.Module, sample: Mapping[str, Any]) -> nn.Module: @@ -43,7 +42,16 @@ def _causal_lm(language: nn.Module) -> nn.Module: @register_edge_spec("groot", "gr00t") class GrootSpec(EdgeSpec): # type: ignore[misc] def apply_patches(self, model=None): - return apply_groot_patches(model) + from ...plugin.attn_patches import apply_patches, patch_attribute + from .helpers import _groot + + with apply_patches(GROOT): + if model is None: + yield + return + eagle_cls = type(_groot(model).backbone.eagle_model) + with patch_attribute(eagle_cls, "forward", _patch_eagle_image_features): + yield def prepare_sample_inputs( self, model: nn.Module, raw: Mapping[str, Any], config: Any diff --git a/tools/hf/exporters/models/nanbeige/__init__.py b/tools/hf/exporters/models/nanbeige/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tools/hf/exporters/models/nanbeige/__init__.py @@ -0,0 +1 @@ + diff --git a/tools/hf/exporters/models/nanbeige/helpers.py b/tools/hf/exporters/models/nanbeige/helpers.py new file mode 100644 index 00000000000..42237655e1c --- /dev/null +++ b/tools/hf/exporters/models/nanbeige/helpers.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import logging + +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +@torch.no_grad() +def ensure_valid_rope_inv_freq(model: nn.Module) -> None: + """Repair deterministic RoPE buffers broken by remote-code version drift.""" + decoder = model if hasattr(model, "layers") else model.model + repaired = 0 + + for layer in decoder.layers: + attention = layer.self_attn + rotary = attention.rotary_emb + device = attention.q_proj.weight.device + dim = int(rotary.dim) + base = float(rotary.base) + exponent = torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim + expected = torch.pow( + torch.tensor(base, device=device, dtype=torch.float32), + -exponent, + ) + + current = getattr(rotary, "inv_freq", None) + valid = ( + isinstance(current, torch.Tensor) + and current.device.type != "meta" + and current.shape == expected.shape + and bool(torch.isfinite(current).all()) + and torch.allclose( + current.to(device=device, dtype=torch.float32), + expected, + ) + ) + if valid: + continue + + if current is None: + rotary.register_buffer("inv_freq", expected, persistent=False) + else: + rotary.inv_freq = expected + repaired += 1 + + if repaired: + logger.warning( + "Reinitialized invalid Nanbeige RoPE buffers in %d attention layers", + repaired, + ) diff --git a/tools/hf/exporters/models/nanbeige/patches.py b/tools/hf/exporters/models/nanbeige/patches.py new file mode 100644 index 00000000000..0011477e0df --- /dev/null +++ b/tools/hf/exporters/models/nanbeige/patches.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, Callable, Iterator + +import torch + +from ...plugin.attn_patches import _patch_language_attention, patch_attribute + + +def _num_loops(config: Any) -> int: + loop_weights = getattr(config, "loop_loss_weights", None) + if loop_weights: + return len(loop_weights) + 1 + return max(int(getattr(config, "num_loops", 1)), 1) + + +def _patch_nanbeige_attention(original: Callable) -> Callable: + """Use the Edge language-attention plugin for the loaded remote-code class.""" + return _patch_language_attention(original) + + +def _patch_nanbeige_language_model(original: Callable) -> Callable: + """Edge prefill when rope is present; otherwise HF forward.""" + + def forward( + self, + inputs_embeds=None, + rope_rotary_cos_sin=None, + context_lengths=None, + kvcache_start_index=None, + last_token_ids=None, + *past_key_values, + **kwargs: Any, + ): + lm = self if hasattr(self, "layers") else self.model + lm_head = getattr(self, "lm_head", None) + + hidden = inputs_embeds.to(dtype=next(lm.parameters()).dtype) + seq_len = inputs_embeds.shape[1] + physical = len(lm.layers) + num_loops = _num_loops(lm.config) + skip_loop_norm = bool(getattr(lm.config, "skip_loop_final_norm", False)) + new_kvs = [] + + for loop_idx in range(num_loops): + for layer_idx, layer in enumerate(lm.layers): + logical = layer_idx + loop_idx * physical + residual = hidden + hidden = layer.input_layernorm(hidden) + hidden, kv = layer.self_attn( + hidden_states=hidden, + rope_rotary_cos_sin=rope_rotary_cos_sin, + past_key_value=past_key_values[logical], + ctx_len=context_lengths, + kvcache_start_index=kvcache_start_index, + ) + hidden = residual + hidden + residual = hidden + hidden = layer.post_attention_layernorm(hidden) + hidden = layer.mlp(hidden) + hidden = residual + hidden + new_kvs.append(kv) + + if not skip_loop_norm or loop_idx == num_loops - 1: + hidden = lm.norm(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, + ] + logits = lm_head(last_hidden).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, hidden, prefix_k, prefix_v + + return forward + + +@contextmanager +def apply_nanbeige_patches(model: Any | None = None) -> Iterator[None]: + """Patch the live trust-remote-code classes only while compiling.""" + # Nanbeige is loaded with trust_remote_code, so its revision-specific module + # path is not stable enough for @register_patch; patch the loaded classes. + if model is None: + yield + return + if not hasattr(model, "lm_head"): + raise TypeError("Nanbeige export expects NanbeigeForCausalLM") + + decoder = model if hasattr(model, "layers") else model.model + if not getattr(decoder, "layers", None): + raise AttributeError("Nanbeige decoder has no layers") + + attn_cls = type(decoder.layers[0].self_attn) + with patch_attribute(attn_cls, "forward", _patch_nanbeige_attention): + with patch_attribute(type(model), "forward", _patch_nanbeige_language_model): + yield diff --git a/tools/hf/exporters/models/nanbeige/spec.py b/tools/hf/exporters/models/nanbeige/spec.py new file mode 100644 index 00000000000..729caab8603 --- /dev/null +++ b/tools/hf/exporters/models/nanbeige/spec.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn + +from ...ops import call_engine +from ...spec import ( + ComponentBundle, + EdgeSpec, + register_edge_spec, +) +from ..common.helpers import kv_kwargs, split_flat_to_kwargs +from .helpers import ensure_valid_rope_inv_freq +from .patches import apply_nanbeige_patches + + +@register_edge_spec("nanbeige") +class NanbeigeSpec(EdgeSpec): # type: ignore[misc] + def apply_patches(self, model=None): + return apply_nanbeige_patches(model) + + def prepare_sample_inputs( + self, model: nn.Module, raw: Mapping[str, Any], config: Any + ) -> MutableMapping[str, Any]: + del config + ensure_valid_rope_inv_freq(model) + if "inputs_embeds" in raw: + sample = dict(raw) + mask = sample.get("attention_mask") + if mask is not None: + sample["attention_mask"] = mask.to( + device=sample["inputs_embeds"].device + ) + return sample + if "input_ids" not in raw: + raise KeyError("Nanbeige inputs require input_ids or inputs_embeds") + + embedding = model.get_input_embeddings() + input_ids = raw["input_ids"] + with torch.no_grad(): + inputs_embeds = embedding(input_ids.to(device=embedding.weight.device)) + + attention_mask = raw.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(device=inputs_embeds.device) + return { + "inputs_embeds": inputs_embeds.contiguous(), + "attention_mask": attention_mask, + } + + 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"], + "use_cache": False, + "return_dict": True, + } + if sample.get("attention_mask") is not None: + kwargs["attention_mask"] = sample["attention_mask"] + + with torch.no_grad(): + output = model(**kwargs) + logits = output.logits + + # The plugin forward returns logits for one selected token, not [B, S, V]. + mask = sample.get("attention_mask") + if mask is None: + last_token_ids = torch.full( + (logits.shape[0],), + logits.shape[1] - 1, + device=logits.device, + dtype=torch.long, + ) + else: + positions = torch.arange(logits.shape[1], device=logits.device) + last_token_ids = ( + positions.expand_as(mask) + .masked_fill(~mask.bool(), -1) + .max(dim=1) + .values + ) + if bool((last_token_ids < 0).any()): + raise ValueError("attention_mask contains an empty sequence") + + batch = torch.arange(logits.shape[0], device=logits.device) + last_logits = logits[batch, last_token_ids] + + if bench is not None: + bench["language"] = cuda_ms(lambda: model(**kwargs).logits) + + return {"language": last_logits} + + def prepare( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + from ...plugin.attention import ContextAttentionMaskType + from ...rope import make_rope_rotary_cos_sin + + decoder = model if hasattr(model, "layers") else model.model + if getattr(decoder.config, "enable_double_loop_split", False): + raise NotImplementedError("Nanbeige LoopSplit export is not supported") + if getattr(decoder.config, "enable_hyper_connection", False): + raise NotImplementedError( + "Nanbeige hyper-connection export is not supported" + ) + if getattr(decoder.config, "enable_depth_attention", False): + raise NotImplementedError( + "Nanbeige depth-attention export is not supported" + ) + if getattr(decoder, "ngram_embeddings", None) is not None: + raise NotImplementedError("Nanbeige n-gram export is not supported") + + embeds = sample["inputs_embeds"] + device, dtype = embeds.device, embeds.dtype + batch_size, seq_len, hidden_size = embeds.shape + + max_seq_len = max(int(config.max_seq_len), seq_len) + physical_layers = len(decoder.layers) + loop_weights = getattr(decoder.config, "loop_loss_weights", None) + num_loops = ( + len(loop_weights) + 1 + if loop_weights + else max(int(getattr(decoder.config, "num_loops", 1)), 1) + ) + logical_layers = physical_layers * num_loops + + num_kv_heads = int(decoder.config.num_key_value_heads) + head_dim = int( + getattr(decoder.config, "head_dim", None) + or hidden_size // int(decoder.config.num_attention_heads) + ) + rope = make_rope_rotary_cos_sin( + decoder.config, + max_seq_len, + device, + language_model=decoder, + ) + + attention_mask = sample.get("attention_mask") + if attention_mask is None: + context_lengths = torch.full( + (batch_size,), + seq_len, + device=device, + dtype=torch.int32, + ) + last_token_ids = torch.full( + (batch_size, 1), + seq_len - 1, + device=device, + dtype=torch.int64, + ) + else: + context_lengths = attention_mask.sum(dim=-1, dtype=torch.int32) + positions = torch.arange(seq_len, device=device) + last_token_ids = ( + positions.expand_as(attention_mask) + .masked_fill(~attention_mask.bool(), -1) + .max(dim=1) + .values.unsqueeze(-1) + ) + + kvcache_start_index = torch.empty(0, device=device, dtype=torch.int32) + past_key_values = [ + torch.zeros( + batch_size, + 2, + num_kv_heads, + max_seq_len, + head_dim, + device=device, + dtype=dtype, + ) + for _ in range(logical_layers) + ] + flat = ( + embeds, + rope, + context_lengths, + kvcache_start_index, + last_token_ids, + *past_key_values, + ) + input_names = [ + "inputs_embeds", + "rope_rotary_cos_sin", + "context_lengths", + "kvcache_start_index", + "last_token_ids", + *[f"past_key_values_{index}" for index in range(logical_layers)], + ] + sample.update(split_flat_to_kwargs(flat, input_names)) + return { + "language": ComponentBundle( + module=model.eval(), + trace_args=flat, + save_args=flat, + input_names=input_names, + output_names=[ + "logits", + "hidden_states", + "prefix_k", + "prefix_v", + ], + parity_output="logits", + context_attention_mask_type=int(ContextAttentionMaskType.CAUSAL), + model_type="nanbeige", + 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: + return call_engine( + engines["language"], + "language", + sample["inputs_embeds"], + sample["rope_rotary_cos_sin"], + sample["context_lengths"], + sample["kvcache_start_index"], + sample["last_token_ids"], + *kv_kwargs(sample), + ) diff --git a/tools/hf/run_nanbeige_export.py b/tools/hf/run_nanbeige_export.py new file mode 100644 index 00000000000..9bfd5f4d3a6 --- /dev/null +++ b/tools/hf/run_nanbeige_export.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Smoke EdgeExporter on pi05 model. + +Pass the LeRobot PI05Policy, not policy.model — prepare_sample_inputs +needs the preprocessor on the policy wrapper. +""" + +from __future__ import annotations + +import torch +import torch_tensorrt +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt +from exporters.utils import force_hf_attention +from transformers import AutoModelForCausalLM, AutoTokenizer + +checkpoint = "Nanbeige/Nanbeige4.2-3B" + + +def load_policy(device: torch.device, dtype: torch.dtype): + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + torch_dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + # if pad token is not set, set it to eos token + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + return model, tokenizer + + +def main() -> None: + load_plugins_for_trt() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float16 + + model, tokenizer = load_policy(device, dtype) + encoded = tokenizer("Hello, how are you?", return_tensors="pt") + sample_inputs = { + "input_ids": encoded["input_ids"].to(device), + "attention_mask": encoded["attention_mask"].to(device), + } + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="nanbeige", + engine_dir="/tmp/nanbeige", + max_seq_len=1024, + ) + + program = exporter.export(model, sample_inputs, config=config) + + print("engines:", exporter.engines) + + with torch.no_grad(): + 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())) + print_bench(exporter.bench) + + +if __name__ == "__main__": + main() diff --git a/tools/hf/run_pi05_export.py b/tools/hf/run_pi05_export.py index f7abd318bc1..9fe027ea611 100644 --- a/tools/hf/run_pi05_export.py +++ b/tools/hf/run_pi05_export.py @@ -18,7 +18,7 @@ from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE -def load_pi05(device: torch.device) -> PI05Policy: +def load_policy(device: torch.device) -> PI05Policy: policy = PI05Policy.from_pretrained("lerobot/pi05_libero_base").eval() cfg = policy.config cfg.device = str(device) @@ -53,7 +53,7 @@ def main() -> None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float16 - policy = load_pi05(device) + policy = load_policy(device) policy.model.to(device=device, dtype=dtype).eval() paligemma = policy.model.paligemma_with_expert.paligemma.model From 8c254f0236a0ffdf7435acf5cd796156ce93e715 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 15 Sep 2026 18:38:49 -0700 Subject: [PATCH 10/11] Add Kimi-K3 text prefill exporter --- tools/hf/exporters/models/__init__.py | 1 + tools/hf/exporters/models/kimi/__init__.py | 1 + tools/hf/exporters/models/kimi/helpers.py | 78 ++++ tools/hf/exporters/models/kimi/patches.py | 350 ++++++++++++++++++ tools/hf/exporters/models/kimi/spec.py | 155 ++++++++ tools/hf/exporters/plugin/kimi_kda.py | 49 +++ tools/hf/exporters/plugin/plugin_converter.py | 32 ++ tools/hf/exporters/plugin/plugin_utils.py | 2 + tools/hf/run_kimi_export.py | 66 ++++ 9 files changed, 734 insertions(+) create mode 100644 tools/hf/exporters/models/kimi/__init__.py create mode 100644 tools/hf/exporters/models/kimi/helpers.py create mode 100644 tools/hf/exporters/models/kimi/patches.py create mode 100644 tools/hf/exporters/models/kimi/spec.py create mode 100644 tools/hf/exporters/plugin/kimi_kda.py create mode 100644 tools/hf/run_kimi_export.py diff --git a/tools/hf/exporters/models/__init__.py b/tools/hf/exporters/models/__init__.py index eaebd1daeb8..182bbd3efa4 100644 --- a/tools/hf/exporters/models/__init__.py +++ b/tools/hf/exporters/models/__init__.py @@ -1,5 +1,6 @@ """Model families. Importing this package registers EdgeSpecs.""" from .groot import spec as _groot # noqa: F401 +from .kimi import spec as _kimi # noqa: F401 from .nemotron import spec as _nemotron # noqa: F401 from .pi05 import spec as _pi05 # noqa: F401 diff --git a/tools/hf/exporters/models/kimi/__init__.py b/tools/hf/exporters/models/kimi/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tools/hf/exporters/models/kimi/__init__.py @@ -0,0 +1 @@ + diff --git a/tools/hf/exporters/models/kimi/helpers.py b/tools/hf/exporters/models/kimi/helpers.py new file mode 100644 index 00000000000..d28e9f43933 --- /dev/null +++ b/tools/hf/exporters/models/kimi/helpers.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def language_model(model: nn.Module) -> nn.Module: + language = getattr(model, "language_model", None) + if not isinstance(language, nn.Module): + raise AttributeError(f"{type(model).__name__} has no language_model") + return language + + +def decoder_model(model: nn.Module) -> nn.Module: + decoder = getattr(language_model(model), "model", None) + if not isinstance(decoder, nn.Module) or not hasattr(decoder, "layers"): + raise AttributeError(f"{type(model).__name__} has no Kimi decoder layers") + return decoder + + +def kda_layer_indices(model: nn.Module) -> list[int]: + return [ + index + for index, layer in enumerate(decoder_model(model).layers) + if bool(getattr(layer, "is_linear_attn", False)) + ] + + +def allocate_kda_states( + model: nn.Module, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[tuple[torch.Tensor, ...], list[str]]: + """Allocate Q/K/V convolution state plus V-first recurrent state per KDA layer.""" + states: list[torch.Tensor] = [] + names: list[str] = [] + decoder = decoder_model(model) + + for layer_index in kda_layer_indices(model): + attention = decoder.layers[layer_index].self_attn + num_heads = int(attention.num_heads) + head_dim = int(attention.head_dim) + projection_size = num_heads * head_dim + conv_size = int(attention.conv_size) + + for suffix in ("q", "k", "v"): + states.append( + torch.zeros( + batch_size, + projection_size, + conv_size, + device=device, + dtype=dtype, + ) + ) + names.append(f"kda_conv_{suffix}_{layer_index}") + + states.append( + torch.zeros( + batch_size, + num_heads, + head_dim, + head_dim, + device=device, + dtype=torch.float32, + ) + ) + names.append(f"kda_recurrent_{layer_index}") + + return tuple(states), names + + +def text_config(model: nn.Module) -> Any: + return language_model(model).config diff --git a/tools/hf/exporters/models/kimi/patches.py b/tools/hf/exporters/models/kimi/patches.py new file mode 100644 index 00000000000..346fa209d13 --- /dev/null +++ b/tools/hf/exporters/models/kimi/patches.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +from contextlib import ExitStack, contextmanager +from typing import Any, Callable, Iterator + +import torch +import torch.nn.functional as F + +from ...plugin.attn_patches import patch_attribute +from .helpers import decoder_model, language_model + + +def _conv_bias(conv, projected: torch.Tensor) -> torch.Tensor: + bias = getattr(conv, "bias", None) + if bias is not None: + return bias.to(device=projected.device, dtype=projected.dtype) + return torch.zeros( + int(conv.weight.shape[0]), + device=projected.device, + dtype=projected.dtype, + ) + + +def _plugin_convolution( + conv, + projected: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kernel_size = int(conv.weight.shape[-1]) + output, next_state = torch.ops.trt.causal_conv1d.default( + projected, + conv.weight, + _conv_bias(conv, projected), + state, + context_lengths, + 1, + kernel_size - 1, + 1, + int(projected.shape[-1]), + ) + return F.silu(output), next_state + + +def _patch_kda_attention(original: Callable) -> Callable: + def forward( + self, + hidden_states, + attention_mask=None, + cache_params=None, + *, + context_lengths=None, + conv_state_q=None, + conv_state_k=None, + conv_state_v=None, + recurrent_state=None, + **kwargs, + ): + if recurrent_state is None: + return original( + self, + hidden_states, + attention_mask=attention_mask, + cache_params=cache_params, + **kwargs, + ) + + del attention_mask, cache_params, kwargs + if context_lengths is None: + raise ValueError("Kimi KDA export requires context_lengths") + if conv_state_q is None or conv_state_k is None or conv_state_v is None: + raise ValueError("Kimi KDA export requires Q/K/V convolution states") + + q, _ = _plugin_convolution( + self.q_conv1d, + self.q_proj(hidden_states), + conv_state_q, + context_lengths, + ) + k, _ = _plugin_convolution( + self.k_conv1d, + self.k_proj(hidden_states), + conv_state_k, + context_lengths, + ) + v, _ = _plugin_convolution( + self.v_conv1d, + self.v_proj(hidden_states), + conv_state_v, + context_lengths, + ) + + batch_size, seq_len, _ = hidden_states.shape + num_heads = int(self.num_heads) + head_dim = int(self.head_dim) + q = q.reshape(batch_size, seq_len, num_heads, head_dim) + k = k.reshape(batch_size, seq_len, num_heads, head_dim) + v = v.reshape(batch_size, seq_len, num_heads, head_dim) + gate = self.f_b_proj(self.f_a_proj(hidden_states)).reshape( + batch_size, seq_len, num_heads, head_dim + ) + beta = self.b_proj(hidden_states).to(dtype=hidden_states.dtype) + lower_bound = getattr(self, "gate_lower_bound", None) + + output, _ = torch.ops.trt.kimi_kda_plugin.default( + q, + k, + v, + gate, + beta, + self.A_log.float(), + self.dt_bias.reshape(num_heads, head_dim).to(hidden_states.dtype), + recurrent_state, + context_lengths, + float(lower_bound if lower_bound is not None else -5.0), + lower_bound is not None, + ) + + output_gate = ( + self.g_proj(hidden_states) + if self.use_full_rank_gate + else self.g_b_proj(self.g_a_proj(hidden_states)) + ).reshape(batch_size, seq_len, num_heads, head_dim) + + output_float = output.float() + eps = float( + getattr( + self.o_norm, + "variance_epsilon", + getattr(self.o_norm, "eps", 1e-5), + ) + ) + output = ( + output_float + * torch.rsqrt(output_float.square().mean(-1, keepdim=True) + eps) + ).to(hidden_states.dtype) + output = output * self.o_norm.weight * output_gate.sigmoid() + return self.o_proj(output.reshape(batch_size, seq_len, -1)) + + return forward + + +def _patch_mla_attention(original: Callable) -> Callable: + def forward( + self, + hidden_states, + attention_mask=None, + position_ids=None, + past_key_values=None, + **kwargs, + ): + del position_ids, past_key_values, kwargs + batch_size, seq_len, _ = hidden_states.shape + query_shape = (batch_size, seq_len, -1, self.q_head_dim) + key_shape = ( + batch_size, + seq_len, + -1, + self.qk_nope_head_dim + self.v_head_dim, + ) + + if self.q_lora_rank is not None: + query = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + query = self.q_proj(hidden_states) + query = query.reshape(query_shape).transpose(1, 2) + query_nope, query_position = torch.split( + query, + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=-1, + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + key_latent, key_position = torch.split( + compressed_kv, + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + key_value = self.kv_b_proj(self.kv_a_layernorm(key_latent)) + key_value = key_value.reshape(key_shape).transpose(1, 2) + key_nope, value = torch.split( + key_value, + [self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + key_position = key_position.reshape( + batch_size, 1, seq_len, self.qk_rope_head_dim + ).expand(*key_nope.shape[:-1], -1) + query = torch.cat((query_nope, query_position), dim=-1) + key = torch.cat((key_nope, key_position), dim=-1) + + scores = torch.matmul(query, key.transpose(-2, -1)) * float(self.scaling) + if attention_mask is not None: + scores = scores + attention_mask[..., : key.shape[-2]] + probabilities = scores.softmax(dim=-1, dtype=torch.float32).to(query.dtype) + output = torch.matmul(probabilities, value).transpose(1, 2).contiguous() + output = output.reshape(batch_size, seq_len, -1) + if self.use_output_gate: + output = output * self.g_proj(hidden_states).sigmoid() + return self.o_proj(output) + + return forward + + +def _patch_sparse_moe(original: Callable) -> Callable: + def forward(self, hidden_states): + identity = hidden_states + original_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + flattened = hidden_states.reshape(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + flattened = self.routed_expert_down_proj(flattened) + + expert_outputs = torch.stack( + [expert(flattened) for expert in self.experts], + dim=1, + ) + routing = F.one_hot( + topk_indices, + num_classes=int(self.num_experts), + ).to(topk_weights.dtype) + routing = (routing * topk_weights.unsqueeze(-1)).sum(dim=1) + output = (expert_outputs * routing.unsqueeze(-1)).sum(dim=1) + output = output.to(flattened.dtype) + + if self.use_latent_moe: + if self.latent_moe_use_norm: + output = self.routed_expert_norm(output) + output = self.routed_expert_up_proj(output) + output = output.reshape(original_shape) + if self.config.num_shared_experts is not None: + output = output + self.shared_experts(identity).to(output.dtype) + return output + + return forward + + +def _causal_mask( + context_lengths: torch.Tensor, + seq_len: int, + dtype: torch.dtype, +) -> torch.Tensor: + positions = torch.arange(seq_len, device=context_lengths.device) + causal = positions.unsqueeze(0) <= positions.unsqueeze(1) + valid_keys = positions.unsqueeze(0) < context_lengths.unsqueeze(1) + allowed = causal.unsqueeze(0) & valid_keys.unsqueeze(1) + mask = torch.zeros( + context_lengths.shape[0], + 1, + seq_len, + seq_len, + device=context_lengths.device, + dtype=dtype, + ) + return mask.masked_fill(~allowed.unsqueeze(1), torch.finfo(dtype).min) + + +def _patch_language_forward(original: Callable) -> Callable: + def forward( + self, + inputs_embeds=None, + context_lengths=None, + last_token_ids=None, + *kda_states, + **kwargs, + ): + if context_lengths is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + if inputs_embeds is None or last_token_ids is None: + raise ValueError("Kimi export requires embeddings and last-token indices") + + decoder = self.model + hidden = inputs_embeds.to(dtype=next(decoder.parameters()).dtype) + batch_size, seq_len, hidden_size = hidden.shape + causal_mask = _causal_mask(context_lengths, seq_len, hidden.dtype) + block_residual = hidden.new_zeros(batch_size * seq_len, 0, hidden_size) + state_index = 0 + + for layer in decoder.layers: + layer_kwargs: dict[str, Any] = {} + layer_mask = causal_mask + if layer.is_linear_attn: + if state_index + 4 > len(kda_states): + raise ValueError("Missing flattened Kimi KDA states") + layer_kwargs = { + "context_lengths": context_lengths, + "conv_state_q": kda_states[state_index], + "conv_state_k": kda_states[state_index + 1], + "conv_state_v": kda_states[state_index + 2], + "recurrent_state": kda_states[state_index + 3], + } + state_index += 4 + layer_mask = None + + hidden, block_residual = layer( + hidden, + attention_mask=layer_mask, + block_residual=block_residual, + **layer_kwargs, + ) + + hidden = decoder._apply_output_attn_res(hidden, block_residual) + hidden = decoder.norm(hidden) + indices = ( + last_token_ids if last_token_ids.ndim == 1 else last_token_ids.squeeze(-1) + ) + selected = hidden[ + torch.arange(batch_size, device=hidden.device, dtype=torch.long), + indices, + ] + return self.lm_head(selected).float() + + return forward + + +@contextmanager +def apply_kimik3_patches(model: Any | None = None) -> Iterator[None]: + if model is None: + yield + return + + language = language_model(model) + decoder = decoder_model(model) + kda_class = mla_class = moe_class = None + for layer in decoder.layers: + if layer.is_linear_attn: + kda_class = type(layer.self_attn) + else: + mla_class = type(layer.self_attn) + if hasattr(layer, "block_sparse_moe"): + moe_class = type(layer.block_sparse_moe) + + with ExitStack() as stack: + stack.enter_context( + patch_attribute(type(language), "forward", _patch_language_forward) + ) + if kda_class is not None: + stack.enter_context( + patch_attribute(kda_class, "forward", _patch_kda_attention) + ) + if mla_class is not None: + stack.enter_context( + patch_attribute(mla_class, "forward", _patch_mla_attention) + ) + if moe_class is not None: + stack.enter_context( + patch_attribute(moe_class, "forward", _patch_sparse_moe) + ) + yield diff --git a/tools/hf/exporters/models/kimi/spec.py b/tools/hf/exporters/models/kimi/spec.py new file mode 100644 index 00000000000..8fd310f0f26 --- /dev/null +++ b/tools/hf/exporters/models/kimi/spec.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn + +from ...ops import call_engine +from ...spec import ComponentBundle, EdgeSpec, register_edge_spec +from .helpers import allocate_kda_states, language_model +from .patches import apply_kimik3_patches + + +@register_edge_spec("kimi", "kimi-k3", "kimi_k3") +class KimiK3Spec(EdgeSpec): # type: ignore[misc] + def apply_patches(self, model=None): + return apply_kimik3_patches(model) + + def prepare_sample_inputs( + self, + model: nn.Module, + raw: Mapping[str, Any], + config: Any, + ) -> MutableMapping[str, Any]: + del config + if "inputs_embeds" in raw: + inputs_embeds = raw["inputs_embeds"] + elif "input_ids" in raw: + embedding = model.get_input_embeddings() + input_ids = raw["input_ids"].to(embedding.weight.device) + with torch.no_grad(): + inputs_embeds = embedding(input_ids) + else: + raise KeyError("Kimi K3 requires input_ids or inputs_embeds") + + attention_mask = raw.get("attention_mask") + if attention_mask is None: + attention_mask = torch.ones( + inputs_embeds.shape[:2], + device=inputs_embeds.device, + dtype=torch.long, + ) + else: + attention_mask = attention_mask.to(inputs_embeds.device) + return { + "inputs_embeds": inputs_embeds.contiguous(), + "attention_mask": attention_mask.contiguous(), + } + + def capture_eager_outputs( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + bench: dict[str, float] | None = None, + ) -> dict[str, torch.Tensor]: + del config + from ...measure import cuda_ms + + kwargs = { + "inputs_embeds": sample["inputs_embeds"], + "attention_mask": sample["attention_mask"], + "use_cache": False, + "return_dict": True, + } + with torch.no_grad(): + logits = language_model(model)(**kwargs).logits + + context_lengths = sample["attention_mask"].sum(dim=-1, dtype=torch.int64) + last_token_ids = context_lengths - 1 + batch = torch.arange(logits.shape[0], device=logits.device) + selected_logits = logits[batch, last_token_ids] + + if bench is not None: + bench["language"] = cuda_ms(lambda: language_model(model)(**kwargs).logits) + return {"language": selected_logits} + + def prepare( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + del config + embeddings = sample["inputs_embeds"] + batch_size = int(embeddings.shape[0]) + context_lengths = sample["attention_mask"].sum(dim=-1, dtype=torch.int32) + if bool((context_lengths <= 0).any()): + raise ValueError("Kimi attention_mask contains an empty sequence") + last_token_ids = (context_lengths - 1).to(torch.int64).unsqueeze(-1) + states, state_names = allocate_kda_states( + model, + batch_size=batch_size, + device=embeddings.device, + dtype=embeddings.dtype, + ) + args = (embeddings, context_lengths, last_token_ids, *states) + input_names = [ + "inputs_embeds", + "context_lengths", + "last_token_ids", + *state_names, + ] + sample.update(dict(zip(input_names, args))) + + return { + "language": ComponentBundle( + module=language_model(model).eval(), + trace_args=args, + save_args=args, + input_names=input_names, + output_names=["logits"], + parity_output="logits", + model_type="kimi_k3", + 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: + state_names = sorted( + ( + name + for name in sample + if name.startswith("kda_conv_") or name.startswith("kda_recurrent_") + ), + key=_state_sort_key, + ) + return call_engine( + engines["language"], + "language", + sample["inputs_embeds"], + sample["context_lengths"], + sample["last_token_ids"], + *(sample[name] for name in state_names), + ) + + +def _state_sort_key(name: str) -> tuple[int, int]: + layer_index = int(name.rsplit("_", 1)[1]) + kind = name[: name.rfind("_")] + order = { + "kda_conv_q": 0, + "kda_conv_k": 1, + "kda_conv_v": 2, + "kda_recurrent": 3, + } + return layer_index, order[kind] diff --git a/tools/hf/exporters/plugin/kimi_kda.py b/tools/hf/exporters/plugin/kimi_kda.py new file mode 100644 index 00000000000..6642d1b3d29 --- /dev/null +++ b/tools/hf/exporters/plugin/kimi_kda.py @@ -0,0 +1,49 @@ +"""Kimi KDA custom op lowered to the Edge-LLM ``kimi_kda`` plugin.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def register_kimi_kda_plugin_op() -> None: + """Register ``trt::kimi_kda_plugin`` for Dynamo export.""" + if hasattr(torch.ops, "trt") and hasattr(torch.ops.trt, "kimi_kda_plugin"): + return + + @torch.library.custom_op("trt::kimi_kda_plugin", mutates_args=()) + def kimi_kda_plugin( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, + lower_bound: float, + use_lower_bound: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del q, k, gate, beta, a_log, dt_bias + del context_lengths, lower_bound, use_lower_bound + return torch.empty_like(v), torch.empty_like(state) + + @kimi_kda_plugin.register_fake + def _( + q, + k, + v, + gate, + beta, + a_log, + dt_bias, + state, + context_lengths, + lower_bound, + use_lower_bound, + ): + del q, k, gate, beta, a_log, dt_bias + del context_lengths, lower_bound, use_lower_bound + return torch.empty_like(v), torch.empty_like(state) diff --git a/tools/hf/exporters/plugin/plugin_converter.py b/tools/hf/exporters/plugin/plugin_converter.py index efb6db8ebdd..dbec1886833 100644 --- a/tools/hf/exporters/plugin/plugin_converter.py +++ b/tools/hf/exporters/plugin/plugin_converter.py @@ -296,6 +296,38 @@ def convert_update_ssm_state(ctx: ConversionContext, target, args, kwargs, name) return layer.get_output(0), layer.get_output(1) +@dynamo_tensorrt_converter( + torch.ops.trt.kimi_kda_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_kimi_kda_plugin(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + tensors = args[:9] + lower_bound = float(args[9]) + use_lower_bound = int(bool(args[10])) + + creator = get_trt_plugin_creator("kimi_kda", "1", "") + if creator is None: + raise RuntimeError("kimi_kda not found in TensorRT plugin registry") + + plugin = _create_trt_plugin( + creator, + name, + [ + _float_field("lower_bound", lower_bound), + _int_field("use_lower_bound", use_lower_bound), + ], + ) + if plugin is None: + raise RuntimeError("Failed to create kimi_kda plugin") + + inputs = _as_plugin_inputs(ctx, tensors, name) + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0), layer.get_output(1) + + def _convert_nvfp4_moe(ctx: ConversionContext, args, name: str, plugin_name: str): args = list(args) tensors = args[:11] diff --git a/tools/hf/exporters/plugin/plugin_utils.py b/tools/hf/exporters/plugin/plugin_utils.py index 31a3c58a82c..e983f178cf6 100644 --- a/tools/hf/exporters/plugin/plugin_utils.py +++ b/tools/hf/exporters/plugin/plugin_utils.py @@ -340,6 +340,7 @@ def load_plugin(): def load_plugins_for_trt(): + from .kimi_kda import register_kimi_kda_plugin_op from .mamba import register_mamba_plugin_ops from .moe import register_moe_plugin_ops @@ -347,6 +348,7 @@ def load_plugins_for_trt(): _register_vit_attention_plugin_op() register_mamba_plugin_ops() register_moe_plugin_ops() + register_kimi_kda_plugin_op() from . import attn_patches as _attn_patches # noqa: F401,E402 load_plugin() diff --git a/tools/hf/run_kimi_export.py b/tools/hf/run_kimi_export.py new file mode 100644 index 00000000000..3c756a46267 --- /dev/null +++ b/tools/hf/run_kimi_export.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""EdgeExporter on kimi-k3 model.""" + +from __future__ import annotations + +import torch +import torch_tensorrt +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt +from transformers import AutoModelForCausalLM, AutoTokenizer + +checkpoint = "inference-optimization/Kimi-K3-0.40B" + + +def load_policy(device: torch.device, dtype: torch.dtype): + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + # if pad token is not set, set it to eos token + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + return model, tokenizer + + +def main() -> None: + load_plugins_for_trt() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float16 + + model, tokenizer = load_policy(device, dtype) + encoded = tokenizer("Hello, how are you?", return_tensors="pt") + sample_inputs = { + "input_ids": encoded["input_ids"].to(device), + "attention_mask": encoded["attention_mask"].to(device), + } + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="kimi", + engine_dir="/tmp/kimi", + max_seq_len=1024, + ) + + program = exporter.export(model, sample_inputs, config=config) + + print("engines:", exporter.engines) + + with torch.no_grad(): + 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())) + print_bench(exporter.bench) + + +if __name__ == "__main__": + main() From aa090bb4272ac4847667f82958e2b534d74a76a4 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 16 Sep 2026 14:53:48 -0700 Subject: [PATCH 11/11] Unify Edge exporter smoke runners --- tools/hf/exporters/models/groot/export.py | 59 +++++++++++++ tools/hf/exporters/models/kimi/export.py | 39 +++++++++ tools/hf/exporters/models/nanbeige/export.py | 39 +++++++++ tools/hf/exporters/models/nemotron/export.py | 41 +++++++++ tools/hf/exporters/models/pi05/export.py | 65 +++++++++++++++ tools/hf/run_export.py | 84 +++++++++++++++++++ tools/hf/run_groot_export.py | 80 ------------------ tools/hf/run_kimi_export.py | 66 --------------- tools/hf/run_nanbeige_export.py | 71 ---------------- tools/hf/run_nemotron_export.py | 80 ------------------ tools/hf/run_pi05_export.py | 88 -------------------- 11 files changed, 327 insertions(+), 385 deletions(-) create mode 100644 tools/hf/exporters/models/groot/export.py create mode 100644 tools/hf/exporters/models/kimi/export.py create mode 100644 tools/hf/exporters/models/nanbeige/export.py create mode 100644 tools/hf/exporters/models/nemotron/export.py create mode 100644 tools/hf/exporters/models/pi05/export.py create mode 100644 tools/hf/run_export.py delete mode 100644 tools/hf/run_groot_export.py delete mode 100644 tools/hf/run_kimi_export.py delete mode 100644 tools/hf/run_nanbeige_export.py delete mode 100644 tools/hf/run_nemotron_export.py delete mode 100644 tools/hf/run_pi05_export.py diff --git a/tools/hf/exporters/models/groot/export.py b/tools/hf/exporters/models/groot/export.py new file mode 100644 index 00000000000..325a2cddd48 --- /dev/null +++ b/tools/hf/exporters/models/groot/export.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import argparse + +import torch +from lerobot.configs import FeatureType, PolicyFeature +from lerobot.policies.groot import GrootPolicy +from lerobot.policies.groot.configuration_groot import GrootConfig +from lerobot.utils.constants import ACTION, OBS_STATE + +from ...config import EdgeConfig +from ...utils import force_hf_attention + +DEFAULT_CHECKPOINT = "nvidia/GR00T-N1.5-3B" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + policy_config = GrootConfig( + base_model_path=args.checkpoint or DEFAULT_CHECKPOINT, + device=str(device), + embodiment_tag="new_embodiment", + chunk_size=50, + n_action_steps=50, + max_state_dim=64, + max_action_dim=32, + image_size=(224, 224), + tokenizer_assets_repo="lerobot/eagle2hg-processor-groot-n1p5", + input_features={ + "observation.images.image": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + "observation.images.image2": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,))}, + ) + policy = GrootPolicy(policy_config).to(device).eval() + model = policy._groot_model.to(device=device, dtype=dtype).eval() + eagle = model.backbone.eagle_model + force_hf_attention(eagle.vision_model, "eager") + force_hf_attention(eagle.language_model, "eager") + + export_config = EdgeConfig( + model_type="groot", + engine_dir=args.engine_dir or "/tmp/groot_edge_exporter", + max_seq_len=args.max_seq_len or 968, + ) + return ( + policy, + {"device": device, "dtype": dtype}, + export_config, + "velocity", + ) diff --git a/tools/hf/exporters/models/kimi/export.py b/tools/hf/exporters/models/kimi/export.py new file mode 100644 index 00000000000..81e7a7148e2 --- /dev/null +++ b/tools/hf/exporters/models/kimi/export.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from ...config import EdgeConfig + +DEFAULT_CHECKPOINT = "inference-optimization/Kimi-K3-0.40B" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + checkpoint = args.checkpoint or DEFAULT_CHECKPOINT + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + + encoded = tokenizer(args.prompt, return_tensors="pt") + sample_inputs = {name: tensor.to(device) for name, tensor in encoded.items()} + config = EdgeConfig( + model_type="kimi", + engine_dir=args.engine_dir or "/tmp/kimi", + max_seq_len=args.max_seq_len or 1024, + ) + return model, sample_inputs, config, "logits" diff --git a/tools/hf/exporters/models/nanbeige/export.py b/tools/hf/exporters/models/nanbeige/export.py new file mode 100644 index 00000000000..0bb8b8659b5 --- /dev/null +++ b/tools/hf/exporters/models/nanbeige/export.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from ...config import EdgeConfig + +DEFAULT_CHECKPOINT = "Nanbeige/Nanbeige4.2-3B" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + checkpoint = args.checkpoint or DEFAULT_CHECKPOINT + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + + encoded = tokenizer(args.prompt, return_tensors="pt") + sample_inputs = {name: tensor.to(device) for name, tensor in encoded.items()} + config = EdgeConfig( + model_type="nanbeige", + engine_dir=args.engine_dir or "/tmp/nanbeige", + max_seq_len=args.max_seq_len or 1024, + ) + return model, sample_inputs, config, "logits" diff --git a/tools/hf/exporters/models/nemotron/export.py b/tools/hf/exporters/models/nemotron/export.py new file mode 100644 index 00000000000..1d548ee8e81 --- /dev/null +++ b/tools/hf/exporters/models/nemotron/export.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from ...config import EdgeConfig +from .mamba_stub import apply as apply_mamba_stub + +DEFAULT_CHECKPOINT = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + checkpoint = args.checkpoint or DEFAULT_CHECKPOINT + apply_mamba_stub() + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + + encoded = tokenizer(args.prompt, return_tensors="pt") + sample_inputs = {name: tensor.to(device) for name, tensor in encoded.items()} + config = EdgeConfig( + model_type="nemotron_h", + engine_dir=args.engine_dir or "/tmp/nemotron_edge_exporter", + max_seq_len=args.max_seq_len or 128, + ) + return model, sample_inputs, config, "logits" diff --git a/tools/hf/exporters/models/pi05/export.py b/tools/hf/exporters/models/pi05/export.py new file mode 100644 index 00000000000..7b9137dfcd7 --- /dev/null +++ b/tools/hf/exporters/models/pi05/export.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import argparse + +import torch +from lerobot.configs import FeatureType, PolicyFeature +from lerobot.policies.pi05 import PI05Policy +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE + +from ...config import EdgeConfig +from ...utils import force_hf_attention + +DEFAULT_CHECKPOINT = "lerobot/pi05_libero_base" + + +def prepare_export( + args: argparse.Namespace, + device: torch.device, + dtype: torch.dtype, +): + policy = PI05Policy.from_pretrained(args.checkpoint or DEFAULT_CHECKPOINT).eval() + config = policy.config + config.device = str(device) + config.chunk_size = 50 + config.n_action_steps = 50 + config.max_state_dim = 32 + config.max_action_dim = 32 + config.input_features = { + f"{OBS_IMAGES}.image": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image2": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image3": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + f"{OBS_IMAGES}.image4": PolicyFeature( + type=FeatureType.VISUAL, shape=(3, 224, 224) + ), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(32,)), + } + config.output_features = { + ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,)) + } + config.empty_cameras = 0 + config.validate_features() + + policy.model.to(device=device, dtype=dtype).eval() + paligemma = policy.model.paligemma_with_expert.paligemma.model + force_hf_attention(paligemma.vision_tower, "eager") + force_hf_attention(paligemma.language_model, "eager") + force_hf_attention(policy.model.paligemma_with_expert.gemma_expert.model, "eager") + + export_config = EdgeConfig( + model_type="pi05", + engine_dir=args.engine_dir or "/tmp/pi05", + max_seq_len=args.max_seq_len or 968, + ) + return ( + policy, + {"device": device, "dtype": dtype}, + export_config, + "velocity", + ) diff --git a/tools/hf/run_export.py b/tools/hf/run_export.py new file mode 100644 index 00000000000..b62b75caea1 --- /dev/null +++ b/tools/hf/run_export.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Run an EdgeExporter smoke test for a supported model family.""" + +from __future__ import annotations + +import argparse +import importlib +from collections.abc import Callable, MutableMapping +from typing import Any + +import torch +import torch.nn as nn +import torch_tensorrt # noqa: F401 +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt + +PreparedExport = tuple[nn.Module, MutableMapping[str, Any], EdgeConfig, str] +ExportPreparer = Callable[ + [argparse.Namespace, torch.device, torch.dtype], PreparedExport +] + +EXPORT_PREPARERS = { + "groot": "exporters.models.groot.export:prepare_export", + "pi05": "exporters.models.pi05.export:prepare_export", + "nemotron": "exporters.models.nemotron.export:prepare_export", + "nanbeige": "exporters.models.nanbeige.export:prepare_export", + "kimi": "exporters.models.kimi.export:prepare_export", +} + + +def get_export_preparer(model_type: str) -> ExportPreparer: + module_name, function_name = EXPORT_PREPARERS[model_type].split(":") + module = importlib.import_module(module_name) + return getattr(module, function_name) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model_type", choices=EXPORT_PREPARERS) + parser.add_argument("--checkpoint") + parser.add_argument("--prompt", default="Hello, how are you?") + parser.add_argument("--engine-dir") + parser.add_argument("--max-seq-len", type=int) + parser.add_argument("--device") + parser.add_argument( + "--dtype", + choices=("float16", "bfloat16"), + default="float16", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + load_plugins_for_trt() + + device = torch.device( + args.device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + dtype = getattr(torch, args.dtype) + prepare_export = get_export_preparer(args.model_type) + model, sample_inputs, config, output_name = prepare_export(args, device, dtype) + + exporter = EdgeExporter() + program = exporter.export(model, sample_inputs, config=config) + + print("engines:", exporter.engines) + print("runtime keys:", sorted(exporter.sample)) + with torch.no_grad(): + result = program.module()(**exporter.sample) + + output = result[0] if isinstance(result, (tuple, list)) else result + print( + output_name, + tuple(output.shape), + "mean", + float(output.float().mean()), + ) + print_bench(exporter.bench) + + +if __name__ == "__main__": + main() diff --git a/tools/hf/run_groot_export.py b/tools/hf/run_groot_export.py deleted file mode 100644 index d21007e8250..00000000000 --- a/tools/hf/run_groot_export.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/usr/bin/env python3 -"""Smoke EdgeExporter on GR00T (4 engines: vision, language, context_projection, action). - -Pass the LeRobot GrootPolicy, not policy._groot_model — prepare_sample_inputs -needs GrootEagleEncodeStep / embodiment_id from the policy wrapper. -""" - -from __future__ import annotations - -import torch # noqa: E402 -import torch_tensorrt # noqa: E402 -from exporters import EdgeConfig, EdgeExporter -from exporters.measure import print_bench -from exporters.plugin.plugin_utils import load_plugins_for_trt -from exporters.utils import force_hf_attention -from lerobot.configs import FeatureType, PolicyFeature -from lerobot.policies.groot import GrootPolicy -from lerobot.policies.groot.configuration_groot import GrootConfig -from lerobot.utils.constants import ACTION, OBS_STATE - - -def load_groot(device: torch.device) -> GrootPolicy: - config = GrootConfig( - base_model_path="nvidia/GR00T-N1.5-3B", - device=str(device), - embodiment_tag="new_embodiment", - chunk_size=50, - n_action_steps=50, - max_state_dim=64, - max_action_dim=32, - image_size=(224, 224), - tokenizer_assets_repo="lerobot/eagle2hg-processor-groot-n1p5", - input_features={ - "observation.images.image": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - "observation.images.image2": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), - }, - output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,))}, - ) - return GrootPolicy(config).to(device).eval() - - -def main() -> None: - load_plugins_for_trt() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float16 - - policy = load_groot(device) - model = policy._groot_model.to(device=device, dtype=dtype).eval() - eagle = model.backbone.eagle_model - force_hf_attention(eagle.vision_model, "eager") - force_hf_attention(eagle.language_model, "eager") - - exporter = EdgeExporter() - 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} - program = exporter.export(policy, sample_inputs, config=config) - - print("engines:", exporter.engines) - print("runtime keys:", sorted(exporter.sample)) - - with torch.no_grad(): - 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())) - print_bench(exporter.bench) - - -if __name__ == "__main__": - main() diff --git a/tools/hf/run_kimi_export.py b/tools/hf/run_kimi_export.py deleted file mode 100644 index 3c756a46267..00000000000 --- a/tools/hf/run_kimi_export.py +++ /dev/null @@ -1,66 +0,0 @@ -#!/usr/bin/env python3 -"""EdgeExporter on kimi-k3 model.""" - -from __future__ import annotations - -import torch -import torch_tensorrt -from exporters import EdgeConfig, EdgeExporter -from exporters.measure import print_bench -from exporters.plugin.plugin_utils import load_plugins_for_trt -from transformers import AutoModelForCausalLM, AutoTokenizer - -checkpoint = "inference-optimization/Kimi-K3-0.40B" - - -def load_policy(device: torch.device, dtype: torch.dtype): - model = ( - AutoModelForCausalLM.from_pretrained( - checkpoint, - trust_remote_code=True, - dtype=dtype, - ) - .to(device=device, dtype=dtype) - .eval() - ) - tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) - # if pad token is not set, set it to eos token - if tokenizer.pad_token_id is None: - tokenizer.pad_token = tokenizer.eos_token - return model, tokenizer - - -def main() -> None: - load_plugins_for_trt() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float16 - - model, tokenizer = load_policy(device, dtype) - encoded = tokenizer("Hello, how are you?", return_tensors="pt") - sample_inputs = { - "input_ids": encoded["input_ids"].to(device), - "attention_mask": encoded["attention_mask"].to(device), - } - - exporter = EdgeExporter() - config = EdgeConfig( - model_type="kimi", - engine_dir="/tmp/kimi", - max_seq_len=1024, - ) - - program = exporter.export(model, sample_inputs, config=config) - - print("engines:", exporter.engines) - - with torch.no_grad(): - 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())) - print_bench(exporter.bench) - - -if __name__ == "__main__": - main() diff --git a/tools/hf/run_nanbeige_export.py b/tools/hf/run_nanbeige_export.py deleted file mode 100644 index 9bfd5f4d3a6..00000000000 --- a/tools/hf/run_nanbeige_export.py +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env python3 -"""Smoke EdgeExporter on pi05 model. - -Pass the LeRobot PI05Policy, not policy.model — prepare_sample_inputs -needs the preprocessor on the policy wrapper. -""" - -from __future__ import annotations - -import torch -import torch_tensorrt -from exporters import EdgeConfig, EdgeExporter -from exporters.measure import print_bench -from exporters.plugin.plugin_utils import load_plugins_for_trt -from exporters.utils import force_hf_attention -from transformers import AutoModelForCausalLM, AutoTokenizer - -checkpoint = "Nanbeige/Nanbeige4.2-3B" - - -def load_policy(device: torch.device, dtype: torch.dtype): - model = ( - AutoModelForCausalLM.from_pretrained( - checkpoint, - trust_remote_code=True, - torch_dtype=dtype, - ) - .to(device=device, dtype=dtype) - .eval() - ) - tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) - # if pad token is not set, set it to eos token - if tokenizer.pad_token_id is None: - tokenizer.pad_token = tokenizer.eos_token - return model, tokenizer - - -def main() -> None: - load_plugins_for_trt() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float16 - - model, tokenizer = load_policy(device, dtype) - encoded = tokenizer("Hello, how are you?", return_tensors="pt") - sample_inputs = { - "input_ids": encoded["input_ids"].to(device), - "attention_mask": encoded["attention_mask"].to(device), - } - - exporter = EdgeExporter() - config = EdgeConfig( - model_type="nanbeige", - engine_dir="/tmp/nanbeige", - max_seq_len=1024, - ) - - program = exporter.export(model, sample_inputs, config=config) - - print("engines:", exporter.engines) - - with torch.no_grad(): - 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())) - print_bench(exporter.bench) - - -if __name__ == "__main__": - main() diff --git a/tools/hf/run_nemotron_export.py b/tools/hf/run_nemotron_export.py deleted file mode 100644 index 14da09de3e8..00000000000 --- a/tools/hf/run_nemotron_export.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/usr/bin/env python3 -"""Smoke EdgeExporter on Nemotron-H (one language engine: attn + mamba + MoE). - -Pass the HF causal LM. Collation is tokenizer → input_ids; the spec embeds and -pads to max_seq_len. apply_mamba_stub() must run before from_pretrained. -""" - -from __future__ import annotations - -import argparse - -import torch -import torch_tensorrt -from exporters import EdgeConfig, EdgeExporter -from exporters.measure import print_bench -from exporters.models.nemotron.mamba_stub import apply as apply_mamba_stub -from exporters.plugin.plugin_utils import load_plugins_for_trt -from transformers import AutoModelForCausalLM, AutoTokenizer - - -def load_nemotron(checkpoint: str, device: torch.device, dtype: torch.dtype): - apply_mamba_stub() - model = ( - AutoModelForCausalLM.from_pretrained( - checkpoint, - trust_remote_code=True, - torch_dtype=dtype, - ) - .to(device=device, dtype=dtype) - .eval() - ) - tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) - if tokenizer.pad_token_id is None: - tokenizer.pad_token = tokenizer.eos_token - return model, tokenizer - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--checkpoint", - default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", - ) - parser.add_argument("--prompt", default="Hello.") - parser.add_argument("--max-seq-len", type=int, default=128) - args = parser.parse_args() - - load_plugins_for_trt() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float16 - - model, tokenizer = load_nemotron(args.checkpoint, device, dtype) - encoded = tokenizer(args.prompt, return_tensors="pt") - sample_inputs = { - "input_ids": encoded["input_ids"].to(device), - "attention_mask": encoded["attention_mask"].to(device), - } - - exporter = EdgeExporter() - config = EdgeConfig( - model_type="nemotron_h", - engine_dir="/tmp/nemotron_edge_exporter", - max_seq_len=args.max_seq_len, - ) - program = exporter.export(model, sample_inputs, config=config) - - print("engines:", exporter.engines) - print("runtime keys:", sorted(exporter.sample)) - - with torch.no_grad(): - 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())) - print_bench(exporter.bench) - - -if __name__ == "__main__": - main() diff --git a/tools/hf/run_pi05_export.py b/tools/hf/run_pi05_export.py deleted file mode 100644 index 9fe027ea611..00000000000 --- a/tools/hf/run_pi05_export.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -"""Smoke EdgeExporter on pi05 model. - -Pass the LeRobot PI05Policy, not policy.model — prepare_sample_inputs -needs the preprocessor on the policy wrapper. -""" - -from __future__ import annotations - -import torch -import torch_tensorrt -from exporters import EdgeConfig, EdgeExporter -from exporters.measure import print_bench -from exporters.plugin.plugin_utils import load_plugins_for_trt -from exporters.utils import force_hf_attention -from lerobot.configs import FeatureType, PolicyFeature -from lerobot.policies.pi05 import PI05Policy -from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE - - -def load_policy(device: torch.device) -> PI05Policy: - policy = PI05Policy.from_pretrained("lerobot/pi05_libero_base").eval() - cfg = policy.config - cfg.device = str(device) - cfg.chunk_size = 50 - cfg.n_action_steps = 50 - cfg.max_state_dim = 32 - cfg.max_action_dim = 32 - cfg.input_features = { - f"{OBS_IMAGES}.image": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - f"{OBS_IMAGES}.image2": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - f"{OBS_IMAGES}.image3": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - f"{OBS_IMAGES}.image4": PolicyFeature( - type=FeatureType.VISUAL, shape=(3, 224, 224) - ), - OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(32,)), - } - cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(32,))} - cfg.empty_cameras = 0 - cfg.validate_features() - return policy - - -def main() -> None: - load_plugins_for_trt() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float16 - - policy = load_policy(device) - - policy.model.to(device=device, dtype=dtype).eval() - paligemma = policy.model.paligemma_with_expert.paligemma.model - force_hf_attention(paligemma.vision_tower, "eager") - force_hf_attention(paligemma.language_model, "eager") - force_hf_attention(policy.model.paligemma_with_expert.gemma_expert.model, "eager") - - exporter = EdgeExporter() - config = EdgeConfig( - model_type="pi05", # optional; inferred from paligemma_with_expert - engine_dir="/tmp/pi05", - max_seq_len=968, - ) - - 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, …). - runtime_kwargs = exporter.sample - print("runtime keys:", sorted(runtime_kwargs)) - - with torch.no_grad(): - 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())) - print_bench(exporter.bench) - - -if __name__ == "__main__": - main()