[Model] Pipeline parallel support for Mixtral

[comaniac]

Take from #6403. Co-authored by @binxuan

cc @youkaichao

[github-actions]

👋 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 trigger fastcheck CI to run, which consists only a small and essential subset of tests to quickly catch errors with the flexibility to run extra individual tests on top (you can do this by unblocking test steps in the Buildkite run).

Full CI run is still required to merge this PR so once the PR is ready to go, please make sure to run it. If you need all test signals in between PR commits, you can trigger full CI as well.

To run full CI, you can do one of these:

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

🚀

[comaniac]

Tested locally with PP=8 and worked.

[youkaichao]

can you test the correctness locally, using https://github.com/vllm-project/vllm/blob/main/tests/distributed/test_pipeline_parallel.py ?

[comaniac]

Passed with the following configures. Note that I tested it on 8xL4 so I have to use 8 GPUs to host the model.

    "TP_SIZE, PP_SIZE, EAGER_MODE, CHUNKED_PREFILL, MODEL_NAME",
    [
        (2, 4, 0, 1, "mistralai/Mixtral-8x7B-Instruct-v0.1"),
        (2, 4, 1, 0, "mistralai/Mixtral-8x7B-Instruct-v0.1"),
        (1, 8, 0, 1, "mistralai/Mixtral-8x7B-Instruct-v0.1"),
        (1, 8, 1, 0, "mistralai/Mixtral-8x7B-Instruct-v0.1"),
    ])

Also fixed some issues in the test file:

• Use TP_SIZE x PP_SIZE as the TP size of the reference. The current max(TP_SIZE, 2) doesn't work for larger models.
• Do not use 0's as the token ID. This may generate random outputs for certain models/tokenizers.

[youkaichao]

it's not going to work. this will run in multi-node tests with mp backend, and we can use at most 2 GPUs.

you can revert this change, keep it only for your local testing.

[comaniac]

Reverted with comments.

[youkaichao]

LGTM if tests pass