[WIP] Qwen3 MoE support #2820
[WIP] Qwen3 MoE support #2820
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2820
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
| if "experts.0" in key: | ||
| new_key = get_mapped_key_moe(key, _FROM_HF) | ||
| converted_state_dict[new_key] = torch.stack( | ||
| [ | ||
| state_dict[str(i).join(key.rsplit("0", 1))].T | ||
| for i in range(num_experts) | ||
| ] | ||
| ) | ||
| elif "experts" in key: | ||
| continue |
There was a problem hiding this comment.
This is kinda hacky and slow, but stacking experts into a single tensor lets us use the existing MoE implementation and _grouped_mm for a very significant speed boost.
There was a problem hiding this comment.
I found this to be prohibitively slow for larger models when testing with DeepSeek. I ended up following the torchtitan approach which uses a nn.ModuleDict for storing expert weights and grouping activated experts on-the-fly for grouped gemm.
There was a problem hiding this comment.
As far as I can see, torchtitan also needs to stack expert weights to use grouped gemm, they just do it at a slightly later point. They also don't support grouped gemm for training for now.
There was a problem hiding this comment.
Oh you're totally right - I was thinking of EP.
If this doesn't take ages to convert the state dict it may be fine, otherwise only stacking the activated experts at runtime does make more sense to me.
|
|
||
| # top scores shape (bs*slen, top_k) | ||
| top_scores, selected_experts_indices = torch.topk( | ||
| scores, k=self.experts_per_token, dim=1 | ||
| ) | ||
| self.selected_experts_indices = selected_experts_indices | ||
| # top_scores /= top_scores.sum(dim=-1, keep_dim=True).to(x.dtype) | ||
| if self.norm_topk_prob: | ||
| top_scores /= top_scores.sum(dim=-1, keepdim=True).to(x.dtype) |
There was a problem hiding this comment.
YMMV but I ended up having to do something like
denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
top_scores /= denominator
This PR adds support for Qwen3 MoE (30B-A3B and 235B-A22B) models. Loss looked reasonable from a simple test with 30B-A3B on the Alpaca dataset.
TODO: