mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-10 09:34:01 +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>
85 lines
2.8 KiB
Python
85 lines
2.8 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.
|
|
|
|
from functools import partial
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from _test_utils.torch.distributed.utils import synchronize_state_dict
|
|
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP # noqa: N817
|
|
from torch.distributed.fsdp.wrap import ModuleWrapPolicy
|
|
|
|
import modelopt.torch.quantization as mtq
|
|
from modelopt.torch.opt.dynamic import _pytorch_managed
|
|
|
|
|
|
def _test_fsdp_simple_linear(rank, size):
|
|
model = nn.Linear(128, 128).cuda(rank)
|
|
inputs = torch.randn(2, 2, 128).cuda(rank)
|
|
|
|
synchronize_state_dict(model)
|
|
|
|
model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
|
|
|
|
manager = model._get_dm_attribute_manager()
|
|
assert "weight" in manager.da_keys()
|
|
assert model._get_dm_attribute_manager().get_da_value("weight") is _pytorch_managed
|
|
|
|
out_ref = model(inputs)
|
|
|
|
fsdp_model = FSDP(model)
|
|
|
|
assert "weight" in manager.da_keys()
|
|
assert isinstance(fsdp_model._get_dm_attribute_manager().get_da_value("weight"), torch.Tensor)
|
|
|
|
# Lets make sure backward works fine
|
|
fsdp_model(inputs).sum().backward()
|
|
|
|
# Lets compare outputs before and after wrapping in FSDP
|
|
out_test = fsdp_model(inputs)
|
|
|
|
assert torch.allclose(out_ref, out_test)
|
|
|
|
|
|
def _test_nested_fsdp(use_orig_params, rank, size):
|
|
model = nn.Sequential(
|
|
nn.Sequential(nn.Linear(128, 128), nn.Linear(128, 128)),
|
|
nn.Sequential(nn.Linear(128, 128), nn.Linear(128, 128)),
|
|
).cuda()
|
|
inputs = torch.randn(2, 2, 128).cuda()
|
|
synchronize_state_dict(model)
|
|
|
|
model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
|
|
out_ref = model(inputs)
|
|
|
|
fsdp_model = FSDP(
|
|
model, auto_wrap_policy=ModuleWrapPolicy((nn.Sequential,)), use_orig_params=use_orig_params
|
|
)
|
|
|
|
fsdp_model(inputs).sum().backward()
|
|
out_test = fsdp_model(inputs)
|
|
|
|
assert torch.allclose(out_ref, out_test)
|
|
|
|
|
|
def test_fsdp_simple_linear(dist_workers):
|
|
dist_workers.run(_test_fsdp_simple_linear)
|
|
|
|
|
|
@pytest.mark.parametrize("use_orig_params", [True, False])
|
|
def test_nested_fsdp(use_orig_params, dist_workers):
|
|
dist_workers.run(partial(_test_nested_fsdp, use_orig_params))
|