Commit 103343ca authored by Rick Ho's avatar Rick Ho
Browse files

strong compatibility to different python and pytorch versions

parent d6169eb1
......@@ -116,7 +116,11 @@ class HackNCCLGroup: public c10d::ProcessGroupNCCL {
public:
ncclComm_t getcomm(at::Device dev) {
auto key = std::to_string(dev.index());
#ifdef ENABLE_NCCL_P2P_SUPPORT
auto v = getNCCLComm(key, {dev}, c10d::OpType::ALLTOALL);
#else
auto v = getNCCLComm(key, {dev});
#endif
if (v.size() == 0) {
std::cerr << "PyTorch has nothing\n";
return 0;
......
......@@ -20,8 +20,12 @@ fi
mkdir -p logs
PYTHON_EXEC=python3
PYTHON_VERSION=$($PYTHON_EXEC --version)
PYTHON_REVISION=${PYTHON_VERSION:7:3}
SCRIPT_PATH=$(dirname $(dirname $(realpath $0)))
export PYTHONPATH=$SCRIPT_PATH:$SCRIPT_PATH/build/lib.linux-x86_64-3.7:$PYTHONPATH
export LD_LIBRARY_PATH=/home/laekov/.local/lib/python3.7/site-packages/torch/lib:$LD_LIBRARY_PATH
export PYTHONPATH=$SCRIPT_PATH:$SCRIPT_PATH/build/lib.linux-x86_64-$PYTHON_REVISION:$PYTHONPATH
export LD_LIBRARY_PATH=/home/laekov/.local/lib/python$PYTHON_REVISION/site-packages/torch/lib:$LD_LIBRARY_PATH
exec python3 $@ 2>logs/$RANK.log
exec $PYTHON_EXEC $@ 2>logs/$RANK.log
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment