`dist_nn.all_reduce` return value not captured in `Physics_Attention_1D_Eidetic`, causing incorrect eidetic state aggregation in multi-GPU setting

[anonymous-author-conference]

Description

In Physics_Attention_1D_Eidetic.forward(), torch.distributed.nn.all_reduce is called without capturing its return value. Unlike torch.distributed.all_reduce (which is in-place), torch.distributed.nn.all_reduce returns a new tensor and does not modify the input in-place. As a result, slice_norm and slice_token remain as per-GPU partial sums, and the global aggregation across GPUs never actually takes effect.

Bug Location

Physics_Attention_1D_Eidetic.forward() in Transolver_plus.py

Current (incorrect) code

dist_nn.all_reduce(slice_norm, op=dist_nn.ReduceOp.SUM)
dist_nn.all_reduce(slice_token, op=dist_nn.ReduceOp.SUM)

Fix

slice_norm = dist_nn.all_reduce(slice_norm, op=dist_nn.ReduceOp.SUM)
slice_token = dist_nn.all_reduce(slice_token, op=dist_nn.ReduceOp.SUM)

Impact

• In multi-GPU training, eidetic states are computed from local partial sums only, making the parallelism framework silently incorrect.
• Gradient flow through the AllReduce operation is also broken, since the returned differentiable tensor is discarded.
• In single-GPU setting, the bug has no effect (no communication needed), which may have masked this issue during testing.