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.

Sign up for GitHub

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

[Model] PP support for embedding models and update docs #9090

Merged
merged 23 commits into from
Oct 6, 2024
Merged

[Model] PP support for embedding models and update docs #9090

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
20b63ee
Add PP support for embeddings models
DarkLight1337 Oct 5, 2024
3a51d81
Add PP support for embedding model runner
DarkLight1337 Oct 5, 2024
49c8216
Test PP embedding models
DarkLight1337 Oct 5, 2024
d995174
Reduce the number of tests
DarkLight1337 Oct 5, 2024
b385f8b
Fix missing PP handling
DarkLight1337 Oct 5, 2024
c6ea4b1
Fix for real
DarkLight1337 Oct 5, 2024
0e1e0b0
Fix wrong default
DarkLight1337 Oct 5, 2024
d139c54
Fix mistral loading
DarkLight1337 Oct 5, 2024
6222fdb
Update imports
DarkLight1337 Oct 5, 2024
d5f2e15
Fix weight loading again
DarkLight1337 Oct 5, 2024
b77ab5c
Update test definitions
DarkLight1337 Oct 5, 2024
1b9f8d4
Fix import
DarkLight1337 Oct 6, 2024
0b6b931
Up the TP size
DarkLight1337 Oct 6, 2024
519f695
Add missing models to docs
DarkLight1337 Oct 6, 2024
59bf961
Add note for tests
DarkLight1337 Oct 6, 2024
dfad8fb
More explicit NOTE
DarkLight1337 Oct 6, 2024
9adc7c0
Fix wrong model name
DarkLight1337 Oct 6, 2024
21212ee
Downgrade to note
DarkLight1337 Oct 6, 2024
79f9c4e
Fix missing org
DarkLight1337 Oct 6, 2024
6e49835
Remove an L
DarkLight1337 Oct 6, 2024
e19fda6
Consolidate multimodal models
DarkLight1337 Oct 6, 2024
4e5cab5
Update docs/source/models/supported_models.rst
DarkLight1337 Oct 6, 2024
90ecb43
Update supported_models.rst
DarkLight1337 Oct 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 108 additions & 29 deletions tests/distributed/test_pipeline_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
"""
import os
from dataclasses import dataclass
from typing import List, NamedTuple, Optional
from typing import List, Literal, NamedTuple, Optional

import pytest

Expand Down Expand Up @@ -104,15 +104,13 @@ def iter_params(self, model_name: str):
# [FAST TESTS]
# Uses Llama
# "BAAI/AquilaChat-7B": PPTestSettings.fast(),
# TODO: Test on larger GPU
# "Snowflake/snowflake-arctic-instruct": PPTestSettings.fast(trust_remote_code=True), # noqa: E501
"Snowflake/snowflake-arctic-instruct": PPTestSettings.fast(tp_base=4, trust_remote_code=True), # noqa: E501
"baichuan-inc/Baichuan-7B": PPTestSettings.fast(trust_remote_code=True),
"baichuan-inc/Baichuan2-13B-Chat": PPTestSettings.fast(trust_remote_code=True), # noqa: E501
"bigscience/bloomz-1b1": PPTestSettings.fast(),
"THUDM/chatglm3-6b": PPTestSettings.fast(trust_remote_code=True),
"CohereForAI/c4ai-command-r-v01": PPTestSettings.fast(tp_base=2, trust_remote_code=True), # noqa: E501
# TODO: Test on larger GPU
# "databricks/dbrx-instruct": PPTestSettings.fast(),
"databricks/dbrx-instruct": PPTestSettings.fast(tp_base=4),
Copy link
Collaborator

Choose a reason for hiding this comment

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

This hardcode of 4 will only work on 8 GPU machines. Might be a bit confusing, or comments should be added to adjust this as necessary.

Copy link
Collaborator

Choose a reason for hiding this comment

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

We might not need all of them as well, maybe a few as examples are sufficient with some comments otherwise needs to be kept up to date.

Copy link
Member Author

Choose a reason for hiding this comment

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

I have added some comments in the code. See if it works for you.

Copy link
Collaborator

Choose a reason for hiding this comment

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

The comments look good to me, but would it be possible to run this test with tp_base of 8? I.e. does this test automatically work with 2 nodes? This doesn't need to block but just something to think about.

Copy link
Member Author

Choose a reason for hiding this comment

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

It requires 16 GPUs to run. What is the setup you used to run those models? I have a comment indicating that tp_base is just an indication of the model size and may have to be adjusted further.

"Deci/DeciLM-7B-instruct": PPTestSettings.fast(trust_remote_code=True),
"deepseek-ai/deepseek-llm-7b-chat": PPTestSettings.fast(),
"deepseek-ai/DeepSeek-V2-Lite-Chat": PPTestSettings.fast(trust_remote_code=True), # noqa: E501
Expand Down Expand Up @@ -161,8 +159,9 @@ def iter_params(self, model_name: str):

EMBEDDING_MODEL_SETTINGS = { # type: ignore[var-annotated]
# [FAST TESTS]
# Uses Llama
# "intfloat/e5-mistral-7b-instruct": PPTestSettings.fast(),
"intfloat/e5-mistral-7b-instruct": PPTestSettings.fast(),
"BAAI/bge-multilingual-gemma2": PPTestSettings.fast(),
"Qwen/Qwen2.5-Math-RM-72B": PPTestSettings.fast(tp_base=4, trust_remote_code=True), # noqa: E501
}

MULTIMODAL_MODEL_SETTINGS = {
Expand Down Expand Up @@ -192,36 +191,31 @@ def iter_params(self, model_name: str):
}
# yapf: enable

MODEL_SETTINGS = {
**GENERATION_MODEL_SETTINGS,
**EMBEDDING_MODEL_SETTINGS,
**MULTIMODAL_MODEL_SETTINGS,
}

# You can update this on your local machine to run specific tests
TEST_MODELS = [
# [LANGUAGE GENERATION]
"meta-llama/Meta-Llama-3-8B",
"facebook/chameleon-7b",
"ibm/PowerLM-3b",
# [LANGUAGE EMBEDDING]
"intfloat/e5-mistral-7b-instruct",
"BAAI/bge-multilingual-gemma2",
# [MULTIMODAL GENERATION]
"OpenGVLab/InternVL2-1B",
"microsoft/Phi-3-vision-128k-instruct",
"mistralai/Pixtral-12B-2409",
"fixie-ai/ultravox-v0_3",
]


@pytest.mark.parametrize(
("model_name", "parallel_setup", "distributed_backend",
"trust_remote_code", "tokenizer_mode"),
[
params for model_name, settings in MODEL_SETTINGS.items()
for params in settings.iter_params(model_name)
if model_name in TEST_MODELS
],
)
@fork_new_process_for_each_test
def test_compare_tp(model_name: str, parallel_setup: ParallelSetup,
distributed_backend: str, trust_remote_code: bool,
tokenizer_mode: Optional[str], num_gpus_available):
def _compare_tp(
model_name: str,
parallel_setup: ParallelSetup,
distributed_backend: str,
trust_remote_code: bool,
tokenizer_mode: Optional[str],
num_gpus_available: int,
*,
method: Literal["generate", "encode"] = "encode",
):
tp_size, pp_size, eager_mode, chunked_prefill = parallel_setup

if num_gpus_available < tp_size:
Expand Down Expand Up @@ -286,10 +280,95 @@ def test_compare_tp(model_name: str, parallel_setup: ParallelSetup,
]

try:
compare_two_settings(model_name, pp_args, tp_args, pp_env)
compare_two_settings(model_name,
pp_args,
tp_args,
pp_env,
method=method)
except Exception:
if pp_env is None:
raise
else:
# Ray ADAG tests are flaky, so we don't want to fail the test
logger.exception("Ray ADAG tests failed")


@pytest.mark.parametrize(
("model_name", "parallel_setup", "distributed_backend",
"trust_remote_code", "tokenizer_mode"),
[
params for model_name, settings in GENERATION_MODEL_SETTINGS.items()
for params in settings.iter_params(model_name)
if model_name in TEST_MODELS
],
)
@fork_new_process_for_each_test
def test_tp_language_generation(
model_name: str,
parallel_setup: ParallelSetup,
distributed_backend: str,
trust_remote_code: bool,
tokenizer_mode: Optional[str],
num_gpus_available,
):
_compare_tp(model_name,
parallel_setup,
distributed_backend,
trust_remote_code,
tokenizer_mode,
num_gpus_available,
method="generate")


@pytest.mark.parametrize(
("model_name", "parallel_setup", "distributed_backend",
"trust_remote_code", "tokenizer_mode"),
[
params for model_name, settings in EMBEDDING_MODEL_SETTINGS.items()
for params in settings.iter_params(model_name)
if model_name in TEST_MODELS
],
)
@fork_new_process_for_each_test
def test_tp_language_embedding(
model_name: str,
parallel_setup: ParallelSetup,
distributed_backend: str,
trust_remote_code: bool,
tokenizer_mode: Optional[str],
num_gpus_available,
):
_compare_tp(model_name,
parallel_setup,
distributed_backend,
trust_remote_code,
tokenizer_mode,
num_gpus_available,
method="encode")


@pytest.mark.parametrize(
("model_name", "parallel_setup", "distributed_backend",
"trust_remote_code", "tokenizer_mode"),
[
params for model_name, settings in MULTIMODAL_MODEL_SETTINGS.items()
for params in settings.iter_params(model_name)
if model_name in TEST_MODELS
],
)
@fork_new_process_for_each_test
def test_tp_multimodal_generation(
model_name: str,
parallel_setup: ParallelSetup,
distributed_backend: str,
trust_remote_code: bool,
tokenizer_mode: Optional[str],
num_gpus_available,
):
_compare_tp(model_name,
parallel_setup,
distributed_backend,
trust_remote_code,
tokenizer_mode,
num_gpus_available,
method="generate")
Loading
Loading