Skip to content

Commit 2b16a71

Browse files
authored
[FSDP] Don't use detach in _rebuild_full_params (#4813)
Summary: Somehow the detach_copy implementation has some issues with FSDP that it messed up with FSDP's memory management mechanism. Let's not use it. Test Plan: CI
1 parent 46e6238 commit 2b16a71

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1398,7 +1398,7 @@ def _rebuild_full_params(self,
13981398

13991399
for p, p_shard in zip(self.full_params, self.sharded_params):
14001400
if not p._has_full_param:
1401-
p_shard_data = p_shard.detach()
1401+
p_shard_data = p_shard
14021402
if apply_opt_barrier:
14031403
self.optimization_barrier_op([p_shard_data])
14041404
if p_shard_data.dtype != self.compute_dtype:

0 commit comments

Comments
 (0)