feat: improve engine caching and fix bugs

[zewenli98]

Description

As I requested, TensorRT 10.14 added an argument trt.SerializationFlag.INCLUDE_REFIT to allow refitted engines to keep refittable. That means engines can be refitted multiple times. Based on the capability, this PR enhances the existing engine caching and refitting features as follows:

1. To save hard disk space, engine caching will only save weight-stripped engines on disk regardless of compilation_settings.strip_engine_weights. Then, when users pull out the cached engine, it will be automatically refitted and kept refittable.
2. Compiled TRT modules can be refitted multiple times with refit_module_weights(). e.g.:

for _ in range(3):
    trt_gm = refit_module_weights(trt_gm, exp_program)

3. Due to some changes, the insertion and pulling of cached engines are located in different places, which causes 🐛 [Bug] Engine cache failed on torch.compile backend=tensorrt #3909. This PR unified the insertion and pulling in _conversion.py.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)

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]

@cehongwang please take a pass so we have multiple eyes on this PR

[narendasan]

When would a torch.compile use try to use strip weights?

[zewenli98]

added the warning back. Not sure why strip_engine_weights arg doesn't work for torch.compile()

[narendasan]

It just doenst make sense. Torch compile is not serializable. So why would you ever want a callable that doesnt have the weights in it

[narendasan]

I feel like strip weights should only apply to the returned engine and not to the cache directly. So a returned cache engine with strip weights == True wont be refit. but you always only save stripped engine.

[zewenli98]

The current design is to save stripped engine for using less hard disk. A returned cache engine with strip weights == True wont be refit as well. Only strip weights == False will be refit.

[narendasan]

yeah but what i mean is every engine we save in the cache should be a stripped weights engine right? We arent doing weight matching in the hash function for pulling from the cache so we will need to refit anyway. So the serialization config should always have serialization_config.set_flag(trt.SerializationFlag.EXCLUDE_WEIGHTS) regardless if the user tells us settings.strip_engine_weights == True or settings.strip_engine_weights == False

[narendasan]

User Setting	Engine Saved to Cache	Engine Returned to User
settings.strip_engine_weights == True	Weightless	Weightless
settings.strip_engine_weights == False	Weightless	Refit Weights

[narendasan]

Or is it that we dont need to explictly set the setting if the weight was built with stripped weights? Is there any harm in doing so? Then there is only one code path

[narendasan]

Print the hash

[narendasan]

I think this code is unclear. Would recommend something like this

hash_val = engine_cache.get_hash(module, inputs, settings) if (settings.cache_built_engines or settings.reuse_cached_engines) else None

if settings.reuse_cached_engines:
    serialized_interpreter_result = pull_cached_engine(
        hash_val, module, engine_cache, settings, inputs
    )
    if serialized_interpreter_result is not None:  # hit the cache
        return serialized_interpreter_result

...

if (
    ENABLED_FEATURES.refit
    and not settings.immutable_weights
    and settings.cache_built_engines
    and engine_cache is not None
):
    _ = insert_engine_to_cache(
        hash_val, interpreter_result, engine_cache, settings, inputs
    )

    serialized_engine = interpreter_result.engine.serialize()
    

[narendasan]

We should probably throw a warning or something if engine_cache is None and settings.cache_built_engines == True