fix(moe): auto-degrade FP8 block_n/k from 128 to 64 on alignment mismatch (#442)#1488
Open
haowu1234 wants to merge 1 commit into
Open
fix(moe): auto-degrade FP8 block_n/k from 128 to 64 on alignment mismatch (#442)#1488haowu1234 wants to merge 1 commit into
haowu1234 wants to merge 1 commit into
Conversation
…nment mismatch Models with intermediate_size=1536 at TP=8 produce intermediate_size_per_partition=192, which is not divisible by block_n=128 or block_k=128 (the defaults for per_1x128 block quantisation). Both CompressedTensorsFp8MoEMethod and Fp8MoEMethod raise ValueError, forcing users to fall back to TP=4 and waste 4 GPUs (MiniMax M2.5, Qwen3-235B, etc.). Fix (Option B from ROCm#442): When intermediate_size_per_partition is not divisible by 128 but is divisible by 64, auto-degrade block_n and/or block_k to 64 and emit a logger warning. The ValueErrors are preserved when alignment is still impossible. Changes: - CompressedTensorsFp8MoEMethod.create_weights: auto-degrade self.block_n and self.block_k before alignment checks. - Fp8MoEMethod: initialise self.block_n/self.block_k in __init__ (was local variables in create_weights only); auto-degrade in create_weights; use them in get_fused_moe_quant_config instead of the hardcoded [128,128] that would mismatch the degraded scale shapes. - tests/test_fp8_moe_block_align.py: 13 unit tests covering degrade path, aligned path, impossible alignment, per_1x32 no-op, scale shape correctness, block_k degradation at TP>1, and full TP=8/1536 end-to-end scenarios. Validated: 13/13 passed on MI300X (gfx942, ROCm 7.1, Python 3.12). Refs: ROCm#442
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.
Problem
Models with
intermediate_size=1536(MiniMax M2.5, Qwen3-235B) at TP=8 produceintermediate_size_per_partition=192, which is not divisible byblock_n=128(the default forper_1x128block quantisation). BothCompressedTensorsFp8MoEMethodandFp8MoEMethodraiseValueError, forcing users to fall back to TP=4 and waste 4 GPUs.The existing error:
Additionally, after fixing
block_n,block_k=128also fails at TP=8 (192 % 128 != 0), which was masked by the earlierblock_nerror.Fix (Option B from #442)
When
intermediate_size_per_partitionis not divisible by 128 but is divisible by 64, auto-degradeblock_nand/orblock_kto 64 and emit alogger.warning. TheValueErroris preserved when alignment is truly impossible (e.g.190 % 64 != 0).Changes
CompressedTensorsFp8MoEMethod.create_weights: auto-degradeself.block_nandself.block_kbefore alignment checks.Fp8MoEMethod.__init__: initialiseself.block_n/self.block_kinstance attributes (were local variables increate_weightsonly).Fp8MoEMethod.create_weights: auto-degrade both block sizes using instance attributes.Fp8MoEMethod.get_fused_moe_quant_config: use[self.block_n, self.block_k]instead of hardcoded[128, 128]that would mismatch the degraded scale shapes.tests/test_fp8_moe_block_align.py: 13 unit tests covering degrade path, aligned path, impossible alignment,per_1x32no-op, scale shape correctness,block_kdegradation at TP>1, and full TP=8/1536 end-to-end scenarios.Validation
block_n=block_k=64→w13_scale=(E, 6, 112),w2_scale=(E, 112, 3)for 192 intermediate at TP=8Refs: #442