How to train Mixture-of-Experts (MoE) model with Fully Sharded Data Parallel (FSDP)

liuslnlp: <p>I attempted to replace the FFN in Transformer with MoE (implemented by fairscale). I am curious about how to integrate MoE and FSDP together. The data parallel groups for different parameters in the model are not the same, and FSDP does not provide an interface to assign different dp groups to different parameters.</p>
<p>For example, consider the diagram below: the model has four experts, with two placed on each rank, and the parallelism of the experts is 2. In this case, the data parallel group for non-expert parameters (such as attention and word embedding) is [0, 1, 2, 3], while the dp groups for experts are [0, 2] and [1, 3].<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699.png" data-download-href="https://discuss.pytorch.org/uploads/default/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699" title="fsdp_moe"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699_2_690x334.png" alt="fsdp_moe" data-base62-sha1="wDzsp0T6iUyOKEqsIiEGXcHrcKJ" width="690" height="334" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699_2_690x334.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699_2_1035x501.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e4beb5d3e9c66f6cee8af1dd0beffa3110d0c699_2_1380x668.png 2x" data-dominant-color="E5E5E5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">fsdp_moe</span><span class="informations">1500×728 42.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Reference:</p><aside class="onebox githubissue" data-onebox-src="https://github.com/facebookresearch/fairscale/issues/871">
  <header class="source">

      <a href="https://github.com/facebookresearch/fairscale/issues/871" target="_blank" rel="noopener nofollow ugc">github.com/facebookresearch/fairscale</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/facebookresearch/fairscale/issues/871" target="_blank" rel="noopener nofollow ugc">[MoE] How to do expert sharding?</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-11-25" data-time="09:12:17" data-timezone="UTC">09:12AM - 25 Nov 21 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2022-01-07" data-time="14:23:14" data-timezone="UTC">02:23PM - 07 Jan 22 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/gongbudaizhe" target="_blank" rel="noopener nofollow ugc">
          <img alt="gongbudaizhe" src="https://avatars.githubusercontent.com/u/5178646?v=4" class="onebox-avatar-inline" width="20" height="20">
          gongbudaizhe
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## ❓ Questions and Help

MoE models have a different expert on each gpu device<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">. Therefore, these expert parameters should *not* be allreduced during backward pass. 

To achieve this, p.expert is set to True as shown in:
https://github.com/facebookresearch/fairscale/blob/562542478a145fae410cb23ffd46a8d7c2497abb/fairscale/nn/moe/moe_layer.py#L62-L64

My question is, how is p.expert used in DDP? Is this attribute automatically handled in ShardedDataParallel and FullyShardedDataParallel to disable allreduce?</span></p>
  </div>

  </article>

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

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