Add Fused MoE W8A8 (Int8) Support #6978
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge). To run full CI, you can do one of these:
🚀 |
|
cc @robertgshaw2-neuralmagic feel free to provide any comments or link to any POCs. Thank you! |
| from vllm.model_executor.sampling_metadata import SamplingMetadata | ||
| from vllm.sequence import IntermediateTensors, SamplerOutput | ||
|
|
||
| from vllm.model_executor.layers.fused_moe import FusedMoE |
There was a problem hiding this comment.
Could you try using mixtral.py? We already have FusedMoE used there. This mixtral_quant is kind of a deprecated initial experiment before we decided on FusedMoE for quantization management
There was a problem hiding this comment.
Sure. That makes a lot of sense. I'll try to allocate some time in recent 1-2 weeks to rebase and update.
There was a problem hiding this comment.
Hey @mgoin since I'm also using the compressed tensor format from neural magic to implement this and I saw you already have one reverted pr that contains the compressed tensor format checking and loading. Do you suggest me to wait for the one to be checkin in again with marlin kernels or you have some better suggestions to help checkin this pr? I think one way for me is to check in the kernel part only in this pr. Or I can merge it together into the ExpertInt8 class to have another flag for w8a8 (this will be ). But this needs an extra config passing from user to know whether they wanna use w8a8 or w8a16, which is not good. Besides as it does quantize inplace to get scales quantize_in_place_and_get_scales rather than using GPTQ/AWQ as in llm-compressor, it seems not be not a high quality way of doing weight int8 quantization especially when we have w8a8.
I can also wait for your pr to be merged again with marlin kernels because otherwise I'll need to put some extra checking here to only keep the case for w8a8 + compressed tensor format (and avoid int4 format). Thank you!
There was a problem hiding this comment.
@mgoin I made a modification and if you think this solution is good as a temporary solution (before your more general version of code + marlin checked in), please feel free to provide comments and I can add some small testing + a bit config tuning as well before checking in. Thank you!
There was a problem hiding this comment.
Sorry for missing these comments @qingquansong ! The Marlin MoE PR should be in now so I think it should be good for you to build on. For now I am willing to live with separate w8a8 and w8a16 configs in the triton kernel to get this feature landed, but we should replace this with at least an enum so we can reduce it to a single argument soon.
There was a problem hiding this comment.
Thanks for the response! Do you suggest to remove this class https://github.com/vllm-project/vllm/pull/6978/files#diff-eddafffeb6f159f8c75f635d18a502fcfbf662a562b1ae7a8683a9790161a10bR135 if the llm compressor layer loading/setup schema has been fully checked in? 🤔
qingquansong
changed the title
There was a problem hiding this comment.
I think we don't have many INT8 models with per-tensor scales because the accuracy drop is too severe. We generally target static per-channel scales for the weights and dynamic per-token scales for the activations for INT8 https://github.com/neuralmagic/compressed-tensors/blob/86211a6b6f30dae9cfa7ee48b0994fed261902a2/src/compressed_tensors/quantization/quant_scheme.py#L115-L130 so I would like to see that implemented. FP8 is able to get away with static per-tensor scales for weights/activations much more easily.
cc @dsikka since the weight loading could also be simplified using the recent refactor
| b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk + | ||
| offs_bn[None, :] * stride_bn) | ||
|
|
||
| if use_int8_w8a8: # only support static per expert input activation whole tensor quantization now |
There was a problem hiding this comment.
Unfortunately I don't think per-tensor scales for activations will be accurate enough for int8. Do you think dynamic per-token quantization could be implemented?
There was a problem hiding this comment.
yeah, that could be done but I'm a bit confused on the llm-compressor recipe setup, do you have an example recipe yml file to correctly set it (without using the class to config it in code)?
| device_name = current_platform.get_device_name().replace(" ", "_") | ||
| device_name = torch.cuda.get_device_name().replace(" ", "_") |
There was a problem hiding this comment.
What was the issue with using current_platform? We should really try to keep using that interface as we support other accelerators
There was a problem hiding this comment.
oh, sure, will revert it, it's just for some local testing purpose
| if self.is_static_input_scheme: | ||
| scale = torch.nn.Parameter(torch.ones(num_experts, | ||
| dtype=torch.float32), | ||
| requires_grad=False) | ||
| set_weight_attrs(scale, { | ||
| "needs_scalar_to_array": True, | ||
| **layer_kwargs | ||
| }) | ||
| set_weight_attrs(scale, {"is_int8_input_scale": True}) | ||
| layer.register_parameter("w13_input_scale", scale) |
There was a problem hiding this comment.
This seems to assume per-tensor scale? We need to check the strategy to see if we can support the scheme in the checkpoint
There was a problem hiding this comment.
yes, input is per tensor (each expert has it's own input scale)
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has merge conflicts that must be resolved before it can be |
github-actions
bot
added
the
unstale
2 participants
Add Fused MoE W8A8 (Int8) Support since without having it int8 moe is much slower on A100 compared to bf16 with fused moe.
[Need more tuning to make it faster]
With some minor tuning, current kernel is for 4k context length with TP=8 on 8 A100 with max_num_batched_tokens= 64 * 4096 and 64 batch size is ~13400 toks/s for prefill stage (only slightly faster than bf16 with fused moe 12900 toks/s) and output ~23 toks/s which is much slower than bf16 output: 48.53 toks/s. Need further tuning to make it work better. A probably non-fair but still worthy comparison is that since 8*22b cannot (or at the limit) to serve with 4GPUs but the quantized one can be easily served with 4GPUs. If we serve 2 with 4GPUs on 1 node, each will have ~8500 toks/s for prefill and 24 toks/s for output, so can match the output while have 17000 toks/s per node throughput. Probably TP/EP implementation is the bottleneck that blocking the MOE to work well and Llama 70 without MoE can have 20%+ improvement even using TP8.
The quantization strategy is posted here: vllm-project/llm-compressor#35 (comment) with channel wise 8bit quantization for weights and per tensor static input activation quantization. Expert routing gate layer is not quantized (as cutlass kernel by default only support tensor_row % 16 = 0 input and 8 experts without doing padding does not support this.
One more thing is the chunk strategy here for solving #5938 seems not quite useful as removing it seems also works as expected and can give slightly faster speed for both normal bf16 and int8. Maybe can remove it?
BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE
PR Checklist (Click to Expand)
Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Model]for adding a new model or improving an existing model. Model name should appear in the title.[Frontend]For changes on the vLLM frontend (e.g., OpenAI API server,LLMclass, etc.)[Kernel]for changes affecting CUDA kernels or other compute kernels.[Core]for changes in the core vLLM logic (e.g.,LLMEngine,AsyncLLMEngine,Scheduler, etc.)[Hardware][Vendor]for hardware-specific changes. Vendor name should appear in the prefix (e.g.,[Hardware][AMD]).[Misc]for PRs that do not fit the above categories. Please use this sparingly.Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
format.shto format your code.docs/source/if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.Notes for Large Changes
Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with
rfc-requiredand might not go through the PR.What to Expect for the Reviews
The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
action-requiredlabel on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.Thank You
Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!