diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 95ea9046d5..4c5683cc03 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -11,7 +11,7 @@ /vllm/model_executor/layers/mamba @tdoublep @tomeras91 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn -/vllm/model_executor/layers/batch_invariant.py @yewentao256 +/vllm/model_executor/determinism/ @yewentao256 /vllm/model_executor/custom_op.py @shen-shanshan /vllm/ir @ProExpertProg /vllm/kernels/ @ProExpertProg @tjtanaa diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index 27d3f4b4e6..ea7099962f 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -123,7 +123,7 @@ class TrainModel: """Ray actor that wraps the training model on a dedicated GPU.""" def __init__(self, model_name: str): - from vllm.model_executor.layers.batch_invariant import ( + from vllm.model_executor.determinism.batch_invariant import ( init_batch_invariance, ) diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index 531c5370df..a162511e9e 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -11,8 +11,8 @@ import pytest import torch from utils import skip_unsupported -from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant -from vllm.model_executor.layers.batch_invariant_configs import ( +from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant +from vllm.model_executor.determinism.batch_invariant_configs import ( _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS, _get_tuned_matmul_arch_family, ) diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 232a43b1f9..a9a15d1bf0 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -6,7 +6,7 @@ import pytest import torch from utils import skip_if_not_cuda, skip_unsupported -from vllm.model_executor.layers.batch_invariant import ( +from vllm.model_executor.determinism.batch_invariant import ( rms_norm_batch_invariant, ) from vllm.platforms import current_platform diff --git a/vllm/model_executor/determinism/__init__.py b/vllm/model_executor/determinism/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/model_executor/determinism/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py similarity index 99% rename from vllm/model_executor/layers/batch_invariant.py rename to vllm/model_executor/determinism/batch_invariant.py index 7db629ebad..867b4772a1 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -8,7 +8,7 @@ from typing import Any import torch import vllm.envs as envs -from vllm.model_executor.layers.batch_invariant_configs import ( +from vllm.model_executor.determinism.batch_invariant_configs import ( _get_matmul_config, resolve_tuned_matmul_configs, ) diff --git a/vllm/model_executor/layers/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py similarity index 100% rename from vllm/model_executor/layers/batch_invariant_configs.py rename to vllm/model_executor/determinism/batch_invariant_configs.py diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 5f9e763af0..4a0d8720e4 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -11,7 +11,7 @@ import vllm.kernels # noqa: F401 from vllm import envs, ir from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant +from vllm.model_executor.determinism.batch_invariant import rms_norm_batch_invariant logger = init_logger(__name__) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 1e31a64d68..e723d6fee4 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -21,7 +21,7 @@ from vllm.distributed import ( ) from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.batch_invariant import ( +from vllm.model_executor.determinism.batch_invariant import ( linear_batch_invariant, ) from vllm.model_executor.layers.quantization.base_config import ( diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 79695084e4..3057270c62 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -16,7 +16,7 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, ) from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.batch_invariant import ( +from vllm.model_executor.determinism.batch_invariant import ( linear_batch_invariant, ) from vllm.model_executor.layers.quantization.base_config import ( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 05c03af75e..50417ce6eb 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1416,7 +1416,7 @@ def init_worker_distributed_environment( ) -> None: """Initialize the distributed environment.""" parallel_config = vllm_config.parallel_config - from vllm.model_executor.layers.batch_invariant import init_batch_invariance + from vllm.model_executor.determinism.batch_invariant import init_batch_invariance init_batch_invariance() override_envs_for_eplb(