[release/2.9] Revert tracking of Work status for FlightRecorder in ProcessGroupXCCL

[frost-intel]

See #2076 , this is a cherry-pick for 2.9 release

This is a high impact bug, in many distributed applications this is a large memory leak resulting in OoM error (see #2084)

[riverliuintel]

@CuiYifeng add this fix as torch-xpu-ops cherry-pick fix list for PT2.9

[CuiYifeng]

@CuiYifeng add this fix as torch-xpu-ops cherry-pick fix list for PT2.9

Added.

[zhangxiaoli73]

@frost-intel The issue happened as: 1) you use work in you callback. 2) you use value in Future argument in your callback, and then output tensor will be reference again.

For 1) I think you don't need whole work but some work status. As work status is already updated before return, so you don't need to put it to callback.

torch-xpu-ops/src/xccl/ProcessGroupXCCL.cpp

Line 784 in 77cc792

setEnqueuedPgStatus(work);

For 2) PyTorch provide option "use_future" to decide if you really need Future argument in callback.

So my suggested change is only:

  auto id = work->trace_id_;
  work->future_->addCallback(
      [id](at::ivalue::Future&) {
        FlightRecorderXCCL::get()->retire_id(id, /*compute_duration*/ false);
      },
       /*use_future*/ false);

[frost-intel]

@zhangxiaoli73 Used your suggested change.
@EikanWang This should be ready for review/merge.