Repository navigation
feat: Add support for Groot N1.5 model - #3736
Merged
Merged
Conversation
Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
github-actions
Bot
added
component: lowering
peri044
marked this pull request as draft
July 31, 2025 14:43
Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
narendasan
approved these changes
Aug 19, 2025
narendasan
left a comment
narendasan
left a comment
Collaborator
There was a problem hiding this comment.
Seems mostly good to go to me
| _enums.dtype._from(input.dtype).to(torch.dtype), | ||
| _enums.dtype._from(other.dtype).to(torch.dtype), | ||
| ) | ||
| if ( |
Collaborator
There was a problem hiding this comment.
Can we start noting down these specific type behavior things somewhere, like in the contributor docs?
peri044
marked this pull request as ready for review
August 20, 2025 18:43
Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
github-actions
Bot
added
the
documentation
Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>
tp5uiuc
added a commit
to tp5uiuc/TensorRT
that referenced
this pull request
Sep 2, 2026
…ule offloads
A hybrid graph module does not own its weights. torch.export.export keeps the source
module's Parameter objects rather than copying them, ExportedProgram.module() re-registers
those same objects, and both partitioners share attributes by setattr -- so the compiled
module and the module it was compiled from hold the identical Parameter instances, and
nn.Module.to() rebinds .data on them. Moving one moves the other.
deallocate_module() is a .to("cpu"), and MutableTorchTensorRTModule called it
unconditionally in three places. Any graph the partitioner did not hand to TensorRT in full
therefore had its PyTorch-executed weights dragged off the device while their activations
stayed on it:
RuntimeError: Expected all tensors to be on the same device, but got mat1 is on
cuda:0, different from other tensors on cpu (wrapper_CUDA_addmm)
The TensorRT submodules can spare their weights -- those live inside the engine -- which is
why a fully converted graph never noticed.
This is not a recent regression. pytorch#3418 added the same unconditional deallocate at all four
call sites; pytorch#3736 later put it behind the user-facing offload_module_to_cpu setting in
compile(), and only there. update_refit_condition(), refit_gm() and the static load() kept
offloading on every call regardless of what the user asked for. It surfaced now during a
triage on a card whose capability guards push addmm into PyTorch, which makes partial
fallback the common case; forcing the same fallback with torch_executed_ops reproduces it
byte-for-byte on an L40S, so it is neither card- nor backend-specific.
* Both halves are fixed together, because either alone leaves the bug reachable. The
three missed sites now consult offload_module_to_cpu the way compile() already did,
which is what stops an offload nobody asked for. And every offload now puts the
PyTorch-executed submodules back on the target device afterwards, mirroring what
compile_module already does at the end of a compile -- without that,
offload_module_to_cpu=True still strands a partial-fallback model.
* All three unguarded sites could strand weights, each by a different route.
update_refit_condition() runs the source model, offloads, and on matching output marks
the module LIVE, so the next forward hits CPU weights with no refit having run at all.
refit_gm() offloads after an in-place refit that deliberately left those weights on the
device; on second and later refits it offloads weights it never hoisted, since only
exp_program.module() is moved to the device there. load() offloads after re-exporting,
and nothing afterwards re-creates the graph module or re-attaches its parameters --
save() pickles the source model and the graph module together, so pickle memoisation
preserves the shared-Parameter identity across the round trip.
* Rejected alternatives. Gating alone is a smaller change but leaves
offload_module_to_cpu=True broken for every partial-fallback model. Teaching
deallocate_module() to skip the parameters a graph module still owns is the most
general, but it is a shared utility with eight call sites and would duplicate the
placement compile_module already performs.
* Not fixed here: with use_fast_partitioner=False the PyTorch-executed nodes stay in the
parent graph and nothing ever places them on the device, so offload_module_to_cpu=True
strands them in torch_tensorrt.compile() too. That is a separate, pre-existing gap in
the compiler rather than in this module, and the helper deliberately mirrors the
placement compile_module does instead of half-covering it.
Testing (L40S / SM 8.9 and T4 / SM 7.5, driver 595.58.03, identical stacks):
* test_mutable_torchtrt_module.py with -n 0. On SM 8.9, with addmm and linear forced
into PyTorch it goes from 6 passed / 5 failed to 13 passed; with no forcing at all it
goes from 11 passed to 13 passed. The five failures before the change are test_resnet18,
test_save, test_resnet18_modify_attribute, test_custom_model_with_kwarg and
test_custom_model_with_inplace_init, every one of them on the mixed-device addmm.
* Two new tests. Both force the fallback with torch_executed_ops so they run on any GPU,
refit, and then assert the device of every parameter and buffer rather than only that
the forward did not throw. Each also asserts the graph really did keep something in
PyTorch, so a future lowering change cannot quietly make them vacuous -- that guard
fired during development, when addmm alone turned out not to fall back. Run against the
unfixed module both fail with the reported error.
* The offload still offloads. With offload_module_to_cpu=True, everything TensorRT
absorbed is on the CPU after a refit and only what the graph module still executes in
PyTorch stays on the device. That is asserted by the second new test, not just observed.
* Full runtime/ module on both arms, every test accounted for (pass + fail + skip ==
collected, 220 on each). Zero pass -> fail and zero pass -> skip on either arm.
SM 8.9: every one of the 218 previously recorded tests keeps its status exactly, plus
the two new tests passing. SM 7.5: the nine tests in this file that were failing all
pass, five of them because of this change and four that other changes already on this
branch had fixed; the sweep also shows five weight-streaming tests recovering, which
this commit cannot have touched -- it changes one module that test file never uses.
The three failures left on each arm are test_aliased_io x2 and
test_dynamic_workspace_allocation, identical on both arms and unrelated to this work.
tp5uiuc
added a commit
to tp5uiuc/TensorRT
that referenced
this pull request
Sep 5, 2026
…ule offloads
A hybrid graph module does not own its weights. torch.export.export keeps the source
module's Parameter objects rather than copying them, ExportedProgram.module() re-registers
those same objects, and both partitioners share attributes by setattr -- so the compiled
module and the module it was compiled from hold the identical Parameter instances, and
nn.Module.to() rebinds .data on them. Moving one moves the other.
deallocate_module() is a .to("cpu"), and MutableTorchTensorRTModule called it
unconditionally in three places. Any graph the partitioner did not hand to TensorRT in full
therefore had its PyTorch-executed weights dragged off the device while their activations
stayed on it:
RuntimeError: Expected all tensors to be on the same device, but got mat1 is on
cuda:0, different from other tensors on cpu (wrapper_CUDA_addmm)
The TensorRT submodules can spare their weights -- those live inside the engine -- which is
why a fully converted graph never noticed.
This is not a recent regression. pytorch#3418 added the same unconditional deallocate at all four
call sites; pytorch#3736 later put it behind the user-facing offload_module_to_cpu setting in
compile(), and only there. update_refit_condition(), refit_gm() and the static load() kept
offloading on every call regardless of what the user asked for. It surfaced now during a
triage on a card whose capability guards push addmm into PyTorch, which makes partial
fallback the common case; forcing the same fallback with torch_executed_ops reproduces it
byte-for-byte on an L40S, so it is neither card- nor backend-specific.
* Both halves are fixed together, because either alone leaves the bug reachable. The
three missed sites now consult offload_module_to_cpu the way compile() already did,
which is what stops an offload nobody asked for. And every offload now puts the
PyTorch-executed submodules back on the target device afterwards, mirroring what
compile_module already does at the end of a compile -- without that,
offload_module_to_cpu=True still strands a partial-fallback model.
* All three unguarded sites could strand weights, each by a different route.
update_refit_condition() runs the source model, offloads, and on matching output marks
the module LIVE, so the next forward hits CPU weights with no refit having run at all.
refit_gm() offloads after an in-place refit that deliberately left those weights on the
device; on second and later refits it offloads weights it never hoisted, since only
exp_program.module() is moved to the device there. load() offloads after re-exporting,
and nothing afterwards re-creates the graph module or re-attaches its parameters --
save() pickles the source model and the graph module together, so pickle memoisation
preserves the shared-Parameter identity across the round trip.
* Rejected alternatives. Gating alone is a smaller change but leaves
offload_module_to_cpu=True broken for every partial-fallback model. Teaching
deallocate_module() to skip the parameters a graph module still owns is the most
general, but it is a shared utility with eight call sites and would duplicate the
placement compile_module already performs.
* Not fixed here: with use_fast_partitioner=False the PyTorch-executed nodes stay in the
parent graph and nothing ever places them on the device, so offload_module_to_cpu=True
strands them in torch_tensorrt.compile() too. That is a separate, pre-existing gap in
the compiler rather than in this module, and the helper deliberately mirrors the
placement compile_module does instead of half-covering it.
Testing (L40S / SM 8.9 and T4 / SM 7.5, driver 595.58.03, identical stacks):
* test_mutable_torchtrt_module.py with -n 0. On SM 8.9, with addmm and linear forced
into PyTorch it goes from 6 passed / 5 failed to 13 passed; with no forcing at all it
goes from 11 passed to 13 passed. The five failures before the change are test_resnet18,
test_save, test_resnet18_modify_attribute, test_custom_model_with_kwarg and
test_custom_model_with_inplace_init, every one of them on the mixed-device addmm.
* Two new tests. Both force the fallback with torch_executed_ops so they run on any GPU,
refit, and then assert the device of every parameter and buffer rather than only that
the forward did not throw. Each also asserts the graph really did keep something in
PyTorch, so a future lowering change cannot quietly make them vacuous -- that guard
fired during development, when addmm alone turned out not to fall back. Run against the
unfixed module both fail with the reported error.
* The offload still offloads. With offload_module_to_cpu=True, everything TensorRT
absorbed is on the CPU after a refit and only what the graph module still executes in
PyTorch stays on the device. That is asserted by the second new test, not just observed.
* Full runtime/ module on both arms, every test accounted for (pass + fail + skip ==
collected, 220 on each). Zero pass -> fail and zero pass -> skip on either arm.
SM 8.9: every one of the 218 previously recorded tests keeps its status exactly, plus
the two new tests passing. SM 7.5: the nine tests in this file that were failing all
pass, five of them because of this change and four that other changes already on this
branch had fixed; the sweep also shows five weight-streaming tests recovering, which
this commit cannot have touched -- it changes one module that test file never uses.
The three failures left on each arm are test_aliased_io x2 and
test_dynamic_workspace_allocation, identical on both arms and unrelated to this work.
tp5uiuc
added a commit
to tp5uiuc/TensorRT
that referenced
this pull request
Sep 16, 2026
…ule offloads
A hybrid graph module does not own its weights. torch.export.export keeps the source
module's Parameter objects rather than copying them, ExportedProgram.module() re-registers
those same objects, and both partitioners share attributes by setattr -- so the compiled
module and the module it was compiled from hold the identical Parameter instances, and
nn.Module.to() rebinds .data on them. Moving one moves the other.
deallocate_module() is a .to("cpu"), and MutableTorchTensorRTModule called it
unconditionally in three places. Any graph the partitioner did not hand to TensorRT in full
therefore had its PyTorch-executed weights dragged off the device while their activations
stayed on it:
RuntimeError: Expected all tensors to be on the same device, but got mat1 is on
cuda:0, different from other tensors on cpu (wrapper_CUDA_addmm)
The TensorRT submodules can spare their weights -- those live inside the engine -- which is
why a fully converted graph never noticed.
This is not a recent regression. pytorch#3418 added the same unconditional deallocate at all four
call sites; pytorch#3736 later put it behind the user-facing offload_module_to_cpu setting in
compile(), and only there. update_refit_condition(), refit_gm() and the static load() kept
offloading on every call regardless of what the user asked for. It surfaced now during a
triage on a card whose capability guards push addmm into PyTorch, which makes partial
fallback the common case; forcing the same fallback with torch_executed_ops reproduces it
byte-for-byte on an L40S, so it is neither card- nor backend-specific.
* Both halves are fixed together, because either alone leaves the bug reachable. The
three missed sites now consult offload_module_to_cpu the way compile() already did,
which is what stops an offload nobody asked for. And every offload now puts the
PyTorch-executed submodules back on the target device afterwards, mirroring what
compile_module already does at the end of a compile -- without that,
offload_module_to_cpu=True still strands a partial-fallback model.
* All three unguarded sites could strand weights, each by a different route.
update_refit_condition() runs the source model, offloads, and on matching output marks
the module LIVE, so the next forward hits CPU weights with no refit having run at all.
refit_gm() offloads after an in-place refit that deliberately left those weights on the
device; on second and later refits it offloads weights it never hoisted, since only
exp_program.module() is moved to the device there. load() offloads after re-exporting,
and nothing afterwards re-creates the graph module or re-attaches its parameters --
save() pickles the source model and the graph module together, so pickle memoisation
preserves the shared-Parameter identity across the round trip.
* Rejected alternatives. Gating alone is a smaller change but leaves
offload_module_to_cpu=True broken for every partial-fallback model. Teaching
deallocate_module() to skip the parameters a graph module still owns is the most
general, but it is a shared utility with eight call sites and would duplicate the
placement compile_module already performs.
* Not fixed here: with use_fast_partitioner=False the PyTorch-executed nodes stay in the
parent graph and nothing ever places them on the device, so offload_module_to_cpu=True
strands them in torch_tensorrt.compile() too. That is a separate, pre-existing gap in
the compiler rather than in this module, and the helper deliberately mirrors the
placement compile_module does instead of half-covering it.
Testing (L40S / SM 8.9 and T4 / SM 7.5, driver 595.58.03, identical stacks):
* test_mutable_torchtrt_module.py with -n 0. On SM 8.9, with addmm and linear forced
into PyTorch it goes from 6 passed / 5 failed to 13 passed; with no forcing at all it
goes from 11 passed to 13 passed. The five failures before the change are test_resnet18,
test_save, test_resnet18_modify_attribute, test_custom_model_with_kwarg and
test_custom_model_with_inplace_init, every one of them on the mixed-device addmm.
* Two new tests. Both force the fallback with torch_executed_ops so they run on any GPU,
refit, and then assert the device of every parameter and buffer rather than only that
the forward did not throw. Each also asserts the graph really did keep something in
PyTorch, so a future lowering change cannot quietly make them vacuous -- that guard
fired during development, when addmm alone turned out not to fall back. Run against the
unfixed module both fail with the reported error.
* The offload still offloads. With offload_module_to_cpu=True, everything TensorRT
absorbed is on the CPU after a refit and only what the graph module still executes in
PyTorch stays on the device. That is asserted by the second new test, not just observed.
* Full runtime/ module on both arms, every test accounted for (pass + fail + skip ==
collected, 220 on each). Zero pass -> fail and zero pass -> skip on either arm.
SM 8.9: every one of the 218 previously recorded tests keeps its status exactly, plus
the two new tests passing. SM 7.5: the nine tests in this file that were failing all
pass, five of them because of this change and four that other changes already on this
branch had fixed; the sweep also shows five weight-streaming tests recovering, which
this commit cannot have touched -- it changes one module that test file never uses.
The three failures left on each arm are test_aliased_io x2 and
test_dynamic_workspace_allocation, identical on both arms and unrelated to this work.
tp5uiuc
added a commit
to tp5uiuc/TensorRT
that referenced
this pull request
Oct 6, 2026
…ule offloads
A hybrid graph module does not own its weights. torch.export.export keeps the source
module's Parameter objects rather than copying them, ExportedProgram.module() re-registers
those same objects, and both partitioners share attributes by setattr -- so the compiled
module and the module it was compiled from hold the identical Parameter instances, and
nn.Module.to() rebinds .data on them. Moving one moves the other.
deallocate_module() is a .to("cpu"), and MutableTorchTensorRTModule called it
unconditionally in two places. Any graph the partitioner did not hand to TensorRT in full
therefore had its PyTorch-executed weights dragged off the device while their activations
stayed on it:
RuntimeError: Expected all tensors to be on the same device, but got mat1 is on
cuda:0, different from other tensors on cpu (wrapper_CUDA_addmm)
The TensorRT submodules can spare their weights -- those live inside the engine -- which is
why a fully converted graph never noticed.
This is not a recent regression. pytorch#3418 added the same unconditional deallocate at all four
call sites; pytorch#3736 later put it behind the user-facing offload_module_to_cpu setting in
compile(), and only there. refit_gm() and the static load() kept offloading on every call
regardless of what the user asked for. pytorch#4664 later removed the third call site, in
update_refit_condition(). It surfaced now during a
triage on a card whose capability guards push addmm into PyTorch, which makes partial
fallback the common case; forcing the same fallback with torch_executed_ops reproduces it
byte-for-byte on an L40S, so it is neither card- nor backend-specific.
* Both halves are fixed together, because either alone leaves the bug reachable. The
two missed sites now consult offload_module_to_cpu the way compile() already did,
which is what stops an offload nobody asked for. And every offload now puts the
PyTorch-executed submodules back on the target device afterwards, mirroring what
compile_module already does at the end of a compile -- without that,
offload_module_to_cpu=True still strands a partial-fallback model.
* Both unguarded sites could strand weights, each by a different route.
refit_gm() offloads after an in-place refit that deliberately left those weights on the
device; on second and later refits it offloads weights it never hoisted, since only
exp_program.module() is moved to the device there. load() offloads after re-exporting,
and nothing afterwards re-creates the graph module or re-attaches its parameters --
save() pickles the source model and the graph module together, so pickle memoisation
preserves the shared-Parameter identity across the round trip.
* Rejected alternatives. Gating alone is a smaller change but leaves
offload_module_to_cpu=True broken for every partial-fallback model. Teaching
deallocate_module() to skip the parameters a graph module still owns is the most
general, but it is a shared utility with eight call sites and would duplicate the
placement compile_module already performs.
Testing (L40S / SM 8.9 and T4 / SM 7.5, driver 595.58.03, identical stacks):
* test_mutable_torchtrt_module.py with -n 0. On SM 8.9, with addmm and linear forced
into PyTorch it goes from 6 passed / 5 failed to 13 passed; with no forcing at all it
goes from 11 passed to 13 passed. The five failures before the change are test_resnet18,
test_save, test_resnet18_modify_attribute, test_custom_model_with_kwarg and
test_custom_model_with_inplace_init, every one of them on the mixed-device addmm.
* Two new tests. Both force the fallback with torch_executed_ops so they run on any GPU,
refit, and then assert the device of every parameter and buffer rather than only that
the forward did not throw. Each also asserts the graph really did keep something in
PyTorch, so a future lowering change cannot quietly make them vacuous -- that guard
fired during development, when addmm alone turned out not to fall back. Run against the
unfixed module both fail with the reported error.
* The offload still offloads. With offload_module_to_cpu=True, everything TensorRT
absorbed is on the CPU after a refit and only what the graph module still executes in
PyTorch stays on the device. That is asserted by the second new test, not just observed.
* Full runtime/ module on both arms, every test accounted for (pass + fail + skip ==
collected, 220 on each). Zero pass -> fail and zero pass -> skip on either arm.
SM 8.9: every one of the 218 previously recorded tests keeps its status exactly, plus
the two new tests passing. SM 7.5: the nine tests in this file that were failing all
pass, five of them because of this change and four that other changes already on this
branch had fixed; the sweep also shows five weight-streaming tests recovering, which
this commit cannot have touched -- it changes one module that test file never uses.
The three failures left on each arm are test_aliased_io x2 and
test_dynamic_workspace_allocation, identical on both arms and unrelated to this work.
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.
Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: