fix: support ITensor amex in quantize and add export_torch_mode() wrapper

[zewenli98]

Description

1. support ITensor amex in quantize
2. add export_torch_mode() wrapper in torch_tensorrt.compile to modelopt quant models

Fixes #4429 #2976

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

Can you explain what this context is and maybe add some tests?

[zewenli98]

Can you explain what this context is and maybe add some tests?

1. In quantize(), amex could be trt.ITensor which doesn't have amax.numel()
2. ModelOpt requires with export_torch_mode() wrapper to make inserted quantizers traceable while exporting pytorch models. See: torch_tensorrt.compile() fails with "fake tensor in exported program constant's list" for ModelOpt quantized models unless wrapped in export_torch_mode() #4429

[narendasan]

Can you explain what this context is and maybe add some tests?

1. In quantize(), `amex` could be trt.ITensor which doesn't have `amax.numel()`

2. ModelOpt requires `with export_torch_mode()` wrapper to make inserted quantizers traceable while exporting pytorch models. See: [torch_tensorrt.compile() fails with "fake tensor in exported program constant's list" for ModelOpt quantized models unless wrapped in export_torch_mode() #4429](https://github.com/pytorch/TensorRT/issues/4429)

I see, is there any harm in always having this context? should we catch cases where the modelopt is missing but the context is needed and throw an error/warning?

[zewenli98]

I see, is there any harm in always having this context? should we catch cases where the modelopt is missing but the context is needed and throw an error/warning?

I remember ModelOpt is not required in torch-trt. If we need to check all models with ModelOpt API, the package should be required rather than opt in.

[narendasan]

Should we throw a warning if its a quantized model but model opt is not available? Does this matter for torchao or other sources?

[zewenli98]

well since the function is_quantized is imported from modelopt, it must be available at this point. If we want to throw errors/warnings, I think we might need to write a helper function to detect quantizers to override is_quantized. For other sources like torchao, if they insert quantizers or similar indicators, we can use the same mechanism.

[zewenli98]

@narendasan I optimized codes so that if modelopt is not installed, it will fallback to check module names. And it will error out for quantized model if modelopt is requirement. Also, it's easy to generalize for other sources like torchao.

[narendasan]

Lgtm