Skip to content

Add Nanbeige4.2-3B export support - #4712

Open
micwill755 wants to merge 35 commits into
exporters-dual-profilefrom
nanbeige
Open

micwill755 wants to merge 35 commits into
exporters-dual-profilefrom
nanbeige

Conversation

@micwill755

@micwill755 micwill755 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds TensorRT Edge Exporter support for Nanbeige/Nanbeige4.2-3B.

Nanbeige uses a looped-transformer architecture with 22 physical decoder layers executed twice. This change:

  • Registers a Nanbeige EdgeSpec.
  • Implements loop-aware language export with 44 logical KV-cache slots.
  • Adds live class patching for the trust_remote_code model.
  • Uses the TensorRT attention plugin with causal attention.
  • Repairs invalid deterministic RoPE buffers caused by Transformers remote-code version drift.
  • Adds an end-to-end Nanbeige export example.
  • Moves the live GR00T patch context into its family spec and aligns the PI05 example loader naming.

Motivation and context

Nanbeige4.2-3B reuses the same transformer blocks across multiple loops while requiring independent KV-cache state for each logical layer execution. Existing causal-LM export paths only process each physical layer once and therefore cannot represent this architecture correctly.

The model's remote code also targets an older Transformers version. With the current environment, its non-persistent RoPE frequency buffers may be zero-filled or invalid after loading. The compatibility helper deterministically reconstructs those buffers from the model configuration before eager parity capture.

Validation

Tested end-to-end with Nanbeige/Nanbeige4.2-3B in FP16:

  • Eager/TRT mean absolute error: 0.008464
  • Maximum absolute error: 0.046875
  • Relative L2 error: 0.0016
  • Values within tolerance: 99.4%
  • Eager execution: 15.948 ms
  • TensorRT execution: 9.446 ms
  • Speedup: 1.688x

Targeted Edge Exporter unit tests and pre-commit formatting/lint hooks pass locally.

Dependencies

Type of change

  • New feature (non-breaking change which adds functionality)

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented hard-to-understand compatibility and patching behavior
  • I have added the relevant labels so that relevant reviewers are notified

Torch-TensorRT Github Bot and others added 15 commits September 3, 2026 18:13
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Co-authored-by: apbose <apbose694@gmail.com>
Co-authored-by: apbose <apbose694@gmail.com>
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Signed-off-by: Torch-TensorRT Github Bot <torch-tensorrt.github.bot@nvidia.com>
Drop the Test/trt import path so PI05, GR00T, and Nemotron export is self-contained, and rename GR00T action_context to context_projection.
Each family now has spec, helpers, and patches in one folder, with shared vision/language/action-step code in models/common.
Compile the original HF/LeRobot submodules under apply_patches and stitch them with execute_engine, matching HF DynamoExporter.
…A export.

Keep Hub load stubs next to the Nemotron spec, and mention PI05/GR00T run_vla.py in the LLM tools README.
Document EdgeExporter.export, the returned ExportedProgram, family patches, and plugin converters, and link it from the user guide and HuggingFace tutorials.
@micwill755
micwill755 added this pull request to stack #4727 September 21, 2026 21:58
@github-actions github-actions Bot added documentation Improvements or additions to documentation component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters component: build system Issues re: Build system component: api [Python] Issues re: Python API component: runtime component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: torch_compile labels Sep 22, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: build system Issues re: Build system component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: runtime component: tests Issues re: Tests component: torch_compile documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants