Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Kernel] Expand MoE weight loading + Add Fused Marlin MoE Kernel #7527

Merged
merged 18 commits into from
Aug 21, 2024

Conversation

dsikka
Copy link
Contributor

@dsikka dsikka commented Aug 14, 2024

Summary

  • Expands weight loading to support grouped and per channel weight quantization. Cleans-up fp8 MoE to use the updated weight loading
  • Adds Marlin Fused MoE Kernel for w4a16 by @ElizaWszola
  • Adds CompressedTensorsMoEMethod to support MoE w4a16 models from llm-compressor and compressed-tensors
  • Tested using 2 and 4 TP with Mixtral

TODO:

  • Test with deepseek v2

Next Steps:

  • The CompressedTensorsMoEMethod is not leveraging the scheme structure in-place for compressed-tensors in order to keep the scope of this PR focused on the kernel + updated weight loading. Will be updated in a follow-up to use the scheme structure

co-authored by @ElizaWszola, from Neural Magic

Copy link

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

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:

  • Comment /ready on the PR
  • Add ready label to the PR
  • Enable auto-merge.

🚀

@dsikka dsikka changed the title [Misc] Expand MoE weight loading [Misc] Expand MoE weight loading/update fp8 MoE Aug 14, 2024
@dsikka dsikka changed the title [Misc] Expand MoE weight loading/update fp8 MoE [Kernel] Expand MoE weight loading + Add Fused Marlin MoE Kernel Aug 16, 2024
@dsikka dsikka marked this pull request as ready for review August 16, 2024 15:54
Copy link
Member

@mgoin mgoin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work, I'm liking the start of refactoring in fused_moe/layer.py

@dsikka
Copy link
Contributor Author

dsikka commented Aug 18, 2024

/ready

@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 18, 2024
@dsikka dsikka force-pushed the compressed-tensors-moe-updated branch from ac3c2b9 to 324ecfd Compare August 19, 2024 15:58
@dsikka dsikka requested a review from mgoin August 19, 2024 17:35
dsikka added 14 commits August 21, 2024 18:42
FILL IN THE PR DESCRIPTION HERE

FIX #xxxx (*link existing issues this PR will resolve*)

**BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE
DESCRIPTION ABOVE**

---

<details>
<!-- inside this <details> section, markdown rendering does not work, so
we use raw html here. -->
<summary><b> PR Checklist (Click to Expand) </b></summary>

<p>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.</p>

<h3>PR Title and Classification</h3>
<p>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:</p>
<ul>
    <li><code>[Bugfix]</code> for bug fixes.</li>
<li><code>[CI/Build]</code> for build or continuous integration
improvements.</li>
<li><code>[Doc]</code> for documentation fixes and improvements.</li>
<li><code>[Model]</code> for adding a new model or improving an existing
model. Model name should appear in the title.</li>
<li><code>[Frontend]</code> For changes on the vLLM frontend (e.g.,
OpenAI API server, <code>LLM</code> class, etc.) </li>
<li><code>[Kernel]</code> for changes affecting CUDA kernels or other
compute kernels.</li>
<li><code>[Core]</code> for changes in the core vLLM logic (e.g.,
<code>LLMEngine</code>, <code>AsyncLLMEngine</code>,
<code>Scheduler</code>, etc.)</li>
<li><code>[Hardware][Vendor]</code> for hardware-specific changes.
Vendor name should appear in the prefix (e.g.,
<code>[Hardware][AMD]</code>).</li>
<li><code>[Misc]</code> for PRs that do not fit the above categories.
Please use this sparingly.</li>
</ul>
<p><strong>Note:</strong> If the PR spans more than one category, please
include all relevant prefixes.</p>

<h3>Code Quality</h3>

<p>The PR need to meet the following code quality standards:</p>

<ul>
<li>We adhere to <a
href="https://google.github.io/styleguide/pyguide.html">Google Python
style guide</a> and <a
href="https://google.github.io/styleguide/cppguide.html">Google C++
style guide</a>.</li>
<li>Pass all linter checks. Please use <a
href="https://github.com/vllm-project/vllm/blob/main/format.sh"><code>format.sh</code></a>
to format your code.</li>
<li>The code need to be well-documented to ensure future contributors
can easily understand the code.</li>
<li>Include sufficient tests to ensure the project to stay correct and
robust. This includes both unit tests and integration tests.</li>
<li>Please add documentation to <code>docs/source/</code> if the PR
modifies the user-facing behaviors of vLLM. It helps vLLM user
understand and utilize the new features or changes.</li>
</ul>

<h3>Notes for Large Changes</h3>
<p>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 <code>rfc-required</code>
and might not go through the PR.</p>

<h3>What to Expect for the Reviews</h3>

<p>The goal of the vLLM team is to be a <i>transparent reviewing
machine</i>. 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: </p>

<ul>
<li> After the PR is submitted, the PR will be assigned to a reviewer.
Every reviewer will pick up the PRs based on their expertise and
availability.</li>
<li> After the PR is assigned, the reviewer will provide status update
every 2-3 days. If the PR is not reviewed within 7 days, please feel
free to ping the reviewer or the vLLM team.</li>
<li> After the review, the reviewer will put an <code>
action-required</code> label on the PR if there are changes required.
The contributor should address the comments and ping the reviewer to
re-review the PR.</li>
<li> Please respond to all comments within a reasonable time frame. If a
comment isn't clear or you disagree with a suggestion, feel free to ask
for clarification or discuss the suggestion.
 </li>
</ul>

<h3>Thank You</h3>

<p> 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! </p>

</details>
@dsikka dsikka force-pushed the compressed-tensors-moe-updated branch from 20b5940 to fe4c9e2 Compare August 21, 2024 18:42
@simon-mo simon-mo merged commit 8678a69 into vllm-project:main Aug 21, 2024
57 of 59 checks passed
mgoin added a commit that referenced this pull request Aug 22, 2024
dsikka added a commit to neuralmagic/nm-vllm that referenced this pull request Aug 22, 2024
@fengyang95
Copy link

I tried to run deepseek-v2-fp8(https://huggingface.co/neuralmagic/DeepSeek-Coder-V2-Instruct-FP8) with this PR and 8 L40 GPUs, and encountered the following error:
model:
VllmWorkerProcess pid=6252) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6256) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6255) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6254) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6257) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6251) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
(VllmWorkerProcess pid=6253) INFO 08-22 22:07:58 model_runner.py:890] Loading model weights took 28.2205 GB
ERROR 08-22 22:07:59 multiproc_worker_utils.py:120] Worker VllmWorkerProcess pid 6252 died, exit code: -15
INFO 08-22 22:07:59 multiproc_worker_utils.py:123] Killing local vLLM worker processes
Process SpawnProcess-1:
Traceback (most recent call last):
File "/usr/lib/python3.9/multiprocessing/process.py", line 315, in _bootstrap
self.run()
File "/usr/lib/python3.9/multiprocessing/process.py", line 108, in run
self._target(*self._args, **self._kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/entrypoints/openai/rpc/server.py", line 230, in run_rpc_server
server = AsyncEngineRPCServer(async_engine_args, usage_context, rpc_path)
File "/usr/local/lib/python3.9/dist-packages/vllm/entrypoints/openai/rpc/server.py", line 31, in init
self.engine = AsyncLLMEngine.from_engine_args(
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/async_llm_engine.py", line 740, in from_engine_args
engine = cls(
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/async_llm_engine.py", line 636, in init
self.engine = self._init_engine(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/async_llm_engine.py", line 840, in _init_engine
return engine_class(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/async_llm_engine.py", line 272, in init
super().init(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/llm_engine.py", line 284, in init
self._initialize_kv_caches()
File "/usr/local/lib/python3.9/dist-packages/vllm/engine/llm_engine.py", line 390, in _initialize_kv_caches
self.model_executor.determine_num_available_blocks())
File "/usr/local/lib/python3.9/dist-packages/vllm/executor/distributed_gpu_executor.py", line 38, in determine_num_available_blocks
num_blocks = self._run_workers("determine_num_available_blocks", )
File "/usr/local/lib/python3.9/dist-packages/vllm/executor/multiproc_gpu_executor.py", line 192, in _run_workers
driver_worker_output = driver_worker_method(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/worker/worker.py", line 222, in determine_num_available_blocks
self.model_runner.profile_run()
File "/usr/local/lib/python3.9/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/worker/model_runner.py", line 1097, in profile_run
self.execute_model(model_input, kv_caches, intermediate_tensors)
File "/usr/local/lib/python3.9/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/worker/model_runner.py", line 1415, in execute_model
hidden_or_intermediate_states = model_executable(
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 504, in forward
hidden_states = self.model(input_ids, positions, kv_caches,
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 461, in forward
hidden_states, residual = layer(positions, hidden_states,
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 401, in forward
hidden_states = self.mlp(hidden_states)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/models/deepseek_v2.py", line 148, in forward
final_hidden_states = self.experts(
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/layers/fused_moe/layer.py", line 287, in forward
final_hidden_states = self.quant_method.apply(
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/layers/quantization/fp8.py", line 499, in apply
return fused_experts(x,
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/layers/fused_moe/fused_moe.py", line 544, in fused_experts
moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/layers/fused_moe/fused_moe.py", line 228, in moe_align_block_size
ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
File "/usr/local/lib/python3.9/dist-packages/vllm/_custom_ops.py", line 28, in wrapper
return fn(*args, **kwargs)
File "/usr/local/lib/python3.9/dist-packages/vllm/_custom_ops.py", line 485, in moe_align_block_size
torch.ops._C.moe_align_block_size(topk_ids, num_experts, block_size,
File "/usr/local/lib/python3.9/dist-packages/torch/ops.py", line 1061, in call
return self
._op(*args, **(kwargs or {}))
RuntimeError: CUDA error: invalid argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

@mgoin
Copy link
Member

mgoin commented Aug 22, 2024

@fengyang95 it looks like your issue is due to incorrectly compiled or linked functions, possibly from older builds. Could you build vllm in a fresh environment?

@fengyang95
Copy link

ooks like your issue is due to incorrectly compiled or linked functions, possibly from older builds. Could you build vllm in a fresh environment?

@mgoin I tried recompiling in a new environment but still encountered the same issue. Has this been tested on ds-v2?

@dsikka
Copy link
Contributor Author

dsikka commented Aug 23, 2024

@fengyang95 Do you mind sharing a code snippet of what you're running?
The scope of the fused marlin kernel introduced in this PR is packed w4a16 models. It does not support FP8 models at the moment.

@fengyang95
Copy link

fengyang95 commented Aug 23, 2024

@fengyang95 Do you mind sharing a code snippet of what you're running? The scope of the fused marlin kernel introduced in this PR is packed w4a16 models. It does not support FP8 models at the moment.
Oh, I was using an FP8 model.
https://huggingface.co/neuralmagic/DeepSeek-Coder-V2-Instruct-FP8
By the way, when will FP8 models be supported?

@dsikka
Copy link
Contributor Author

dsikka commented Aug 23, 2024

@fengyang95 FP8 models will use the triton kernels which should work with the DeepSeek V2 model listed.
Adding FP8 Fused MoE Marlin kernel support is not in scope at the moment.

omrishiv pushed a commit to omrishiv/vllm that referenced this pull request Aug 26, 2024
@fengyang95
Copy link

fengyang95 commented Aug 30, 2024

How should I load a w4a16 version of deepseek-v2 by vllm==0.5.5 that was compressed using llm-compressor?
I used quantization=compressed-tensors, but it throws an error:

File "/usr/local/lib/python3.9/dist-packages/vllm/model_executor/layers/fused_moe/layer.py", line 192, in init
assert self.quant_method is not None

Alvant pushed a commit to compressa-ai/vllm that referenced this pull request Oct 26, 2024
Alvant pushed a commit to compressa-ai/vllm that referenced this pull request Oct 26, 2024
KuntaiDu pushed a commit to KuntaiDu/vllm that referenced this pull request Nov 20, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants