InfiniTensor/examples/distributed
panzezhong 6c4dd7b28b fix(front): 将stub改为可以接收GraphProto作为输入,消除分布式脚本保存额外的onnx文件, 采用int64作为index输入类型 2024-04-07 17:15:40 +08:00
..
README.md 针对bert和gpt2模型分布式推理的优化 (#221) 2024-04-01 14:04:28 +08:00
bang_launch.py Bang cncl (#163) 2024-01-03 13:28:03 +08:00
cuda_launch.py fix(front): 将stub改为可以接收GraphProto作为输入,消除分布式脚本保存额外的onnx文件, 采用int64作为index输入类型 2024-04-07 17:15:40 +08:00
launch_kunlun.py XCCL support (#171) 2024-02-29 11:48:35 +08:00
launch_kvcache.py Support kvcache (#134) 2023-09-18 14:17:02 +08:00
parallel.py impl distributed launch with NCCL (#106) 2023-09-05 09:47:35 +08:00
parallel_opt.py fix(front): 将stub改为可以接收GraphProto作为输入,消除分布式脚本保存额外的onnx文件, 采用int64作为index输入类型 2024-04-07 17:15:40 +08:00
placement.py tensor parallel for transformer (#125) 2023-09-14 14:19:45 +08:00
run_pytorch.py fix(front): 将stub改为可以接收GraphProto作为输入,消除分布式脚本保存额外的onnx文件, 采用int64作为index输入类型 2024-04-07 17:15:40 +08:00

README.md

分布式脚本

1. 运行pytorch模型并生成输入和标准输出可选择导出onnx

使用 --export_onnx 设置导出onnx的目录默认为当前路径 ./不使用这个flag则只进行计算和生成输入输出。

python run_pytorch.py --model gpt2  --batch_size 1  --length 1 --export_onnx ./

会在当前目录下生成输入输出文件test_inputs.npytest_results.npy,目前只支持单一输入输出。

2. 运行InfiniTensor分布式脚本

python cuda_launch.py --model "/XXX/XXX.onnx" --nproc_per_node 4