Fix DiLoCo with DTensor

H-Huang · H-Huang · commit 783777d6dcb9 · 2025-05-20T14:32:13.000-07:00
diff --git a/torchft/local_sgd.py b/torchft/local_sgd.py
@@ -249,7 +249,7 @@ def _restore_parameters(self) -> None:
                     # we averaged the local version of the tensor so need to copy it back as a DTensor
                     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()
                         ),
                         non_blocking=False,
                     )

Original file line number	Diff line number	Diff line change
`@@ -249,7 +249,7 @@ def _restore_parameters(self) -> None:`
`249`	`249`	`# we averaged the local version of the tensor so need to copy it back as a DTensor`
`250`	`250`	`p.data.copy_(`
`251`	`251`	`DTensor.from_local(`
`252`		`- self.original_parameters[name], p.device_mesh, p.placements`
	`252`	`+ self.original_parameters[name], p.device_mesh, p.placements, shape=p.shape, stride=p.stride()`
`253`	`253`	`),`
`254`	`254`	`non_blocking=False,`
`255`	`255`	`)`