Fix DiLoCo with DTensor #197
Conversation
facebook-github-bot
added
the
CLA Signed
torchft/local_sgd.py
Outdated
| p.data.copy_( | ||
| DTensor.from_local( | ||
| self.original_parameters[name], p.device_mesh, p.placements | ||
| self.original_parameters[name], p.device_mesh, p.placements, shape=p.shape, stride=p.stride() |
There was a problem hiding this comment.
I'm curious why shape changes?
There was a problem hiding this comment.
For 10 ranks, the original parameter was shape [128256, 4096]
After converting to dtensor each rank had a local tensor of [12826, 4096], except for rank 9 which was [12822, 4096]
I saw this note in the docs about needing to add shape arg if tensors are not even across ranks. Though it is strange because I cannot repro this in a test yet. However this does fix the issue when running on the cluster
There was a problem hiding this comment.
Uneven sharding part makes sense. But I recalled the issue didn't happen at the first iteration? If it is an uneven sharding issue, shouldn't it happen at the first iteration?
There was a problem hiding this comment.
it would happen at the iteration which the semi-sync methods perform sync, so its dependent on the sync_every config
3 participants
Hitting an error with
RuntimeError: The size of tensor a (128256) must match the size of tensor b (128220) at non-singleton dimension 0when restoring the parameters in DiLoCo.Depends on this fix: pytorch/pytorch#154156, tests pass with pytorch nightly