Skip to content

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

Open
micwill755 wants to merge 1 commit into
exporters-dual-profilefrom
nanbeige
Open

micwill755 wants to merge 1 commit 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

@meta-cla meta-cla Bot added the cla signed label Sep 10, 2026
@micwill755
micwill755 added this pull request to stack #4727 September 21, 2026 21:58

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant