Skip to content

Commit dd98938

Browse files
whoisjwhoisj
committed
add torch.export.export generation support to libtorch
1 parent 16c6c0a commit dd98938

2 files changed

Lines changed: 214 additions & 9 deletions

File tree

‎qa/common/gen_qa_model_repository‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -241,9 +241,11 @@ set -e
241241
set -x
242242
PATH=$PATH:/usr/local/cuda-13.0/bin
243243
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --models_dir=$TRITON_MDLS_QA_MODEL
244+
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch-pt2 --models_dir=$TRITON_MDLS_QA_MODEL
244245
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --torch-aoti --models_dir=$TRITON_MDLS_QA_MODEL
245246
chmod -R 777 $TRITON_MDLS_QA_MODEL
246247
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
248+
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --libtorch-pt2 --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
247249
python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --torch-aoti --variable --models_dir=$TRITON_MDLS_QA_VARIABLE_MODEL
248250
chmod -R 777 $TRITON_MDLS_QA_VARIABLE_MODEL
249251
python3 $TRITON_MDLS_SRC_DIR/gen_qa_identity_models.py --libtorch --models_dir=$TRITON_MDLS_QA_IDENTITY_MODEL
@@ -437,7 +439,11 @@ if which docker ; then
437439
$UBUNTU_IMAGE \
438440
bash -xe $TRITON_MDLS_SRC_DIR/$OPENVINOSCRIPT
439441

440-
if [ $? -ne 0 ]; then
442+
exit_code=$?
443+
444+
rm $OPENVINOSCRIPT
445+
446+
if [ $exit_code -ne 0 ]; then
441447
log_message.error "docker run: ${OPENVINOSCRIPT} failed"
442448
exit 1
443449
fi
@@ -454,7 +460,11 @@ if which docker ; then
454460
$UBUNTU_IMAGE \
455461
bash -xe $TRITON_MDLS_SRC_DIR/$ONNXSCRIPT
456462

457-
if [ $? -ne 0 ]; then
463+
exit_code=$?
464+
465+
rm $ONNXSCRIPT
466+
467+
if [ $exit_code -ne 0 ]; then
458468
log_message.error "docker run: ${ONNXSCRIPT} failed"
459469
exit 1
460470
fi
@@ -473,7 +483,11 @@ if which docker ; then
473483
$PYTORCH_IMAGE \
474484
bash -xe $TRITON_MDLS_SRC_DIR/$TORCHSCRIPT
475485

476-
if [ $? -ne 0 ]; then
486+
exit_code=$?
487+
488+
rm $TORCHSCRIPT
489+
490+
if [ $exit_code -ne 0 ]; then
477491
log_message.error "docker run: ${TORCHSCRIPT} failed"
478492
exit 1
479493
fi
@@ -494,7 +508,11 @@ if which docker ; then
494508
$TENSORRT_IMAGE \
495509
bash -xe $TRITON_MDLS_SRC_DIR/$TRTSCRIPT
496510

497-
if [ $? -ne 0 ]; then
511+
exit_code=$?
512+
513+
rm $TRTSCRIPT
514+
515+
if [ $exit_code -ne 0 ]; then
498516
log_message.error "docker run: ${TRTSCRIPT} failed"
499517
exit 1
500518
fi

‎qa/common/gen_qa_models.py‎

Lines changed: 192 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1270,7 +1270,7 @@ def forward(self, INPUT0, INPUT1):
12701270
)
12711271
return op0, op1
12721272

1273-
addSubModel = AddSubNet((torch_output0_dtype, torch_output1_dtype, swap))
1273+
addSubModel = AddSubNet((torch_output0_dtype, torch_output1_dtype, swap,))
12741274
traced = torch.jit.script(addSubModel)
12751275

12761276
model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)
@@ -1283,6 +1283,61 @@ def forward(self, INPUT0, INPUT1):
12831283
traced.save(model_version_dir + "/model.pt")
12841284

12851285

1286+
def create_libtorch_pt2_modelfile(
1287+
models_dir,
1288+
max_batch,
1289+
model_version,
1290+
input_shape,
1291+
output0_shape,
1292+
output1_shape,
1293+
input_dtype,
1294+
output0_dtype,
1295+
output1_dtype,
1296+
swap=False,
1297+
):
1298+
if not tu.validate_for_libtorch_model(
1299+
input_dtype,
1300+
output0_dtype,
1301+
output1_dtype,
1302+
input_shape,
1303+
output0_shape,
1304+
output1_shape,
1305+
max_batch,
1306+
):
1307+
return
1308+
1309+
model_name = tu.get_model_name(
1310+
"libtorch_pt2",
1311+
input_dtype,
1312+
output0_dtype,
1313+
output1_dtype,
1314+
)
1315+
# handle for -1 (when variable) since can't create tensor with shape of [-1]
1316+
input_shape = [abs(ips) for ips in input_shape]
1317+
1318+
model_version_dir = models_dir + "/" + model_name + "/" + str(model_version)
1319+
1320+
try:
1321+
os.makedirs(model_version_dir)
1322+
except OSError:
1323+
pass # ignore existing dir
1324+
1325+
class AddSubNet(nn.Module):
1326+
def __init__(self, swap):
1327+
self.swap = swap
1328+
super(AddSubNet, self).__init__()
1329+
1330+
def forward(self, INPUT0, INPUT1):
1331+
op0 = (INPUT0 - INPUT1) if self.swap else (INPUT0 + INPUT1)
1332+
op1 = (INPUT0 + INPUT1) if self.swap else (INPUT0 - INPUT1)
1333+
return op0, op1
1334+
1335+
ep = torch.export.export(
1336+
AddSubNet(swap), (torch.randn(*input_shape), torch.randn(*input_shape),),
1337+
)
1338+
torch.export.save(ep, model_version_dir + "/model.pt2")
1339+
1340+
12861341
def create_torch_aoti_modelfile(
12871342
models_dir,
12881343
max_batch,
@@ -1333,12 +1388,11 @@ def forward(self, INPUT0, INPUT1):
13331388
return op0, op1
13341389

13351390
ep = torch.export.export(
1336-
AddSubNet(swap), (torch.randn(*input_shape), torch.randn(*input_shape))
1391+
AddSubNet(swap), (torch.randn(*input_shape), torch.randn(*input_shape),),
13371392
)
13381393
torch._inductor.aoti_compile_and_package(
1339-
ep, package_path=model_version_dir + "/model.pt2"
1394+
ep, package_path=model_version_dir + "/model.pt2",
13401395
)
1341-
torch._inductor.aoti_load_package
13421396

13431397

13441398
def create_libtorch_modelconfig(
@@ -1442,6 +1496,107 @@ def create_libtorch_modelconfig(
14421496
lfile.write("label" + str(l) + "\n")
14431497

14441498

1499+
def create_libtorch_pt2_modelconfig(
1500+
models_dir,
1501+
max_batch,
1502+
model_version,
1503+
input_shape,
1504+
output0_shape,
1505+
output1_shape,
1506+
input_dtype,
1507+
output0_dtype,
1508+
output1_dtype,
1509+
output0_label_cnt,
1510+
version_policy,
1511+
):
1512+
if not tu.validate_for_libtorch_model(
1513+
input_dtype,
1514+
output0_dtype,
1515+
output1_dtype,
1516+
input_shape,
1517+
output0_shape,
1518+
output1_shape,
1519+
max_batch,
1520+
):
1521+
return
1522+
1523+
# Unpack version policy
1524+
version_policy_str = "{ latest { num_versions: 1 }}"
1525+
if version_policy is not None:
1526+
type, val = version_policy
1527+
if type == "latest":
1528+
version_policy_str = "{{ latest {{ num_versions: {} }}}}".format(val)
1529+
elif type == "specific":
1530+
version_policy_str = "{{ specific {{ versions: {} }}}}".format(val)
1531+
else:
1532+
version_policy_str = "{ all { }}"
1533+
1534+
# Use a different model name for the non-batching variant
1535+
model_name = tu.get_model_name(
1536+
"libtorch_pt2_nobatch" if max_batch == 0 else "libtorch_pt2",
1537+
input_dtype,
1538+
output0_dtype,
1539+
output1_dtype,
1540+
)
1541+
config_dir = models_dir + "/" + model_name
1542+
config = """
1543+
backend: "pytorch"
1544+
name: "{}"
1545+
platform: "pytorch_libtorch"
1546+
max_batch_size: {}
1547+
version_policy: {}
1548+
input [
1549+
{{
1550+
name: "INPUT0"
1551+
data_type: {}
1552+
dims: [ {} ]
1553+
}},
1554+
{{
1555+
name: "INPUT1"
1556+
data_type: {}
1557+
dims: [ {} ]
1558+
}}
1559+
]
1560+
output [
1561+
{{
1562+
name: "OUTPUT__0"
1563+
data_type: {}
1564+
dims: [ {} ]
1565+
label_filename: "output0_labels.txt"
1566+
}},
1567+
{{
1568+
name: "OUTPUT__1"
1569+
data_type: {}
1570+
dims: [ {} ]
1571+
}}
1572+
]
1573+
""".format(
1574+
model_name,
1575+
max_batch,
1576+
version_policy_str,
1577+
np_to_model_dtype(input_dtype),
1578+
tu.shape_to_dims_str(input_shape),
1579+
np_to_model_dtype(input_dtype),
1580+
tu.shape_to_dims_str(input_shape),
1581+
np_to_model_dtype(output0_dtype),
1582+
tu.shape_to_dims_str(output0_shape),
1583+
np_to_model_dtype(output1_dtype),
1584+
tu.shape_to_dims_str(output1_shape),
1585+
)
1586+
1587+
try:
1588+
os.makedirs(config_dir)
1589+
except OSError as ex:
1590+
pass # ignore existing dir
1591+
1592+
with open(config_dir + "/config.pbtxt", "w") as cfile:
1593+
cfile.write(config)
1594+
1595+
with open(config_dir + "/output0_labels.txt", "w") as lfile:
1596+
for l in range(output0_label_cnt):
1597+
lfile.write("label" + str(l) + "\n")
1598+
1599+
14451600
def create_torch_aoti_modelconfig(
14461601
models_dir,
14471602
max_batch,
@@ -1905,6 +2060,32 @@ def create_models(
19052060
output1_dtype,
19062061
)
19072062

2063+
if FLAGS.libtorch_pt2:
2064+
create_libtorch_pt2_modelconfig(
2065+
models_dir,
2066+
0,
2067+
model_version,
2068+
input_shape,
2069+
output0_shape,
2070+
output1_shape,
2071+
input_dtype,
2072+
output0_dtype,
2073+
output1_dtype,
2074+
output0_label_cnt,
2075+
version_policy,
2076+
)
2077+
create_libtorch_pt2_modelfile(
2078+
models_dir,
2079+
8,
2080+
model_version,
2081+
input_shape,
2082+
output0_shape,
2083+
output1_shape,
2084+
input_dtype,
2085+
output0_dtype,
2086+
output1_dtype,
2087+
)
2088+
19082089
if FLAGS.torch_aoti:
19092090
# max-batch 8
19102091
create_torch_aoti_modelconfig(
@@ -2122,6 +2303,12 @@ def create_fixed_models(
21222303
action="store_true",
21232304
help="Generate Pytorch LibTorch models",
21242305
)
2306+
parser.add_argument(
2307+
"--libtorch-pt2",
2308+
required=False,
2309+
action="store_true",
2310+
help="Generate Pytorch LibTorch PT2 models",
2311+
)
21252312
parser.add_argument(
21262313
"--torch-aoti",
21272314
required=False,
@@ -2154,7 +2341,7 @@ def create_fixed_models(
21542341
import tensorrt as trt
21552342
if FLAGS.onnx:
21562343
import onnx
2157-
if FLAGS.libtorch or FLAGS.torch_aoti:
2344+
if FLAGS.libtorch or FLAGS.libtorch_pt2 or FLAGS.torch_aoti:
21582345
import torch
21592346
from torch import nn
21602347
if FLAGS.openvino:

0 commit comments

Comments
 (0)