❓ [Question] how to reduce peak memory usage when loading a Torch-TensorRT module

[patrick-botco]

❓ Question

hey team,

happy new year!

i am looking to reduce peak memory usage when loading a Torch-TensorRT module. given an ExportedProgram, is there something that i can be doing that is more optimal? this is on a Jetson.

What you have already tried

1. compiled_model = torch_tensorrt.dynamo.compile(...), using https://docs.pytorch.org/TensorRT/contributors/resource_management.html
2. torch.save(compiled_model, path, pickle_protocol=5)
3. torch.load(path, weights_only=False) (cold start, nothing else running)

the result is really high peak memory (a few times more than the serialized engine size).

thank you!! cc @peri044 @lanluo-nvidia @narendasan

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• PyTorch Version (e.g., 1.0):
• CPU Architecture:
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, libtorch, source):
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version:
• CUDA version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[narendasan]

@cehongwang can you provide some advice here?

[cehongwang]

Can you share your config inside compiled_model = torch_tensorrt.dynamo.compile(...) with us?

[patrick-botco]

@cehongwang pretty vanilla - single graph, no torch executed ops:

    enabled_precisions: Set[torch.dtype] = {torch.float, torch.bfloat16},
    workspace_size: int = 4 << 30,
    min_block_size: int = 7,
    offload_module_to_cpu: bool = False,

in particular, no refitting ("immutable_weights" is disabled). is there anything in particular you are looking for?

[cehongwang]

Did you also use export TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=1 before executing the script?

Also, just want to clarify: the result is really high peak memory do you mean in the cold start the system memory usage is really high? Or it's the compilation process?

Can you send me the repro script if it's possible?

[patrick-botco]

Did you also use export TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=1 before executing the script?

Also, just want to clarify: the result is really high peak memory do you mean in the cold start the system memory usage is really high? Or it's the compilation process?

Can you send me the repro script if it's possible?

yea im setting TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM.

really high peak memory when loading a compiled model (compile the model -> exit -> load the compiled model from disk (high peak memory here))

ill try to get you a repro script asap