File tree Expand file tree Collapse file tree 2 files changed +17
-11
lines changed
py/torch_tensorrt/dynamo/backend Expand file tree Collapse file tree 2 files changed +17
-11
lines changed Original file line number Diff line number Diff line change 10
10
from torch ._dynamo .backends .common import aot_autograd
11
11
from torch ._dynamo .utils import detect_fake_mode
12
12
from torch ._functorch .aot_autograd import aot_export_joint_simple
13
- from torch .distributed .tensor import DTensor
14
13
from torch_tensorrt .dynamo import CompilationSettings
15
14
from torch_tensorrt .dynamo ._compiler import compile_module
16
15
from torch_tensorrt .dynamo .lowering import (
@@ -89,6 +88,11 @@ def aot_torch_tensorrt_aten_backend(
89
88
logger .warning (
90
89
"It is recommended to run the model with use_distributed_mode_trace = True since there are distributed tensors in the input which is not supported in aot_export_joint_simple"
91
90
)
91
+
92
+ if settings .offload_module_to_cpu :
93
+ logger .warning (
94
+ "`offload_module_to_cpu` is not supported for `torch_compile` backend."
95
+ )
92
96
return _pretraced_backend (gm , sample_inputs , settings , engine_cache )
93
97
94
98
Original file line number Diff line number Diff line change @@ -79,11 +79,12 @@ def test_resnet18_cpu_offload(ir):
79
79
}
80
80
81
81
trt_mod = torchtrt .compile (model , ** compile_spec )
82
- assertions .assertTrue (
83
- get_model_device (model ).type == "cpu" ,
84
- msg = "Model should be offloaded to CPU" ,
85
- )
86
- model .cuda ()
82
+ if ir == "dynamo" :
83
+ assertions .assertTrue (
84
+ get_model_device (model ).type == "cpu" ,
85
+ msg = "Model should be offloaded to CPU" ,
86
+ )
87
+ model .cuda ()
87
88
cos_sim = cosine_similarity (model (input ), trt_mod (input ))
88
89
assertions .assertTrue (
89
90
cos_sim > COSINE_THRESHOLD ,
@@ -286,11 +287,12 @@ def test_bert_base_uncased_cpu_offload(ir):
286
287
"offload_module_to_cpu" : True ,
287
288
}
288
289
trt_mod = torchtrt .compile (model , ** compile_spec )
289
- assertions .assertTrue (
290
- get_model_device (model ).type == "cpu" ,
291
- msg = "Model should be offloaded to CPU" ,
292
- )
293
- model .cuda ()
290
+ if ir == "dynamo" :
291
+ assertions .assertTrue (
292
+ get_model_device (model ).type == "cpu" ,
293
+ msg = "Model should be offloaded to CPU" ,
294
+ )
295
+ model .cuda ()
294
296
295
297
model_outputs = model (input , input2 )
296
298
trt_model_outputs = trt_mod (input , input2 )
You can’t perform that action at this time.
0 commit comments