CPU Memory Consumption Reduction

[narendasan]

[cehongwang]

In Qwen, TRT builder uses 1x to build live engine.
Sometimes TRT builder uses more than 2x even for building live engine: FLUX, GPT-OSS

[cehongwang]

INetworkDefinition does not take any memory actually, it is the lowered graph and constand folding takes the up to 1x (0-1x) memory. Code here:

TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

Line 39 in 763fd08

replace_node_with_constant(

INetwork is just hold the reference to the weights in the lowered graph

[cehongwang]

Summary

Add four opt-in memory modes to make Torch-TensorRT predictable and usable on constrained machines. Modes are selectable via Python API and env vars.

Modes

1) standard (default)

Behavior: No extra memory optimization.

Lifecycle: Models/engines stay resident on GPU for compile + run.

Use when: You have ample CPU and GPU memory and want a balance between CPU and GPU memory consumption.

Consumption: CPU memory would use ~4x of the model size. GPU will use 2x of the model size.

2) low_CPU_ram

Goal: Cut host memory consumption during compile/run.

Techniques:

After engine building, we use malloc_trim to release memory. Before serialization, the CPU memory consumption is reduced to a minimum (less than 1x), and serialization takes up to 2x. So the peak memory usage will be less than 3x in general.

Use when: you have enough GPU memory (>2x model size) and limited CPU memory (<3x model size).
Consumption: CPU memory would use 2-3x of the model size. GPU will use 2x of the model size.

Risk: The stablilty of malloc_trim is experimental

3) low_GPU_vRAM

Goal: Lower peak GPU VRAM.

Techniques:

1. Lazy engine initialization for graphs with breaks (build when first used).
2. Engine offload to CPU memory before building the engine.

Use when: you don't have enough GPU memory (<2x model size) and enough CPU memory (>5x model size).
Consumption: CPU memory would use 5x of the model size. GPU will use 1x of the model size.

4) all_on

Goal: Run under tight CPU & GPU budgets.

Techniques: low_CPU_ram + low_GPU_vRAM combined
Use when: you don't have enough GPU memory (<2x model size) and don't have much CPU memory (<5x model size).
Consumption: CPU memory would use 3-4x of the model size. GPU will use 1x of the model size.

[narendasan]

Is all on conflicting? Like what is the expected consumption if cpu offload and we do malloc trim?