mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-12-23 23:30:36 +08:00
[test] add allreduce test on npu (#9619)
Co-authored-by: frozenleaves <frozen@Mac.local>
This commit is contained in:
浮梦
GitHub
93
tests_v1/accelerator/test_allreduce.py
Normal file
93
tests_v1/accelerator/test_allreduce.py
Normal file
@@ -0,0 +1,93 @@
|
||||
# Copyright 2025 the LlamaFactory team.
|
||||
#
|
||||
# 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 pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from llamafactory.v1.accelerator.helper import ReduceOp, all_reduce, is_torch_cuda_available, is_torch_npu_available
|
||||
from llamafactory.v1.utils.utils import find_available_port
|
||||
|
||||
|
||||
def _dist_worker(rank, world_size):
|
||||
if is_torch_cuda_available():
|
||||
backend = "nccl"
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(rank)
|
||||
elif is_torch_npu_available():
|
||||
backend = "hccl"
|
||||
device = torch.device(f"npu:{rank}")
|
||||
torch.npu.set_device(rank)
|
||||
else:
|
||||
backend = "gloo"
|
||||
device = torch.device("cpu")
|
||||
|
||||
dist.init_process_group(
|
||||
backend=backend,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
|
||||
# --------------------
|
||||
# Test all_reduce SUM
|
||||
# --------------------
|
||||
y = torch.tensor(rank + 1.0, device=device)
|
||||
y_sum = all_reduce(y.clone(), op=ReduceOp.SUM)
|
||||
assert y_sum.item() == 3.0
|
||||
|
||||
# --------------------
|
||||
# Test all_reduce MEAN
|
||||
# --------------------
|
||||
y_mean = all_reduce(y.clone(), op=ReduceOp.MEAN)
|
||||
assert y_mean.item() == pytest.approx(1.5)
|
||||
|
||||
# --------------------
|
||||
# Test all_reduce MAX
|
||||
# --------------------
|
||||
y_max = all_reduce(y.clone(), op=ReduceOp.MAX)
|
||||
assert y_max.item() == 2.0
|
||||
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.runs_on(["npu", "cuda"])
|
||||
@pytest.mark.require_distributed(2)
|
||||
def test_distributed_ops(monkeypatch):
|
||||
monkeypatch.setenv("MASTER_ADDR", "127.0.0.1")
|
||||
monkeypatch.setenv("MASTER_PORT", str(find_available_port()))
|
||||
WORLD_SIZE = 2
|
||||
mp.spawn(
|
||||
_dist_worker,
|
||||
args=(WORLD_SIZE,),
|
||||
nprocs=WORLD_SIZE,
|
||||
join=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.runs_on(["npu", "cuda"])
|
||||
@pytest.mark.require_distributed(4)
|
||||
def test_required_multi():
|
||||
# test require_distributed mark ok
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.runs_on(["npu", "cuda"])
|
||||
@pytest.mark.require_distributed(999)
|
||||
def test_required_invalid():
|
||||
# test require_distributed mark not ok,
|
||||
raise RuntimeError(
|
||||
"this case should not be run, please check whether the require_distributed mark implementation is correct"
|
||||
)
|
||||
@@ -12,18 +12,147 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""LLaMA-Factory test configuration.
|
||||
|
||||
Contains shared fixtures, pytest configuration, and custom markers.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from pytest import Config, Item
|
||||
|
||||
from llamafactory.train.test_utils import patch_valuehead_model
|
||||
from llamafactory.v1.accelerator.helper import get_current_device, get_device_count
|
||||
from llamafactory.v1.utils.packages import is_transformers_version_greater_than
|
||||
from llamafactory.v1.utils.utils import is_env_enabled
|
||||
|
||||
|
||||
try:
|
||||
CURRENT_DEVICE = get_current_device().type # cpu | cuda | npu
|
||||
except Exception:
|
||||
CURRENT_DEVICE = "cpu"
|
||||
|
||||
|
||||
def pytest_configure(config: Config):
|
||||
"""Register custom pytest markers."""
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"slow: marks tests as slow (deselect with '-m \"not slow\"' or set RUN_SLOW=1 to run)",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"runs_on: test requires specific device type, e.g., @pytest.mark.runs_on(['cuda'])",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"require_distributed(num_devices): allow multi-device execution (default: 2)",
|
||||
)
|
||||
|
||||
|
||||
def _handle_runs_on(items: list[Item]):
|
||||
"""Skip tests on specified device TYPES (cpu/cuda/npu)."""
|
||||
for item in items:
|
||||
marker = item.get_closest_marker("runs_on")
|
||||
if not marker:
|
||||
continue
|
||||
|
||||
devices = marker.args[0]
|
||||
if isinstance(devices, str):
|
||||
devices = [devices]
|
||||
|
||||
if CURRENT_DEVICE not in devices:
|
||||
item.add_marker(pytest.mark.skip(reason=f"test requires one of {devices} (current: {CURRENT_DEVICE})"))
|
||||
|
||||
|
||||
def _handle_slow_tests(items: list[Item]):
|
||||
"""Skip slow tests unless RUN_SLOW is enabled."""
|
||||
if not is_env_enabled("RUN_SLOW", "0"):
|
||||
skip_slow = pytest.mark.skip(reason="slow test (set RUN_SLOW=1 to run)")
|
||||
for item in items:
|
||||
if "slow" in item.keywords:
|
||||
item.add_marker(skip_slow)
|
||||
|
||||
|
||||
def _get_visible_devices_env():
|
||||
"""Return device visibility env var name."""
|
||||
if CURRENT_DEVICE == "cuda":
|
||||
return "CUDA_VISIBLE_DEVICES"
|
||||
if CURRENT_DEVICE == "npu":
|
||||
return "ASCEND_RT_VISIBLE_DEVICES"
|
||||
return None
|
||||
|
||||
|
||||
def _handle_device_visibility(items: list[Item]):
|
||||
"""Handle device visibility based on test markers."""
|
||||
env_key = _get_visible_devices_env()
|
||||
if env_key is None or CURRENT_DEVICE == "cpu":
|
||||
return
|
||||
|
||||
# Parse visible devices
|
||||
visible_devices_env = os.environ.get(env_key)
|
||||
if visible_devices_env is None:
|
||||
available = get_device_count()
|
||||
else:
|
||||
visible_devices = [v for v in visible_devices_env.split(",") if v != ""]
|
||||
available = len(visible_devices)
|
||||
|
||||
for item in items:
|
||||
marker = item.get_closest_marker("require_distributed")
|
||||
if not marker:
|
||||
continue
|
||||
|
||||
required = marker.args[0] if marker.args else 2
|
||||
if available < required:
|
||||
item.add_marker(pytest.mark.skip(reason=f"test requires {required} devices, but only {available} visible"))
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config: Config, items: list[Item]):
|
||||
if is_transformers_version_greater_than("4.57.0"):
|
||||
"""Modify test collection based on markers and environment."""
|
||||
# Handle version compatibility (from HEAD)
|
||||
if not is_transformers_version_greater_than("4.57.0"):
|
||||
skip_bc = pytest.mark.skip(reason="Skip backward compatibility tests")
|
||||
for item in items:
|
||||
if "tests_v1" in str(item.fspath):
|
||||
item.add_marker(skip_bc)
|
||||
|
||||
_handle_slow_tests(items)
|
||||
_handle_runs_on(items)
|
||||
_handle_device_visibility(items)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _manage_distributed_env(request, monkeypatch):
|
||||
"""Set environment variables for distributed tests if specific devices are requested."""
|
||||
env_key = _get_visible_devices_env()
|
||||
if not env_key:
|
||||
return
|
||||
|
||||
skip_bc = pytest.mark.skip(reason="Skip backward compatibility tests")
|
||||
# Save old environment for logic checks, monkeypatch handles restoration
|
||||
old_value = os.environ.get(env_key)
|
||||
|
||||
for item in items:
|
||||
if "tests_v1" in str(item.fspath):
|
||||
item.add_marker(skip_bc)
|
||||
marker = request.node.get_closest_marker("require_distributed")
|
||||
if marker:
|
||||
# Distributed test
|
||||
required = marker.args[0] if marker.args else 2
|
||||
specific_devices = marker.args[1] if len(marker.args) > 1 else None
|
||||
|
||||
if specific_devices:
|
||||
devices_str = ",".join(map(str, specific_devices))
|
||||
else:
|
||||
devices_str = ",".join(str(i) for i in range(required))
|
||||
|
||||
monkeypatch.setenv(env_key, devices_str)
|
||||
else:
|
||||
# Non-distributed test
|
||||
if old_value:
|
||||
visible_devices = [v for v in old_value.split(",") if v != ""]
|
||||
monkeypatch.setenv(env_key, visible_devices[0] if visible_devices else "0")
|
||||
else:
|
||||
monkeypatch.setenv(env_key, "0")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fix_valuehead_cpu_loading():
|
||||
"""Fix valuehead model loading."""
|
||||
patch_valuehead_model()
|
||||
|
||||
Reference in New Issue
Block a user