Fix run_demo(demo_model_parallel, world_size) issue

[TheMemoryDealer]

Fixes #1750

Description

Fixes run_demo(demo_model_parallel, world_size) issue as described in #1750

Checklist

• The issue that is being fixed is referred in the description (see above "Fixes #ISSUE_NUMBER")
• Only one issue is addressed in this pull request
• Labels from the issue that this PR is fixing are added to this pull request
• No unnessessary issues are included into this pull request.

cc @mrshenli @osalpekar @H-Huang @kwen2501

[netlify]

✅ Deploy Preview for pytorch-tutorials-preview ready!

Name	Link
🔨 Latest commit	86acfe7
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-tutorials-preview/deploys/647a3bc622d6be00081ef72f
😎 Deploy Preview	https://deploy-preview-2367--pytorch-tutorials-preview.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify site settings.

[subramen]

Changing the value of world_size impacts the values of dev0 and dev1, you will need to update how dev is calculated to account for the new world_size value

[subramen]

Please update the dev calculation in the model class to reflect the new world_size value

[TheMemoryDealer]

@subramen I've updated the calculation in a simple way to take into account the prior division. So now dev0 and dev1 are calculated as

    dev0 = (rank * 2) % (world_size * 2)
    dev1 = (rank * 2 + 1) % (world_size * 2)

[TheMemoryDealer]

@subramen going back on it, does % (world_size * 2) not simply eliminate the world_size = n_gpus//2 ? Would it not make more sense to just ignore world_size = n_gpus//2 and keep

dev0 = rank * 2
dev1 = rank * 2 + 1

?

[subramen]

Yes, you don't actually need world_size anymore :) that is the update I was looking for.