[bugfix] make fp8 scale_inv contiguous before TP/PP collectives - #197
Merged
tastelikefeet merged 1 commit intoSep 15, 2026
Merged
Conversation
`_get_weight` trims the padding that TE adds to `_rowwise_scale_inv`, and slicing the last dim returns a non-contiguous view. Grouped MoE experts keep the weight as a single 3D tensor, so `mg_scale_inv[0]` preserves that view and hands it straight to `_all_gather_tp`, where ProcessGroupNCCL raises `ValueError: Tensors must be contiguous`. This breaks `megatron export --fp8_recipe blockwise --fp8_param_gather true` whenever TP > 1, and would break `dist.broadcast` in `_broadcast_ep_pp` for TP == 1 with PP > 1. Call `.contiguous()` at the slice site so both collectives are covered. It is a no-op when TE did not pad the last dim, so the common path keeps the same storage and adds no copy; when it does copy, the tensor is 1/128 of the fp8 data. Values are unchanged. Multi-shard weights were unaffected because `torch.concat` already returns a contiguous tensor, which is why only grouped-GEMM MoE hit this. Fixes modelscope/ms-swift#10136
hjh0119
approved these changes
Sep 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
_get_weighttrims the padding that TE adds to_rowwise_scale_inv, and slicing the last dim returns a non-contiguous view. Grouped MoE experts keep the weight as a single 3D tensor, somg_scale_inv[0]preserves that view and hands it straight to_all_gather_tp, where ProcessGroupNCCL raisesValueError: Tensors must be contiguous. This breaksmegatron export --fp8_recipe blockwise --fp8_param_gather truewhenever TP > 1, and would breakdist.broadcastin_broadcast_ep_ppfor TP == 1 with PP > 1.Call
.contiguous()at the slice site so both collectives are covered. It is a no-op when TE did not pad the last dim, so the common path keeps the same storage and adds no copy; when it does copy, the tensor is 1/128 of the fp8 data. Values are unchanged.Multi-shard weights were unaffected because
torch.concatalready returns a contiguous tensor, which is why only grouped-GEMM MoE hit this.Fixes modelscope/ms-swift#10136