Fix segment fault for alltoall

[pengwa]

Description: fix segment fault for alltoall.

In MOE model tests, there are random issues reporting segment fault.

The problem is AllToAll autograd function is consuming a process group object in its forward function. when the segmentfault happens, the process group param accepted in the forward function is a 'str' like "ONNX:Slice114" or some wired object pointer.

The root cause is: during PythonOp export, we save the process group (as a non tensor pointer type) object's address only, and save it in PythonOp's attribute. But python did not know there is a reference to that object. In the new env MOE test is trying to onboard, it looks there is more possibility process group is destroyed before the tests completed. If it is teared down before PythonOp/PythonOpGrad completed, an unexpected object (e.g. the str or werid object) is picked up and used as process group type, segement fault occurs.

The fix: the process group as inputs of all PythonOp, they should exist along the training lifetime. So we can track those non-tensor object reference once we are exporting PythonOp. Keeping those object in a global store will avoid that object be released. We will clean up the store before python program exits. For multiple UT cases, cleaning up is not triggered, but it should be fine.

Motivation and Context

• Why is this change required? What problem does it solve?
• If it fixes an open issue, please link to the issue here.

[pengwa]

Thanks @wschin @baijumeswani.