mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-20 05:59:32 +08:00
Keval Morabia
a076e6c218
### What does this PR do? Type of change: Test infra improvement <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> * Instead of spawning multiprocesses per test case (adds ~10s overhead x 150 tests), we spawn once per file (gpu_megatron) or session (gpu) to make the tests much much faster * Remove some unused / unnecessary tests * Fix and enable Minitron NAS test on Hybrid models for Mcore 0.16 Test time (`pytest tests/<test_type>`) on 2x RTXPro 6000 Blackwell: https://github.com/NVIDIA/Model-Optimizer/actions/runs/22625069275 | | gpu | gpu_megatron | |------------|-----|--------------| | **Before** | 27m | 40m | | **After** | 16m | 19m | Full PR-merge CI/CD now only takes ~45mins ### Testing <!-- Mention how have you tested your change if applicable. --> Tested on 1,2,4,8 GPU setup ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, using `torch.load(..., weights_only=True)`, avoiding `pickle`, etc.). - Is this change backward compatible?: ✅ <!--- If ❌, explain why. --> - If you copied code from any other source, did you follow IP policy in [CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?: N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Reworked GPU test infra to use a persistent distributed worker pool with improved per-worker teardown, aggregated error reporting, Megatron-specific cleanup, reduced GPU CI timeouts, and a small build-system dependency change. * **Tests** * Converted many multi-process GPU tests to a fixture-driven distributed runner, updating test entry points and orchestration while preserving test logic; adjusted imports and conditional guards for several GPU tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
115 lines
3.0 KiB
Python
115 lines
3.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import platform
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
from _test_utils.torch.distributed.utils import init_process
|
|
|
|
import modelopt.torch.opt as mto
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def verbose(request):
|
|
return request.config.getoption("verbose")
|
|
|
|
|
|
def pytest_addoption(parser):
|
|
parser.addoption(
|
|
"--run-manual",
|
|
action="store_true",
|
|
default=False,
|
|
help="Run manual tests",
|
|
)
|
|
parser.addoption(
|
|
"--run-release",
|
|
action="store_true",
|
|
default=False,
|
|
help="Run release tests",
|
|
)
|
|
|
|
|
|
def pytest_collection_modifyitems(config, items):
|
|
"""Skip tests with specific markers unless their corresponding flag is provided."""
|
|
skip_marks = [
|
|
("manual", "--run-manual"),
|
|
("release", "--run-release"),
|
|
]
|
|
|
|
for mark_name, option_name in skip_marks:
|
|
if not config.getoption(option_name):
|
|
skipper = pytest.mark.skip(reason=f"Only run when {option_name} is given")
|
|
for item in items:
|
|
if mark_name in item.keywords:
|
|
item.add_marker(skipper)
|
|
|
|
|
|
@pytest.fixture
|
|
def skip_on_windows():
|
|
if platform.system() == "Windows":
|
|
pytest.skip("Skipping on Windows")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def num_gpus():
|
|
return torch.cuda.device_count()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cuda_capability():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA is not available")
|
|
return torch.cuda.get_device_capability()
|
|
|
|
|
|
@pytest.fixture
|
|
def distributed_setup_size_1():
|
|
init_process(rank=0, size=1, backend="nccl")
|
|
yield
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.fixture
|
|
def need_2_gpus():
|
|
if torch.cuda.device_count() < 2:
|
|
pytest.skip("Need at least 2 GPUs to run this test")
|
|
|
|
|
|
@pytest.fixture
|
|
def need_4_gpus():
|
|
if torch.cuda.device_count() < 4:
|
|
pytest.skip("Need at least 4 GPUs to run this test")
|
|
|
|
|
|
@pytest.fixture
|
|
def need_8_gpus():
|
|
if torch.cuda.device_count() < 8:
|
|
pytest.skip("Need at least 8 GPUs to run this test")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def set_torch_dtype(request):
|
|
orig_dtype = torch.get_default_dtype()
|
|
torch.set_default_dtype(request.param)
|
|
yield
|
|
torch.set_default_dtype(orig_dtype)
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def enable_hf_checkpointing():
|
|
mto.enable_huggingface_checkpointing()
|