PyTorch Proposes DTensor Numerical Fix to Prevent NaN Cleanup from Changing Distributed Reduction Results

On November 7, the PyTorch community submitted PR #203933 to correct operation-placement rules for the distributed tensor DTensor: nan_to_num was listed as an operation that could preserve a pending max or min reduction, but that optimization can change the answer when NaNs or custom infinity replacement values are involved. At the time of review, the proposal remained open and had not been merged.

DTensor uses placements to describe how data is distributed. Partial means each process holds a local value that still needs to be reduced, while Replicate means each process holds a full copy. The official documentation explains that converting from the former to the latter runs all_reduce; placement-preservation rules therefore also determine the order of communication and computation. DTensor documentation

The issue can be illustrated with two processes holding NaN and 5, respectively. Under the ordinary tensor semantics reported in the issue, taking the maximum first yields NaN, and then applying nan_to_num gives 0. If NaN is replaced with 0 on each process before taking the maximum, the result is instead 5. Custom posinf and neginf replacement values can also break equivalence across operation orderings. This is a difference in operation semantics, not an acceptable floating-point tolerance. Issue and reproduction code

The fix removes the relevant placement-preservation rules for three nan_to_num overloads, so the system first converts to Replicate(), completes the reduction, and then replaces special values. The author said that four cases were reproduced on CPU with LocalTensorMode in a torch 2.14.1 environment and all matched expectations after the change; adjacent elementwise tests also passed. This is the author’s report, however, and the page showed that the relevant CI checks were still awaiting approval to start. Validation record

Engineering teams should check whether sharded tensors have NaNs or Infs cleaned up immediately after max or min reductions. Inputs containing special values can be used to compare distributed and ordinary tensor results. Based on the documentation and the proposed fix, completing the reduction earlier may also change communication timing, so performance should be measured separately. Follow-up should track maintainer review, real multi-process testing, and inclusion in a release. Current evidence is insufficient to conclude that ordinary training workflows are broadly affected.