The three training loops in distributed/tensor_parallelism call
backward() and optimizer.step() without ever calling zero_grad(), and
none of the files calls it anywhere. PyTorch accumulates gradients into
.grad by default, so each iteration steps on the running sum of every
gradient computed so far rather than on that iteration's own.
With num_iters = 10 the loops still run and still print, so nothing
looks wrong -- but the optimizer is not doing what the example appears
to demonstrate, and these files are a common starting point for real
training code.
Placed after optimizer.step() to match the sibling
distributed/FSDP2/example.py.
The three training loops in
distributed/tensor_parallelismcallbackward()andoptimizer.step()without ever callingzero_grad()— and none of the three files calls it anywhere:PyTorch accumulates into
.gradby default, so iteration n steps on the sum of all gradients from iterations 1..n rather than on its own.With
num_iters = 10the loops still run and still log, so nothing looks broken — but the optimizer isn't doing what the example appears to demonstrate, and these files are a natural starting point for real training code.Fixed in all three (
tensor_parallel_example.py,sequence_parallel_example.py,fsdp_tp_example.py), placed afteroptimizer.step()to match the siblingdistributed/FSDP2/example.py.Found with torch-preflight.