Repository navigation
fix(dynamo): save programs with no inputs, keep attention masks that are not plain causal, fix uint8 shape ops and views left to PyTorch, show the runtime TensorRT build number, keep plugins that run PyTorch out of ExecuTorch programs - #4953
Open
shoumikhin wants to merge 117 commits into
Open
shoumikhin wants to merge 117 commits into
shoumikhin wants to merge 117 commits into
Conversation
…d errors
When the ExecuTorch TensorRT backend cannot read an engine, it logs which
TensorRT exported the engine and which TensorRT is loading it. The exported
version comes from tensorrt.__version__, which includes the build number, but
the runtime version was printed without it:
the engine was exported with TensorRT 11.4.0.106, and this runtime is TensorRT 11.4.0
So two equal versions looked different, and the message pointed at a version
mismatch that was not there.
Print the runtime version with its build number too, using TensorRT's own
getInferLibMajorVersion, getInferLibMinorVersion, getInferLibPatchVersion and
getInferLibBuildVersion. The message now reads:
the engine was exported with TensorRT 11.4.0.106, and this runtime is TensorRT 11.4.0.106
The existing load error test now checks the full runtime version.
torch_tensorrt.save(..., retrace=False), which is the default for ExecuTorch
output, failed for a compiled program with no inputs with a bare AssertionError:
class Draw(torch.nn.Module):
def forward(self):
return torch.randn(50, 32, device="cuda")
engine = torch_tensorrt.dynamo.compile(
torch.export.export(Draw(), ()), arg_inputs=(), min_block_size=1
)
torch_tensorrt.save(engine, "draw.pte", output_format="executorch", retrace=False)
# _exporter.py, in lift: assert fake_mode is not None -> AssertionError
lift() turns the graph's parameters and buffers into inputs, and gives each
lifted tensor a fake value made with the graph's fake mode. It looked for that
mode only on the existing inputs, so a program with no inputs found none, even
when there was nothing to lift.
lift() now finds the fake mode with torch's _detect_fake_mode_from_gm, which
reads the other nodes when there are no inputs. This repository already uses it
in its lowering passes.
Once the fake mode was found, a second problem showed up. Each lifted input was placed
before the first user input, and with no user input that means the front of
the graph. So two or more lifted constants ended up in reverse order, and the
ExportedProgram verifier rejected them:
Input spec name a_1 does not match node name b_1
Each lifted input now goes right after the previous one, which keeps the order
of the input specs with or without user inputs.
A new unit test builds the ExportedProgram for a program with no inputs that
reads a parameter and a buffer, and runs it.
github-actions
Bot
added
component: tests
Saving with retrace=False now gets past lift() for a program with no tensor
inputs. For a program whose only inputs are ints, that exposed a second
lookup. create_trt_exp_program took its fake mode only from tensor inputs
to build range_constraints, so it found none and left out the range of each
int input. The save then reported success, but loading the file failed:
class Prefix(torch.nn.Module):
def forward(self, n: int):
return self.c[:n] * 2
# torch.export.load(...) -> vr must not be None for symbol s0
The range lookup now uses _detect_fake_mode_from_gm, the same lookup lift()
uses, which reads the other nodes when no input is a tensor.
A new unit test saves and loads such a program and checks the range of n.
…verifier The test for a program with no inputs sorted the input specs before comparing them, and its value check, w * c, gives the same result when the two constants are swapped. So a wrong order was caught only by torch's verifier. The test now compares the specs in their order and computes w - c, so a swap also shows up in the result. The write-only copy-back test pointed to "the no-input test above", which now lands on this test, so it and the kwargs test name the test they mean. The comment in lift() keeps only the reason for chaining the lifted placeholders.
The fix for saving a program with no inputs had only a unit test, which builds the graph straight from torch.export, so it never went through inline_trt_modules or lifted a TensorRT engine. The new test compiles a program with no inputs whose TensorRT engine reads a random tensor made in PyTorch, checks that an engine was built, and saves it as a .pte with retrace=False. With no user inputs, the engine is the only thing lift() has to turn into an input. It runs with the other ExecuTorch tests.
A uint8 input that goes through a shape op before a cast, such as `image.permute(2, 0, 1).float()` or `image[None].float()`, failed to compile with "Could not get tensor shape." TensorRT allows uint8 only at engine inputs and outputs, so its shuffle, slice and concatenation layers reject a uint8 tensor, and the next converter cannot read the shape. Both partitioners now run any node that would make a uint8 tensor in PyTorch, next to the existing checks for complex and high-rank tensors. Casts from uint8, and ops that read uint8 but return another type, still convert. A clone, and a slice_scatter that overwrites the whole tensor, can return an engine input as it is, so they are left to their converters. uint8 constants, which are packed FP4 weights, are not affected.
shoumikhin
changed the title
…hape_copy Torch-TensorRT turns every aten.view into aten._reshape_copy before it partitions the graph. A view TensorRT does not take, such as one with more than 8 dimensions in the patchify step of Qwen-VL image processors, stays in PyTorch as aten._reshape_copy. That is not a core ATen op, so a program saved with output_format="executorch" and the ExecuTorch CUDA partitioner fails with "missing fallback kernels: aten::_reshape_copy": AOTInductor has no kernel for it. compile_module now turns aten._reshape_copy back into aten.reshape in the parts that run in PyTorch. aten.reshape decomposes into core ATen ops wherever it goes. TensorRT engines are not affected.
shoumikhin
changed the title
github-actions
Bot
added
the
component: lowering
…is causal force_causal_efficient_attention runs by default (attn_bias_is_causal=True) and replaced the attn_bias of every square efficient attention call with is_causal=True, without looking at the bias. A block-causal mask, where each frame of tokens sees itself and earlier frames, then became plain causal attention and the model silently returned wrong values. This shows up when a graph holds complex numbers, for example complex RoPE. The complex lowering retraces the graph, which turns scaled_dot_product_attention with a mask into _scaled_dot_product_efficient_attention with the mask as attn_bias, and the pass then dropped it. With a block-causal mask built from aranges, the compiled model differed from PyTorch by about 1.0. The pass now replaces a bias only when it is proven to be causal: the bias depends on no graph input, so the pass computes it once, and it keeps exactly the lower triangle (0) and blocks the rest (-inf, or the lowest value of the dtype as Hugging Face uses). Any other bias, including one built from an input such as a padding mask, is kept and handled by the attention converter. Tests: a block-causal mask with complex RoPE now matches PyTorch (it failed without this change), and a causal bias built in the graph is still replaced with is_causal=True.
shoumikhin
changed the title
torch.compile runs the lowering passes under a fake mode. There the bias computed by force_causal_efficient_attention has no values, so comparing it raised DataDependentOutputException and the whole graph fell back to PyTorch. Compute and check the bias outside the fake mode, and treat a fake constant as not proven. The causal bias test now runs the pass under a fake mode; it raised without this change.
…ings guide The compilation settings guide still said that attn_bias_is_causal treats every attention bias as causal. Give it the same text as the docstrings, and say in the attention tests which test checks the replacement and which only checks the numbers. Also drop the shape checks in _is_proven_causal that its caller already makes.
github-actions
Bot
added
the
documentation
…eshape aten._reshape_copy always returns new memory, but aten.reshape can return a view. Lowering may already have dropped a user's clone of an input (for example x.clone().view(...) under torch.compile, or inside a torch.cond branch), so the restored reshape could hand back a tensor that shares memory with the caller's input, where PyTorch returns a fresh tensor. Restore aten._reshape_copy as aten.clone(aten.reshape(...)) instead. Both are core ATen ops, so the ExecuTorch CUDA backend still runs it. Test: x.clone().view() to 10 dimensions under torch.compile returns a tensor that does not share memory with the input (it did before).
…orch The partitioners keep a node that makes a uint8 tensor in PyTorch, since TensorRT allows uint8 only at engine inputs and outputs. They exempted every slice_scatter, but only one that overwrites the whole tensor can hand an engine input through. A partial write, such as torch.slice_scatter(image, patch, 0, 2, 4) or self.cache[1:3] = update on a uint8 buffer, still reached TensorRT and failed with "Could not get tensor shape". Exempt only a full overwrite, using the slice_scatter converter's own resolve_slice_scatter_write. A uint8 buffer written on the KV cache dim then failed the buffer lifting check, because _kv_write_will_alias predicted an aliasing KV cache layer for a write the partitioner now keeps in PyTorch. It now returns False for such a write, as it does for an op in torch_executed_ops. Tests: a partial slice_scatter on a uint8 input, and a uint8 buffer written on dim 2, compile and match PyTorch (both failed before).
…not prove The previous change made force_causal_efficient_attention replace only an attention bias it proves causal. Some causal biases cannot be proven: one built from a graph input (such as cache_position), one with a dynamic sequence length, or one from a module buffer under torch.compile, where buffers are graph inputs. Those lost TensorRT's native causal path, which was about twice as fast and more accurate, and there was no way to get it back. attn_bias_is_causal now has three values: - None, the new default: replace only a bias proven causal. - True: also replace a bias the pass cannot prove, which the caller asserts is causal. A bias the pass computes and finds not causal is still kept: its zeros are not exactly the lower triangle, as in a block-causal mask, or it leaves a position above the diagonal unblocked (anything above -1e4, which softmax would not zero out). - False: never replace. The check now runs once per shared mask instead of once per attention node, and a failure inside it, such as running out of memory, keeps the bias instead of stopping the compile. The docs say that a replaced bias is not read again, so refitting the weights it came from does not change the engine. The converter test harness now uses the default, so the tests that pass a bias as an input exercise the converter's bias path again. Tests: a causal bias from an input is kept by default and replaced with attn_bias_is_causal=True, also with a dynamic sequence length, and a causal buffer is replaced by default.
TensorRT holds some 0-d values, such as i + 1, as a 1-D tensor of one element. Used as the only index, after a slice, as in x[:, i + 1] = v, that extra dim became a dim of size 1 in the shape of x[:, i + 1], so the values were laid out against it and every row but the first was lost, with no error. Read the index's rank from the graph, and give a 0-d index held as shape (1,) back its rank.
… count Values with a dim of size 0 fit only an x[indices] with no elements. An earlier change returned the input for such values without building the scatter, so an index that selected elements at runtime, which PyTorch rejects with a shape mismatch, silently passed. Build the scatter again, with the empty values flattened as they are, so TensorRT checks the count at runtime as before. Laying the values out, as for other values, fails the engine build.
…x decomposition test Adding an attention to this test made it compare only the attention output, with a tolerance for TF32, so the exact check of the rotary output was lost. Return both, compare the rotary output exactly again, and use the looser tolerance for the attention only.
The previous change gives a computed 0-d index that TensorRT holds as shape (1,) back its rank. It did that for a 0-d bool mask too, such as x[gate > 0] = 7, and expand_boolean_indices turns a 0-d mask into an empty index, so the write silently did nothing. Restore the rank of integer indices only.
…dcast Values whose size is known only at runtime were broadcast with expand, which reads as many values as the target size needs and drops the rest. So x[:, index] = values[values > 0] with two values for one selected column wrote the first value and dropped the second, while PyTorch raises a shape mismatch. Check at runtime that each such size is 1 or the size it broadcasts to, as TensorRT did before through a reshape.
The Python Linting job runs black --check over the repository and failed on this file: a with statement holding two context managers kept black's older layout. Reformat it with black 26.3.1, the version pinned in .pre-commit-config.yaml. Formatting only.
PyTorch raises when the index tensors of index_put do not broadcast together, or when the values do not broadcast to input[indices]. torch.export does not check these shapes, so the converter has to, and it wrote the values anyway. For example, values [5., 7.] for two rows of three filled both rows with 5., and an index of shape (1, 2) next to one of shape (1, 3) wrote part of the mesh. Sizes known when the engine is built are now checked during conversion and raise a ValueError. Sizes known only at runtime are checked with an assertion layer, so the engine fails where PyTorch raises. The values' rank comes from the graph, since TensorRT can hold a 0-d value as a 1-D tensor of one element, and values with more dims than input[indices] raise. Values with no elements are checked like any others before they take their own path.
…y index The previous change left two gaps. A values size known when the engine is built was not checked against an input size known only at runtime, and expand keeps such a size, so values of shape (2, 3, 2) went into rows that turned out to be (2, 3) at runtime, in the wrong layout, while PyTorch raises. And an index that is empty when the engine is built returned the input before the values were checked at all. Check each such size at runtime too, and check the values before taking the empty-index shortcut, which now covers any index shape with a size of 0.
With one bool mask and a CPU value of one element, PyTorch runs index_put as masked_fill, which takes the value as a scalar whatever its shape. The previous change rejected such a value when it had more dims than input[indices], for example torch.full((1, 1, 1), 7.0) for x[:, mask] with x of shape (2, 1, 1), which built and ran before. Recognize that case from the graph, as ATen does, and use the value as a scalar.
TensorRT keeps the sum of two 0-d tensors 0-d. The 1-D tensor of one element comes from Torch-TensorRT's own converters, which wrap a Python scalar such as the 1 in value + 1 as a 1-D constant. Say so in the comments, and say when that extra dim matters: when input[indices] is 0-d.
PyTorch takes a 0-d bool mask, such as x[gate > 0] = 7 written as an index_put call, only as masked_fill, which fills all of x when the mask is true. The converter treated the mask as a 1-D mask over the first dim, so it wrote only the first row, which was right only when that dim had size 1. Fill all of the input with a select instead.
This reverts commit e7b630c. The select it used fails to build for int8 inputs, where the scatter path builds, and it skipped PyTorch's check for too many indices. A 0-d mask works as it did before this pull request again.
PyTorch raises when a bool mask's shape is not that of the input dims it covers, and when there are more indices than input dims. The converter checked neither. A mask of two for an input of four wrote the first two elements, and indices past the last dim were ignored. Check both from the graph during conversion, and raise a ValueError.
The previous change counted every None toward the input dims. PyTorch's masked_fill takes its one mask and ignores the Nones after it, so index_put(x, [mask, None], value) with a 2-D mask that covers x and a one-element CPU value is valid, and it built and ran before. Stop counting after the mask in that case.
The check that a mask fits the input dims it covers counted an index whose node has no shape metadata as covering one dim. When that index was a mask, the next mask was compared with the wrong input dims, and a valid index_put raised. Stop the check at such an entry, since how many dims it covers is not known.
…ew checks The masked_fill and mask-shape checks also took uint8 tensors as masks. None reach this converter, since its validator admits int32, int64 and bool indices, and masked_fill itself rejects a uint8 mask, so the case only misdescribed PyTorch. Also correct a comment that still said the scatter checks the count of empty values: the checks before it already match every values dim.
The previous change compared the mask dtype with ==, which mypy types as Any, so the masked_fill check failed mypy's no-any-return rule. torch dtypes are singletons, so is gives the same answer as a plain bool.
The index check counted every None toward the input dims and raised when the indices covered more dims than the input has. PyTorch only raises for an index tensor past the last dim. A None there is fine, for example index_put(x, [i, None, mask, None], v) with x of shape (2, 3, 2, 1), a 0-d index i and a mask of shape (2, 1), which built and ran before. Check only index tensors against the dims they cover. That also covers the masked_fill case the previous change handled on its own, so drop that special case.
With truncate_double, a float64 input of a TensorRT engine gets a cast to float32 in front of the engine, and the cast copied the metadata of the node that makes the input. A partition left to PyTorch has no metadata in the parent graph, neither on its call_module node nor on the getitem nodes that read its outputs. So when such a partition made the input, the cast had no metadata either, and torch_tensorrt.save(..., retrace=False), the default for ExecuTorch, failed in the exported program verifier with "Node.meta _to_copy_default is missing val field". This happens, for example, when int64 positions are cast to float64 before more math: float32 cannot hold every int64, so that cast stays in PyTorch and its float64 result goes into the next engine. Take the cast's metadata from the engine's own placeholder for that input, which always has it. Where the old source had metadata, both are copies of the same node's metadata, so those casts get the same metadata as before.
The previous change also raised for an integer index past the last input dim. PyTorch drops such an index when the indices are not next to each other and fails later otherwise, so index_put(x, [mask, None, i], v) with x of shape (2, 1, 1), a mask of shape (2, 1) and a 0-d index i is valid, and it built and ran before. Check only that masks fit, and the count of indices.
… graph says so The values rank check raised unless the graph showed a masked_fill. When the mask's or the input's node has no shape metadata, the graph cannot show it, so a valid masked_fill such as index_put(x, [None, mask], torch.full((1, 1, 1), 7.0)) raised, where it built and ran before. Make the masked_fill test answer "not known" in that case, and raise only when the graph shows that PyTorch would.
The masked_fill test answered "not known" as soon as any of the input, mask or value nodes lacked metadata, even when the rest already ruled masked_fill out, such as an integer index or a CUDA value. Then a value of too high a rank was laid out instead of rejected. It also answered "no" for a CPU value whose size is known only at runtime, such as values[values > 0].reshape(1, 1, -1).cpu(), which may still hold one element, so a valid masked_fill raised where it built and ran before. Check what rules masked_fill out first, and answer "not known" for a value size that is not known.
On TensorRT-RTX, index_nonbool_validator declines every index_put with a bool index, since nonzero is not supported there, so the mask case of test_shapes_that_do_not_broadcast_raise never reaches the converter check it tests. The op runs in PyTorch, which raises only when the module runs, and the test failed with "not raised". Skip that one case on TensorRT-RTX, as the file already does for its other bool index cases, and keep every case on standard TensorRT as strict as before.
…n accumulating index_put With fallback_pytorch_plugin_ops, an accumulating index_put with one 1-D index runs through TensorRT's ScatterElements plugin. The converter used version 2 of that plugin, an IPluginV3. TensorRT 11.4 on Jetson Thor merges IPluginV3 layers into the compiler region around them, so a model that repeats such a layer many times becomes one large region, and the engine build needs far more host memory. In a TensorRT network of 160 mixture-of-experts blocks, each followed by this plugin, the build needed 32.9 GiB more host memory with version 2 and 4.3 GiB with version 1 on Thor. An A100 kept each plugin as its own layer with either version, at 3.6 GiB. Version 1 is the IPluginV2DynamicExt form of the same plugin. Both versions run the same kernel and accept the same types, and initLibNvInferPlugins registers both, so the ExecuTorch delegate loads either one. Version 1 stays its own layer on Thor as well. So the converter now uses version 1 while TensorRT registers it, and version 2 otherwise. Version 1 reads the reduction name up to a NUL, while version 2 reads exactly the field's length, so each gets its own spelling. Nothing changes without fallback_pytorch_plugin_ops. Tests: the ScatterElements conversion test checks that the engine has a layer of the plugin type that the registered versions call for.
…he Thor trade-off The previous change also let the validator accept version 1 of TensorRT's ScatterElements plugin. Every TensorRT release since 10.3 registers both versions, so this let no new graph through, but it made the validator look up version 1 first. TensorRT logs an error for each lookup of a plugin version it does not have, so a TensorRT build without this plugin would log two errors for each accumulating index_put instead of one. The validator checks for version 2 again, exactly as before, and the converter still prefers version 1. The comments said that TensorRT 11.4 merges version 2 on Thor. TensorRT 10.16 does the same, so they now say TensorRT on Thor. The comment now also states the cost of version 1. Keeping each plugin out of the compiler region is what saves the build memory, and it also stops TensorRT from fusing across the plugin, so inference on Thor can be slower: a stack of 16 MLP blocks compiled with Torch-TensorRT ran about 7% slower, and a TensorRT network of 32 mixture-of-experts blocks about 37% slower. Version 1 is deprecated since TensorRT 10.3. When TensorRT drops it, the converter falls back to version 2, and TensorRT then logs that version 1 is missing.
…ejects input shapes When TensorRT rejects the input shapes of an engine, for example because an assertion in the engine fails, IExecutionContext::inferShapes returns a negative count and writes no tensor names. The runtime still printed its list of names into the error message. Those entries were never set, so the message streamed null pointers. With libstdc++ the text was silently cut off after "The shapes of the inputs:". With the Microsoft C++ library the process crashed with an access violation. The index_put converter adds such an assertion for sizes known only at runtime. TensorRT-RTX runs nonzero in PyTorch, so the index sizes reach the engine as input shapes and the assertion fails during shape inference. That is why the index_put test that expects an error for shapes that do not broadcast crashed its worker on Windows with TensorRT-RTX, while it passed on Linux. Now -1 raises an error that names the engine and points to the TensorRT error, and otherwise only the names TensorRT wrote are printed. Tests: an engine with an assertion on its input lengths raises an error that names the engine, with and without the output allocator, which take the two copies of this check. With the old runtime both messages lost the engine name on Linux.
…r before the engine reads it TensorRT reads every input buffer in row-major order. ExecuTorch can plan a tensor in another dim order, for example a program input exported in channels-last order, or the output of the efficient attention op that ExecuTorch's CUDA backend runs, which has dim order (0, 2, 1, 3). When such a tensor fed a TensorRT engine, the runtime bound its memory as it was, so the engine read the elements in the wrong order and returned wrong values with no error. The runtime now checks the strides of each input. An input that is not row-major is copied into row-major order, into the staging buffer that unaligned inputs already use, and the runtime logs this once per engine. The copy is a series of cudaMemcpy2DAsync calls, so the runtime still needs no CUDA kernel of its own. Each row of a copy is a run of elements that is contiguous in both layouts. An outer dim whose stride spans all the rows gathered so far joins the copy's height, so a channels-last input takes one call per channel rather than one per image row. A dim of size 1 moves no element, so its stride is ignored. A row-major input takes the same path as before, with no extra copy. An input that the engine updates in place must be row-major, and any other is refused with a clear error. Before, the update was written back in the wrong order. Writing it back in the caller's order would need a second reordering copy after every call, and is left out. Outputs need no change. The engine op's fake implementation creates its outputs as contiguous tensors, so the exported program plans them in row-major order. An aliased output follows its input, which must now be row-major. Tests: test inputs now carry the dim order and strides that a program gives them. A new engine with a rank-4 input checks four layouts, both the answer and the number of 2-D copies, which a linker wrapper counts as in other tests here: attention's output with one head needs no copy, with three heads one copy per batch and head, channels-last one per batch and channel, and an order whose outer dim cannot be folded one per row. An aliased input in another dim order is refused. With the old runtime the copied layouts give wrong values and the refused call succeeds.
test_failed_engine_assertion_names_the_engine runs an engine that needs the output allocator. The runtime refuses to run such an engine while CUDA graphs mode is on, and raises its own error instead, which does not name the engine. CUDA graphs mode is a process-wide setting, so the test failed whenever an earlier test in the same pytest worker left it on, for example test_pre_allocated_outputs_dynamic. The test now turns CUDA graphs mode off before it runs the engines, and turns it off again when it ends, as the other runtime tests do.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
lerobot-export compiles a trained policy into a bundle that
`lerobot-rollout --bundle.path=<folder>` runs:
lerobot-export \
--policy.path=outputs/train/act_so101/checkpoints/last/pretrained_model \
--backend.type=executorch_tensorrt
Three pieces, each knowing one thing:
- A recipe knows the policy, never the compiler. It turns the policy into
parts (a module, example inputs, and the names of its inputs and outputs)
and makes the test case. The default recipe makes one part that does
what lerobot-rollout does for one chunk: image conversion, the saved
preprocessor, predict_action_chunk trimmed to the actions rollout plays,
and the saved postprocessor. It takes the robot's frame and the starting
noise, and returns actions in robot units. It covers policies that keep
the chunk_inference_spec contract, draw at most one starting noise, and
compute the same chunk twice from the same frame and noise; it refuses
others with the reason. A policy that reads the task holds it, tokenized
once at export. A policy that needs something else adds its own recipe
in export_<type>.py beside its configuration, and the factory finds it
there. The pieces a recipe shares (the robot's frame from the policy or
its dataset, a random frame, the action scale) are plain functions.
- A backend knows the compiler, never the policy. It is a config class
chosen with --backend.type. executorch_tensorrt compiles each part with
Torch-TensorRT into an ExecuTorch program whose inputs and outputs stay
in GPU memory, with TensorRT's TF32 off so float32 stays float32. It
keeps out of TensorRT any operation whose TensorRT plugin would need
PyTorch at run time (Torch-TensorRT's fallback_pytorch_plugin_ops), since
ExecuTorch programs run without PyTorch. That setting comes with
pytorch/TensorRT#4953, not in a release yet, and the backend refuses an
older Torch-TensorRT, which would ignore the setting without a word. An
installed package named lerobot_backend_* adds a backend, the way
lerobot_runtime_* adds a runtime. A backend given with
--backend.discover_packages_path is loaded in every step that needs it.
- The bundle holds the programs, the manifest (written once, after the
last part is compiled), the test case, and the policy config as a record.
lerobot-export itself never imports PyTorch. It traces the policy in one
process, compiles each part in a process of its own, and checks the bundle
in another, so each step starts with the memory the one before it used
given back. The checkpoint loads like in lerobot-rollout, at
--policy.pretrained_revision and with PEFT adapters folded into the base
weights, onto the CPU; the recipe moves what it compiles to the backend's
device. The load is strict: a weight the checkpoint does not fill would
keep its random value, and the export's own check compares the bundle with
that same policy, so it could not tell. PyTorch's TF32 is off in every step, so the test case is the
policy's float32 arithmetic. Tracing and compiling run without autograd,
which nothing there needs and which some of ExecuTorch's passes cannot
record (complex views). The tolerance is in normalized action units:
each action's scale is measured from the saved postprocessor and stored in
the manifest. The bundle is written to a temporary folder, checked in a new
process the way lerobot-rollout checks it, and only then moved to the
output folder. If the actions differ too much, or the check crashes,
nothing is written, and a step that is killed leaves no folder behind.
Torch-TensorRT is not added to an extra: it pins one torch minor version,
which uv's single lock would impose on every extra. ExecuTorch 1.6, which
the runtime needs, and torchao 0.19 are not on PyPI yet, so the `export`
extra is left out of pyproject.toml for now, with a note.
The tests export a tiny ACT through the command with a test backend, added
as a plugin package, and check that the bundle plays the same actions as
SyncInferenceEngine, that the command never imports PyTorch, and that a checkpoint missing a
weight is refused. They also cover the recipe's
task, noise, image, camera and same-chunk rules on small policies, and
check that the plugin setting reaches Torch-TensorRT, that an older
Torch-TensorRT is refused, and that a backend from
--backend.discover_packages_path reaches the compile step. Also tested
by hand on NVIDIA GPUs: ACT, SmolVLA and pi0 exported with
executorch_tensorrt and passed the check.
shoumikhin
added a commit
to shoumikhin/lerobot
that referenced
this pull request
Nov 2, 2026
A policy that starts from a random noise, such as pi0.5, SmolVLA or GR00T, got that noise from the host: the engine drew it with NumPy and copied it to the GPU for every chunk. On a GPU backend the export now adds a small program without inputs that draws the noise on the GPU, and the bundle runs it first. So a chunk copies only the robot's frame to the GPU. The bundle's own check still gives the test case's recorded noise to the same `noise` input and skips the draw, so the check against PyTorch does not change. ExecuTorch 1.6 cannot seed a draw on the GPU, so such a draw cannot be replayed exactly; replays use noise from the host, as the check does. Saving a program without inputs with `executorch_tensorrt` needs pytorch/TensorRT#4953, which is not in a release yet. The docs say so. A test on a GPU counts every copy between the host and the GPU while the bundle runs: each chunk copies only the frame to the GPU and the actions back.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This fixes the problems below. Graphs that compiled before keep their results and, except as noted for
index_put, their engines.Fixes
retrace=False.truncate_double=True, saving failed withmissing val fieldwhen an op left to PyTorch fed an engine float64.image.permute(2, 0, 1)failed to build. Ops that only move data, andwhere, now widen uint8 inside the engine. Other uint8 ops run in PyTorch.int_tensor / 255gave 0 instead of a float.aten._reshape_copy, which the ExecuTorch CUDA backend cannot run. They now becomeaten.cloneandaten.view.force_causal_efficient_attentionno longer replaces a block-causalattn_biaswithis_causal=True(errors near 1.0).offload_module_to_cpu=Truefailed with a device mismatch when a mask came from constants.torch.condbranches and constant shapes, so another graph could reuse an engine. A refused refit now raises.index_putandindexcould write to the wrong places or fail to build. They now follow PyTorch's indexing rule. Shapes that do not broadcast raise. Engines withindex_put, or anexpandbetween runtime sizes, can change, but right results stay the same.fallback_pytorch_plugin_ops=Trueruns them in PyTorch, and an accumulatingindex_puton ScatterElements version 1, which builds with far less host memory on Jetson Thor but can be slower. Wrong ScatterAdd cases now run in PyTorch.InvalidArgument.What was tested
Linux x86_64, NVIDIA A100, TensorRT 11.4: each new test fails without its fix. Other tests are unchanged except the fixed
index_putcases. A block-causal mask, a channels-last input and masked attention in the CUDA backend now match PyTorch. Rejected shapes were tested only on Jetson Thor, ScatterElements on Thor too. Windows was not tested.