Skip to content

Cache post-lowering ATen graphs to skip decomp on warm compiles. #4540

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
micwill755 wants to merge 1 commit into from
Draft

Cache post-lowering ATen graphs to skip decomp on warm compiles. #4540

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 67 additions & 19 deletions py/torch_tensorrt/dynamo/_compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@
)
from torch_tensorrt.dynamo._engine_cache import BaseEngineCache, DiskEngineCache
from torch_tensorrt.dynamo._exporter import replace_execute_engine_no_op_node
from torch_tensorrt.dynamo._lowering_cache import (
DiskLoweringCache,
LoweringCacheEntry,
repropagate_graph_metadata,
)
from torch_tensorrt.dynamo.conversion import (
CompilationSettings,
UnsupportedOperatorException,
Expand Down Expand Up @@ -455,6 +460,9 @@ def compile(
engine_cache_dir: str = _defaults.ENGINE_CACHE_DIR,
engine_cache_size: int = _defaults.ENGINE_CACHE_SIZE,
custom_engine_cache: Optional[BaseEngineCache] = _defaults.CUSTOM_ENGINE_CACHE,
cache_lowered_graphs: bool = _defaults.CACHE_LOWERED_GRAPHS,
reuse_cached_lowered_graphs: bool = _defaults.REUSE_CACHED_LOWERED_GRAPHS,
lowering_cache_dir: str = _defaults.LOWERING_CACHE_DIR,
use_fp32_acc: bool = _defaults.USE_FP32_ACC,
refit_identical_engine_weights: bool = _defaults.REFIT_IDENTICAL_ENGINE_WEIGHTS,
strip_engine_weights: bool = _defaults.STRIP_ENGINE_WEIGHTS,
Expand Down Expand Up @@ -549,6 +557,9 @@ def compile(
engine_cache_dir (str): Directory to store the cached TRT engines
engine_cache_size (int): Maximum hard-disk space (bytes) to use for the engine cache, default is 1GB. If the cache exceeds this size, the oldest engines will be removed by default
custom_engine_cache (Optional[BaseEngineCache]): Engine cache instance to use for saving and loading engines. Users can provide their own engine cache by inheriting from BaseEngineCache. If used, engine_cache_dir and engine_cache_size will be ignored.
cache_lowered_graphs (bool): Whether to save lowered and partitioned graphs for warm compilation
reuse_cached_lowered_graphs (bool): Whether to reuse lowered and partitioned graphs before decomposition
lowering_cache_dir (str): Directory used by the lowered graph cache
use_fp32_acc (bool): Enable FP32 accumulation for FP16 matmul layers while retaining FP16
inputs and outputs. When combined with ``decompose_attention=True``, the complete
decomposed FP16 scaled dot product attention calculation runs in FP32 and only its
Expand Down Expand Up @@ -743,6 +754,8 @@ def compile(
"lazy_engine_init": lazy_engine_init,
"cache_built_engines": cache_built_engines,
"reuse_cached_engines": reuse_cached_engines,
"cache_lowered_graphs": cache_lowered_graphs,
"reuse_cached_lowered_graphs": reuse_cached_lowered_graphs,
"use_fp32_acc": use_fp32_acc,
"refit_identical_engine_weights": refit_identical_engine_weights,
"strip_engine_weights": strip_engine_weights,
Expand Down Expand Up @@ -772,25 +785,51 @@ def compile(

logger.info("Compilation Settings: %s\n", settings)
exported_program = pre_export_lowering(exported_program, settings)
exported_program = exported_program.run_decompositions(
get_decompositions(
enable_experimental_decompositions,
decompose_attention,
use_distributed_mode_trace,
use_fp32_acc=use_fp32_acc,
lowering_cache = None
lowering_cache_key = None
lowering_cache_entry = None
if cache_lowered_graphs or reuse_cached_lowered_graphs:
if not DiskLoweringCache.can_cache(settings):
logger.warning(
"Lowering cache bypassed: the initial implementation requires "
"require_full_compilation=True, use_fast_partitioner=True, dryrun=False, "
"enable_autocast=False, and use_distributed_mode_trace=False"
)
else:
lowering_cache = DiskLoweringCache(lowering_cache_dir)
lowering_cache_key = lowering_cache.get_hash(
exported_program, trt_arg_inputs, trt_kwarg_inputs, settings
)
if reuse_cached_lowered_graphs:
lowering_cache_entry = lowering_cache.load(lowering_cache_key)

if lowering_cache_entry is not None:
gm = lowering_cache_entry.lowered_module
lifted_buffers = list(lowering_cache_entry.lifted_buffers)
gm = repropagate_graph_metadata(
gm, trt_arg_inputs, trt_kwarg_inputs, settings.device
)
logger.info(
"Reusing cached lowered graph; skipping decomposition and post-lowering"
)
else:
exported_program = exported_program.run_decompositions(
get_decompositions(
enable_experimental_decompositions,
decompose_attention,
use_distributed_mode_trace,
use_fp32_acc=use_fp32_acc,
)
)
)

gm = exported_program.module()
# Move the weights in the state_dict to CPU
logger.debug("Input graph: " + str(gm.graph))
gm = exported_program.module()
logger.debug("Input graph: " + str(gm.graph))

# Lift mutated buffers from get_attr to placeholders BEFORE post_lowering's
# constant_fold runs, so the engine sees them as input bindings.
gm, lifted_buffers = lift_mutated_buffers(gm)
gm = post_lowering(gm, settings)

# Lift mutated buffers from get_attr to placeholders BEFORE post_lowering's
# constant_fold runs, so the engine sees them as input bindings (a
# prerequisite for IKVCacheUpdateLayer / aliased I/O to fire on a
# module-held cache). Returns a fresh GraphModule whose forward signature
# reflects the new placeholders.
gm, lifted_buffers = lift_mutated_buffers(gm)
if lifted_buffers:
# Append each lifted buffer as an engine input AFTER the user inputs.
# Buffer tensors live on the gm's state; prepare an Input spec for
Expand All @@ -804,9 +843,17 @@ def compile(
[b for _, b, _ in lifted_buffers],
)

# Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires
# module parameters to still be on GPU, so we must not deallocate before this call.
gm = post_lowering(gm, settings)
if (
lowering_cache_entry is None
and cache_lowered_graphs
and lowering_cache is not None
and lowering_cache_key is not None
):
gm = lowering_cache.save(
lowering_cache_key,
LoweringCacheEntry(gm, tuple(lifted_buffers)),
)

logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
logger.debug("Lowered Input graph: " + str(gm.graph))

Expand All @@ -823,6 +870,7 @@ def compile(
logger.warning(
"Remaining GPU memory may not be enough to compile the TensorRT engine for this model resulting in an OOM error, Consider setting offload_module_to_cpu=True"
)

trt_gm = compile_module(
gm,
trt_arg_inputs,
Expand Down
3 changes: 3 additions & 0 deletions py/torch_tensorrt/dynamo/_defaults.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@
ENGINE_CACHE_DIR = os.path.join(tempfile.gettempdir(), "torch_tensorrt_engine_cache")
ENGINE_CACHE_SIZE = 5368709120 # 5GB
CUSTOM_ENGINE_CACHE = None
CACHE_LOWERED_GRAPHS = False
REUSE_CACHED_LOWERED_GRAPHS = False
LOWERING_CACHE_DIR = os.path.join(ENGINE_CACHE_DIR, "lowered_graphs")
USE_FP32_ACC = False
REFIT_IDENTICAL_ENGINE_WEIGHTS = False
STRIP_ENGINE_WEIGHTS = False
Expand Down
Loading
Loading