How to let GAN do in fsdp

username: Wei-Hsin
cooked: <p><span class="hashtag">#i</span> want to build cyclegan do in fsdp but it got this error that i have no idea how to let fsdp backward loss in same device<br>
self.genX2Y.train()<br>
self.genY2X.train()<br>
self.discX.train()<br>
self.discY.train()<br>
fake_x = self.genY2X(y_img1)<br>
fake_y = self.genX2Y(x_img1)<br>
self.opt_disc.zero_grad()<br>
D_X_real = self.discX(x_img1)<br>
D_X_loss = -torch.mean( D_X_real)<br>
D_Y_real = self.discY(y_img1)<br>
D_Y_loss = -torch.mean( D_Y_real)<br>
D_loss = (D_X_loss + D_Y_loss) / 2</p>
<p>dist.all_reduce(D_loss, op=dist.ReduceOp.SUM)<br>
D_loss.backward(retain_graph=True)</p>
<p>self.opt_gen.zero_grad()</p>
<h1>
<a name="train-generators-1" class="anchor" href="#train-generators-1"></a>train generators</h1>
<p>D_X_fake1 = self.discX(fake_x)<br>
D_Y_fake1 = self.discY(fake_y)<br>
loss_G_X1 = -torch.mean(D_X_fake1)<br>
loss_G_Y1 = -torch.mean(D_Y_fake1)</p>
<h1>
<a name="cycle-loss-2" class="anchor" href="#cycle-loss-2"></a>cycle loss</h1>
<p>cycle_Y = self.genX2Y(fake_x)<br>
cycle_X = self.genY2X(fake_y)<br>
cycle_Y_loss1 = self.mse(cycle_Y, y_img1)<br>
cycle_X_loss1 = self.mse(cycle_X, x_img1)</p>
<h1>
<a name="identity-loss-3" class="anchor" href="#identity-loss-3"></a>identity loss</h1>
<p>id_Y = self.genX2Y(y_img1)<br>
id_X = self.genY2X(x_img1)<br>
id_Y_loss = self.mse(id_Y,x_img1)<br>
id_X_loss = self.mse(id_X,y_img1)</p>
<h1>
<a name="combined-loss-4" class="anchor" href="#combined-loss-4"></a>combined loss</h1>
<p>LAMBDA_CYCLE = 3<br>
LAMBDA_ID = 1<br>
cycle=(cycle_Y_loss1 + cycle_X_loss1)*LAMBDA_CYCLE<br>
loss_GAN=(loss_G_Y1 + loss_G_X1)/2<br>
loss_identity=(id_Y_loss + id_X_loss)*LAMBDA_ID/2<br>
G_loss = loss_GAN + cycle + loss_identity<br>
dist.all_reduce(G_loss, op=dist.ReduceOp.SUM)<br>
G_loss.backward(retain_graph=True)<br>
self.opt_gen.step()<br>
self.opt_disc.step()</p>
<p><span class="hashtag">#error</span> messages:<br>
File “/root/notebooks/groups/FSDP_torchrun_0420.py”, line 772, in _run_epoch<br>
fake_x = self.genY2X(y_img1)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/nn/modules/module.py”, line 1501, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py”, line 783, in forward<br>
args, kwargs = _pre_forward(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 414, in _pre_forward<br>
unshard_fn()<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 439, in _pre_forward_unshard<br>
_unshard(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 331, in _unshard<br>
handle.post_unshard()<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/flat_param.py”, line 1249, in post_unshard<br>
self._check_on_compute_device(self.flat_param)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/flat_param.py”, line 2259, in _check_on_compute_device<br>
_p_assert(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/utils.py”, line 104, in _p_assert<br>
traceback.print_stack()<br>
Expects tensor to be on the compute device cuda:1<br>
File “/root/notebooks/groups/FSDP_torchrun_0420.py”, line 969, in <br>
main(args.save_every, args.total_epochs, args.batch_size)<br>
File “/root/notebooks/groups/FSDP_torchrun_0420.py”, line 958, in main<br>
trainer.train(total_epochs)<br>
File “/root/notebooks/groups/FSDP_torchrun_0420.py”, line 895, in train<br>
self._run_epoch(epoch)<br>
File “/root/notebooks/groups/FSDP_torchrun_0420.py”, line 772, in _run_epoch<br>
fake_x = self.genY2X(y_img1)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/nn/modules/module.py”, line 1501, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py”, line 783, in forward<br>
args, kwargs = _pre_forward(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 414, in _pre_forward<br>
unshard_fn()<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 439, in _pre_forward_unshard<br>
_unshard(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/_runtime_utils.py”, line 331, in _unshard<br>
handle.post_unshard()<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/flat_param.py”, line 1249, in post_unshard<br>
self._check_on_compute_device(self.flat_param)<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/fsdp/flat_param.py”, line 2259, in _check_on_compute_device<br>
_p_assert(<br>
File “/root/notebooks/anaconda3/envs/aienv/lib/python3.11/site-packages/torch/distributed/utils.py”, line 104, in _p_assert<br>
traceback.print_stack()</p>

username: marksaroufim
cooked: <aside class="quote no-group" data-username="Wei-Hsin" data-post="1" data-topic="177992">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wei-hsin/48/59693_2.png" class="avatar"> Wei-Hsin:</div>
<blockquote>
<p>Expects tensor to be on the compute device cuda:1</p>
</blockquote>
</aside>
<p>Hard to answer for sure if you can’t share a fully functioning repo but I suspect this is the meaningful error <code>Expects tensor to be on the compute device cuda:1</code></p>
<p>Can you make sure all your tensors are on the correct device?</p>