Can't get simple FSDP file to work

Nova-Rift: <p>Trying to use FSDP.</p>
<p>Can’t get this file to run with torchrun:</p>
<p>torchrun --nproc_per_node 2 Aug_27_v2.py</p>
<p>Keeps saying:</p>
<p>warnings.warn(<br>
Traceback (most recent call last):<br>
File “Aug_27_v2.py”, line 57, in <br>
train(rank, world_size)<br>
File “Aug_27_v2.py”, line 43, in train<br>
outputs = fsdp_model(inputs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1190, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py”, line 1440, in forward<br>
outputs = self.module(*args, **kwargs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1190, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/fairscale/nn/misc/flatten_params_wrapper.py”, line 487, in forward<br>
return self.module(*inputs, **kwinputs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1190, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “Aug_27_v2.py”, line 16, in forward<br>
return self.fc(x)<br>
File “/home/clark/.local/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1190, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/clark/.local/lib/python3.8/site-packages/torch/nn/modules/linear.py”, line 114, in forward<br>
return F.linear(input, self.weight, self.bias)<br>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_addmm)</p>
<p>File:</p>
<pre><code class="lang-auto"># Aug_27_v2.py

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

def train(rank, world_size):
    # Set up the device and allow the environment to decide the specific GPU
    device = torch.device(f"cuda:{rank}")

    # Initialize the distributed environment
    dist.init_process_group(
        'nccl',
        init_method='env://',
        rank=rank,
        world_size=world_size
    )

    # Create model and move to appropriate device
    model = SimpleModel().to(device)
    fsdp_model = FSDP(model)

    # Ensure criterion is also on the correct device
    criterion = nn.MSELoss().to(device)
    optimizer = optim.SGD(fsdp_model.parameters(), lr=0.01)

    inputs = torch.randn(20, 10).to(device)
    targets = torch.randn(20, 10).to(device)

    for epoch in range(1000):  # loop over the dataset multiple times
        optimizer.zero_grad()
        outputs = fsdp_model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        if rank == 0 and epoch % 100 == 0:  # Only print from one process to avoid clutter
            print('Epoch {}: Loss: {:.4f}'.format(epoch, loss.item()))

    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = int(os.environ.get('WORLD_SIZE', 1))
    rank = int(os.environ.get('RANK', 0))

    train(rank, world_size)

</code></pre>

Nova-Rift: <p>Couldn’t find a version that worked so I created one here if anyone is looking:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/Nova-Rift/very_simple_fsdp_example">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/Nova-Rift/very_simple_fsdp_example" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/344;"><img src="https://opengraph.githubassets.com/3a99f76a2b9a78c35b258329a7392c8697d1d7c5819b06a2c4f9305e38ebc191/Nova-Rift/very_simple_fsdp_example" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/Nova-Rift/very_simple_fsdp_example" target="_blank" rel="noopener nofollow ugc">GitHub - Nova-Rift/very_simple_fsdp_example: A very simple example of using...</a></h3>

  <p>A very simple example of using fairscales FSDP ( fully sharded data parallel ) - GitHub - Nova-Rift/very_simple_fsdp_example: A very simple example of using fairscales FSDP ( fully sharded data par...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>