wehub-resource-sync
860 lines
33 KiB
Python
860 lines
33 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
import os
|
|
from contextlib import redirect_stderr
|
|
from dataclasses import asdict
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import yaml
|
|
from lightning import Fabric
|
|
from lightning.fabric.loggers import CSVLogger, TensorBoardLogger
|
|
from lightning.fabric.plugins import BitsandbytesPrecision
|
|
from lightning.pytorch.loggers import LitLogger, MLFlowLogger, WandbLogger
|
|
|
|
from litgpt import GPT
|
|
from litgpt.args import TrainArgs
|
|
from litgpt.constants import (
|
|
_LITLOGGER_AVAILABLE,
|
|
_MLFLOW_AVAILABLE,
|
|
_MLFLOW_SKINNY_AVAILABLE,
|
|
_TENSORBOARD_AVAILABLE,
|
|
_WANDB_AVAILABLE,
|
|
)
|
|
from litgpt.parser_config import save_hyperparameters
|
|
from litgpt.utils import (
|
|
CLI,
|
|
CycleIterator,
|
|
_RunIf,
|
|
capture_hparams,
|
|
check_file_size_on_cpu_and_warn,
|
|
check_nvlink_connectivity,
|
|
check_valid_checkpoint_dir,
|
|
choose_logger,
|
|
chunked_cross_entropy,
|
|
copy_config_files,
|
|
extend_checkpoint_dir,
|
|
find_resume_path,
|
|
fix_and_load_json,
|
|
incremental_save,
|
|
init_out_dir,
|
|
instantiate_bnb_optimizer,
|
|
instantiate_torch_optimizer,
|
|
num_parameters,
|
|
parse_devices,
|
|
select_sft_generate_example,
|
|
)
|
|
|
|
|
|
# match fails on windows. why did they have to use backslashes?
|
|
@_RunIf(skip_windows=True)
|
|
def test_check_valid_checkpoint_dir(tmp_path):
|
|
os.chdir(tmp_path)
|
|
|
|
out = StringIO()
|
|
with pytest.raises(SystemExit), redirect_stderr(out):
|
|
check_valid_checkpoint_dir(tmp_path)
|
|
out = out.getvalue().strip()
|
|
expected = f"""
|
|
checkpoint_dir '{str(tmp_path.absolute())}' is missing the files: ['lit_model.pth', 'model_config.yaml', 'tokenizer.json OR tokenizer.model', 'tokenizer_config.json'].
|
|
Find download instructions at https://github.com/Lightning-AI/litgpt/blob/main/tutorials
|
|
|
|
See all download options by running:
|
|
litgpt download
|
|
""".strip()
|
|
assert out == expected
|
|
|
|
out = StringIO()
|
|
checkpoint_dir = tmp_path / "checkpoints" / "stabilityai" / "stablelm-base-alpha-3b"
|
|
with pytest.raises(SystemExit), redirect_stderr(out):
|
|
check_valid_checkpoint_dir(checkpoint_dir)
|
|
out = out.getvalue().strip()
|
|
expected = f"""
|
|
checkpoint_dir '{str(checkpoint_dir.absolute())}' is not a checkpoint directory.
|
|
Find download instructions at https://github.com/Lightning-AI/litgpt/blob/main/tutorials
|
|
|
|
See all download options by running:
|
|
litgpt download
|
|
""".strip()
|
|
assert out == expected
|
|
|
|
out = StringIO()
|
|
checkpoint_dir.mkdir(parents=True)
|
|
foo_checkpoint_dir = tmp_path / "foo"
|
|
with pytest.raises(SystemExit), redirect_stderr(out):
|
|
check_valid_checkpoint_dir(foo_checkpoint_dir)
|
|
out = out.getvalue().strip()
|
|
expected = f"""
|
|
checkpoint_dir '{str(foo_checkpoint_dir.absolute())}' is not a checkpoint directory.
|
|
Find download instructions at https://github.com/Lightning-AI/litgpt/blob/main/tutorials
|
|
|
|
You have downloaded locally:
|
|
'{str(checkpoint_dir.absolute())}'
|
|
|
|
See all download options by running:
|
|
litgpt download
|
|
""".strip()
|
|
assert out == expected
|
|
|
|
|
|
def test_incremental_write(tmp_path):
|
|
sd = {str(k): torch.randn(5, 10) for k in range(3)}
|
|
sd["0"].someattr = 1
|
|
sd_expected = {k: v.clone() for k, v in sd.items()}
|
|
fn = str(tmp_path / "test.pt")
|
|
with incremental_save(fn) as f:
|
|
sd["0"] = f.store_early(sd["0"])
|
|
sd["2"] = f.store_early(sd["2"])
|
|
f.save(sd)
|
|
sd_actual = torch.load(fn)
|
|
assert sd_actual.keys() == sd_expected.keys()
|
|
assert sd_actual["0"].someattr == 1 # requires PyTorch 2.0+
|
|
for k, v_expected in sd_expected.items():
|
|
v_actual = sd_actual[k]
|
|
torch.testing.assert_close(v_expected, v_actual)
|
|
sd_actual = torch.load(fn, weights_only=True)
|
|
assert sd_actual.keys() == sd_expected.keys()
|
|
assert sd_actual["0"].someattr == 1 # requires PyTorch 2.0+
|
|
for k, v_expected in sd_expected.items():
|
|
v_actual = sd_actual[k]
|
|
torch.testing.assert_close(v_expected, v_actual)
|
|
|
|
|
|
@pytest.mark.parametrize("B", (1, 2))
|
|
@pytest.mark.parametrize("ignore_index", (None, -1, -2, -100))
|
|
def test_chunked_cross_entropy(ignore_index, B):
|
|
V = 50
|
|
T = 25
|
|
regular_logits = torch.randn(B, T, V)
|
|
targets = torch.randint(0, V, (B, T))
|
|
|
|
if ignore_index is not None:
|
|
targets[:, [1, 4, 10, 19]] = ignore_index
|
|
|
|
baseline_loss = F.cross_entropy(
|
|
regular_logits.reshape(-1, regular_logits.size(-1)),
|
|
targets.reshape(-1),
|
|
ignore_index=(ignore_index if ignore_index is not None else -100),
|
|
)
|
|
|
|
ignore_index = ignore_index if ignore_index is not None else -100
|
|
regular_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=0, ignore_index=ignore_index)
|
|
assert torch.equal(baseline_loss, regular_loss)
|
|
assert regular_loss.numel() == 1
|
|
|
|
chunked_loss = chunked_cross_entropy(regular_logits, targets, chunk_size=10, ignore_index=ignore_index)
|
|
torch.testing.assert_close(chunked_loss, regular_loss)
|
|
torch.testing.assert_close(chunked_loss, baseline_loss)
|
|
|
|
logit_chunk_size = 6
|
|
assert T % logit_chunk_size != 0 # ensure leftover
|
|
chunked_logits = list(regular_logits.split(logit_chunk_size, dim=1))
|
|
chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=0, ignore_index=ignore_index)
|
|
torch.testing.assert_close(chunked_loss, regular_loss)
|
|
torch.testing.assert_close(chunked_loss, baseline_loss)
|
|
|
|
chunked_loss = chunked_cross_entropy(chunked_logits, targets, chunk_size=10, ignore_index=ignore_index)
|
|
torch.testing.assert_close(chunked_loss, regular_loss)
|
|
torch.testing.assert_close(chunked_loss, baseline_loss)
|
|
|
|
|
|
def test_num_parameters():
|
|
model = torch.nn.Linear(2, 2)
|
|
assert num_parameters(model) == 6
|
|
assert num_parameters(model, requires_grad=True) == 6
|
|
assert num_parameters(model, requires_grad=False) == 0
|
|
|
|
model = torch.nn.Linear(2, 2)
|
|
model.bias.requires_grad = False
|
|
assert num_parameters(model) == 6
|
|
assert num_parameters(model, requires_grad=True) == 4
|
|
assert num_parameters(model, requires_grad=False) == 2
|
|
|
|
|
|
@_RunIf(min_cuda_gpus=1)
|
|
@pytest.mark.parametrize("mode", ["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"])
|
|
def test_num_parameters_bitsandbytes(mode):
|
|
plugin = BitsandbytesPrecision(mode=mode)
|
|
fabric = Fabric(plugins=plugin, accelerator="cuda", devices=1)
|
|
|
|
model = torch.nn.Linear(10, 10)
|
|
model = fabric.setup(model)
|
|
assert num_parameters(model) == 110
|
|
|
|
with fabric.init_module(empty_init=True):
|
|
model = GPT.from_name("pythia-14m")
|
|
assert num_parameters(model) == 14067712
|
|
|
|
|
|
def test_cycle_iterator():
|
|
iterator = CycleIterator([])
|
|
with pytest.raises(StopIteration):
|
|
next(iterator)
|
|
|
|
iterator = CycleIterator(range(3))
|
|
assert iterator.epoch == 0
|
|
assert next(iterator) == 0
|
|
assert iterator.epoch == 0
|
|
assert next(iterator) == 1
|
|
assert iterator.epoch == 0
|
|
assert next(iterator) == 2
|
|
assert iterator.epoch == 0
|
|
assert next(iterator) == 0
|
|
assert iterator.epoch == 1
|
|
|
|
|
|
def test_parse_devices():
|
|
with pytest.raises(ValueError, match="must be 'auto' or a positive integer"):
|
|
assert parse_devices(0)
|
|
with pytest.raises(ValueError, match="must be 'auto' or a positive integer"):
|
|
assert parse_devices(-2)
|
|
|
|
with mock.patch("litgpt.utils.torch.cuda.device_count", return_value=0):
|
|
assert parse_devices("auto") == 1 # CPU
|
|
assert parse_devices(10) == 10 # leave validation up to Fabric later on
|
|
with mock.patch("litgpt.utils.torch.cuda.device_count", return_value=1):
|
|
assert parse_devices("auto") == 1 # CUDA
|
|
with mock.patch("litgpt.utils.torch.cuda.device_count", return_value=3):
|
|
assert parse_devices("auto") == 3
|
|
assert parse_devices(-1) == 3
|
|
|
|
assert parse_devices(5) == 5
|
|
|
|
|
|
def test_copy_config_files(fake_checkpoint_dir, tmp_path):
|
|
copy_config_files(fake_checkpoint_dir, tmp_path)
|
|
expected = {"model_config.yaml", "tokenizer_config.json", "tokenizer.json"}
|
|
contents = set(os.listdir(tmp_path))
|
|
assert expected.issubset(contents)
|
|
|
|
|
|
def test_capture_hparams():
|
|
integer = 1
|
|
string = "string"
|
|
boolean = True
|
|
none = None
|
|
path = Path("/path")
|
|
dataclass = TrainArgs()
|
|
other = torch.nn.Linear(1, 1)
|
|
hparams = capture_hparams()
|
|
assert hparams == {
|
|
"integer": integer,
|
|
"string": string,
|
|
"boolean": boolean,
|
|
"none": none,
|
|
"path": path,
|
|
"dataclass": asdict(dataclass),
|
|
"other": str(other),
|
|
}
|
|
|
|
|
|
def _test_function(out_dir: Path, foo: bool = False, bar: int = 1):
|
|
save_hyperparameters(_test_function, out_dir)
|
|
|
|
|
|
def test_save_hyperparameters(tmp_path):
|
|
with mock.patch("sys.argv", ["any.py", str(tmp_path), "--foo", "True"]):
|
|
CLI(_test_function)
|
|
|
|
with open(tmp_path / "hyperparameters.yaml", encoding="utf-8") as file:
|
|
hparams = yaml.full_load(file)
|
|
|
|
assert hparams["out_dir"] == str(tmp_path)
|
|
assert hparams["foo"] is True
|
|
assert hparams["bar"] == 1
|
|
|
|
|
|
def _test_function2(out_dir: Path, foo: bool = False, bar: int = 1):
|
|
assert False, "I only exist as a signature, but I should not run."
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"command",
|
|
[
|
|
"any.py",
|
|
"litgpt finetune",
|
|
"litgpt finetune_full",
|
|
"litgpt finetune_lora",
|
|
"litgpt finetune_adapter",
|
|
"litgpt finetune_adapter_v2",
|
|
"litgpt pretrain",
|
|
],
|
|
)
|
|
def test_save_hyperparameters_known_commands(command, tmp_path):
|
|
with mock.patch("sys.argv", [*command.split(" "), str(tmp_path), "--foo", "True"]):
|
|
save_hyperparameters(_test_function2, tmp_path)
|
|
|
|
with open(tmp_path / "hyperparameters.yaml", encoding="utf-8") as file:
|
|
hparams = yaml.full_load(file)
|
|
|
|
assert hparams["out_dir"] == str(tmp_path)
|
|
assert hparams["foo"] is True
|
|
assert hparams["bar"] == 1
|
|
|
|
|
|
def test_choose_logger(tmp_path):
|
|
assert isinstance(choose_logger("csv", out_dir=tmp_path, name="csv"), CSVLogger)
|
|
if _TENSORBOARD_AVAILABLE:
|
|
assert isinstance(choose_logger("tensorboard", out_dir=tmp_path, name="tb"), TensorBoardLogger)
|
|
if _WANDB_AVAILABLE:
|
|
assert isinstance(choose_logger("wandb", out_dir=tmp_path, name="wandb"), WandbLogger)
|
|
if _MLFLOW_AVAILABLE or _MLFLOW_SKINNY_AVAILABLE:
|
|
assert isinstance(choose_logger("mlflow", out_dir=tmp_path, name="wandb"), MLFlowLogger)
|
|
if _LITLOGGER_AVAILABLE:
|
|
assert isinstance(choose_logger("litlogger", out_dir=tmp_path, name="litlogger"), LitLogger)
|
|
with pytest.raises(ValueError, match="`--logger_name=foo` is not a valid option."):
|
|
choose_logger("foo", out_dir=tmp_path, name="foo")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path_type, input_path, expected",
|
|
[
|
|
("relative", "some/relative/path", "some/relative/path"),
|
|
("absolute", "/usr/absolute/path", "/usr/absolute/path"),
|
|
("env_relative", "some/relative/path", "prefix/some/relative/path"),
|
|
("env_absolute", "/usr/absolute/path", "/usr/absolute/path"),
|
|
],
|
|
)
|
|
def test_init_out_dir(path_type, input_path, expected):
|
|
if path_type.startswith("env_"):
|
|
with mock.patch.dict(os.environ, {"LIGHTNING_ARTIFACTS_DIR": "prefix"}):
|
|
result = init_out_dir(input_path)
|
|
assert result == Path(expected), f"Failed for {path_type} with input {input_path} (result {result})"
|
|
else:
|
|
result = init_out_dir(input_path)
|
|
if "LIGHTNING_ARTIFACTS_DIR" not in os.environ:
|
|
assert result == Path(expected), f"Failed for {path_type} with input {input_path} (result {result})"
|
|
else:
|
|
assert result == Path(os.getenv("LIGHTNING_ARTIFACTS_DIR")) / expected, (
|
|
f"Failed for {path_type} with input {input_path} (result {result})"
|
|
)
|
|
|
|
|
|
def test_find_resume_path(tmp_path):
|
|
assert find_resume_path(resume=None, out_dir=Path("does/not/exist")) is None
|
|
assert find_resume_path(resume=Path("does/not/exist"), out_dir=Path("does/not/matter")) == Path("does/not/exist")
|
|
assert find_resume_path(resume=(tmp_path / "checkpoint.pt"), out_dir=Path("does/not/matter")) == (
|
|
tmp_path / "checkpoint.pt"
|
|
)
|
|
|
|
# `resume='auto'` does not enforce the checkpoint to exist
|
|
assert find_resume_path(resume="auto", out_dir=Path("does/not/exist")) is None
|
|
|
|
# `resume=True` requires a checkpoint to exist
|
|
with pytest.raises(FileNotFoundError, match="You passed `--resume=True`, but no checkpoint file was found"):
|
|
find_resume_path(resume=True, out_dir=Path("does/not/exist"))
|
|
with pytest.raises(FileNotFoundError, match="You passed `--resume=True`, but no checkpoint file was found"):
|
|
find_resume_path(resume=True, out_dir=tmp_path)
|
|
|
|
(tmp_path / "step-001").mkdir()
|
|
(tmp_path / "step-001" / "lit_model.pth").touch()
|
|
(tmp_path / "step-002").mkdir()
|
|
(tmp_path / "step-002" / "lit_model.pth").touch()
|
|
(tmp_path / "step-003").mkdir()
|
|
(tmp_path / "step-003" / "lit_model.pth").touch()
|
|
|
|
assert find_resume_path(resume=True, out_dir=tmp_path) == (tmp_path / "step-003" / "lit_model.pth")
|
|
assert find_resume_path(resume="auto", out_dir=tmp_path) == (tmp_path / "step-003" / "lit_model.pth")
|
|
|
|
|
|
@pytest.fixture
|
|
def model_parameters():
|
|
return [torch.nn.Parameter(torch.randn(2, 2))]
|
|
|
|
|
|
def test_instantiate_bnb_optimizer_with_str(model_parameters):
|
|
import bitsandbytes as bnb
|
|
|
|
with mock.patch("litgpt.utils.get_argument_names", return_value={"lr", "eps", "weight_decay"}):
|
|
optimizer = instantiate_bnb_optimizer("AdamW", model_parameters)
|
|
assert isinstance(optimizer, bnb.optim.adamw.PagedAdamW)
|
|
|
|
|
|
def test_instantiate_bnb_optimizer_with_dict(model_parameters):
|
|
import bitsandbytes as bnb
|
|
|
|
optimizer_dict = {"class_path": "AdamW", "init_args": {"lr": 0.01}}
|
|
with mock.patch("litgpt.utils.get_argument_names", return_value={"lr", "eps", "weight_decay"}):
|
|
optimizer = instantiate_bnb_optimizer(optimizer_dict, model_parameters)
|
|
assert isinstance(optimizer, bnb.optim.adamw.PagedAdamW)
|
|
assert optimizer.param_groups[0]["lr"] == 0.01
|
|
|
|
|
|
def test_instantiate_bnb_optimizer_with_invalid_str(model_parameters):
|
|
with pytest.raises(ValueError, match="only supports the AdamW"):
|
|
instantiate_bnb_optimizer("SGD", model_parameters)
|
|
|
|
|
|
def test_instantiate_torch_optimizer_with_str(model_parameters):
|
|
optimizer = instantiate_torch_optimizer("Adam", model_parameters, lr=0.01)
|
|
assert isinstance(optimizer, torch.optim.Adam)
|
|
assert optimizer.param_groups[0]["lr"] == 0.01
|
|
|
|
|
|
def test_instantiate_torch_optimizer_with_class(model_parameters):
|
|
optimizer = instantiate_torch_optimizer(
|
|
{"class_path": "torch.optim.Adam", "init_args": {"lr": 123}}, model_parameters, lr=0.02
|
|
)
|
|
assert isinstance(optimizer, torch.optim.Adam)
|
|
# init args gets overridden
|
|
assert optimizer.param_groups[0]["lr"] == 0.02
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_path, expected",
|
|
[
|
|
(Path("checkpoints/my_model"), Path("checkpoints/my_model")),
|
|
(Path("checkpoints/my_model"), Path("./checkpoints/my_model")),
|
|
],
|
|
)
|
|
def test_extend_checkpoint_dir_is_prefixed(input_path, expected):
|
|
original_dir = Path.cwd() # Save the current directory
|
|
with TemporaryDirectory() as tmp_dir:
|
|
os.chdir(tmp_dir)
|
|
|
|
try:
|
|
if not input_path.is_absolute():
|
|
input_path = Path(tmp_dir) / input_path
|
|
if not expected.is_absolute():
|
|
expected = Path(tmp_dir) / expected
|
|
input_path.parent.mkdir(parents=True, exist_ok=True)
|
|
input_path.touch(exist_ok=True)
|
|
assert extend_checkpoint_dir(input_path) == expected
|
|
finally:
|
|
os.chdir(original_dir) # Reset the current directory
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_path, expected",
|
|
[
|
|
(Path("my_model"), Path("checkpoints/my_model")),
|
|
(Path("my_model"), Path("./checkpoints/my_model")),
|
|
],
|
|
)
|
|
def test_extend_checkpoint_dir(input_path, expected):
|
|
original_dir = Path.cwd() # Save the current directory
|
|
with TemporaryDirectory() as tmp_dir:
|
|
os.chdir(tmp_dir)
|
|
|
|
try:
|
|
if not input_path.is_absolute():
|
|
input_path = Path(tmp_dir) / "checkpoints" / input_path
|
|
if not expected.is_absolute():
|
|
expected = Path(tmp_dir) / expected
|
|
input_path.parent.mkdir(parents=True, exist_ok=True)
|
|
input_path.touch(exist_ok=True)
|
|
assert extend_checkpoint_dir(input_path) == expected
|
|
finally:
|
|
os.chdir(original_dir) # Reset the current directory
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_path, expected",
|
|
[
|
|
(Path("my_model"), Path("my_model")),
|
|
(Path("/my_model"), Path("/my_model")),
|
|
],
|
|
)
|
|
def test_extend_checkpoint_dir_dont_exist(input_path, expected):
|
|
assert extend_checkpoint_dir(input_path) == expected
|
|
|
|
|
|
def test_file_size_below_limit_on_cpu():
|
|
# Test file size below limit on CPU
|
|
with NamedTemporaryFile() as temp_file:
|
|
with mock.patch("os.path.getsize", return_value=4_000_000_000):
|
|
size = check_file_size_on_cpu_and_warn(temp_file.name, "cpu")
|
|
assert size == 4_000_000_000
|
|
|
|
|
|
def test_file_size_above_limit_on_cpu():
|
|
# Test file size above limit on CPU
|
|
with NamedTemporaryFile() as temp_file:
|
|
with mock.patch("os.path.getsize", return_value=4_600_000_000):
|
|
with pytest.warns(UserWarning) as record:
|
|
size = check_file_size_on_cpu_and_warn(temp_file.name, "cpu")
|
|
assert size == 4_600_000_000
|
|
assert "over 4.2 GB" in str(record[0].message)
|
|
|
|
|
|
def test_file_size_above_limit_on_gpu():
|
|
# Test file size above limit on GPU should not warn
|
|
with NamedTemporaryFile() as temp_file:
|
|
with mock.patch("os.path.getsize", return_value=4_600_000_000):
|
|
size = check_file_size_on_cpu_and_warn(temp_file.name, "gpu")
|
|
assert size == 4_600_000_000
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_cuda_is_available_true(monkeypatch):
|
|
"""Fixture to mock torch.cuda.is_available() to return True."""
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_nvidia_device_properties(monkeypatch):
|
|
"""Fixture to mock torch.cuda.get_device_properties() for NVIDIA GPUs."""
|
|
mock_device_properties = mock.MagicMock(name="GPU Device", spec=["name"])
|
|
mock_device_properties.name = "NVIDIA RTX A6000"
|
|
monkeypatch.setattr(torch.cuda, "get_device_properties", lambda idx: mock_device_properties)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_amd_device_properties(monkeypatch):
|
|
"""Fixture to mock torch.cuda.get_device_properties() for AMD GPUs."""
|
|
mock_device_properties = mock.MagicMock(name="GPU Device", spec=["name"])
|
|
mock_device_properties.name = "AMD Instinct MI250X"
|
|
monkeypatch.setattr(torch.cuda, "get_device_properties", lambda idx: mock_device_properties)
|
|
|
|
|
|
@pytest.fixture
|
|
def all_nvlink_connected_output():
|
|
return mock.MagicMock(
|
|
stdout=""" GPU0 GPU1 GPU2 GPU3
|
|
GPU0 X NV12 NV12 NV12
|
|
GPU1 NV12 X NV12 NV12
|
|
GPU2 NV12 NV12 X NV12
|
|
GPU3 NV12 NV12 NV12 X""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_all_nvlink_connected(
|
|
mock_run, all_nvlink_connected_output, mock_cuda_is_available_true, mock_nvidia_device_properties
|
|
):
|
|
mock_run.return_value = all_nvlink_connected_output
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("All GPUs are fully connected via NVLink.")
|
|
|
|
|
|
@pytest.fixture
|
|
def nvlink_partially_connected_output():
|
|
return mock.MagicMock(
|
|
stdout=""" GPU0 GPU1 GPU2 GPU3 CPU Affinity
|
|
GPU0 X NV1 SYS SYS 0-7
|
|
GPU1 NV1 X SYS SYS 0-7
|
|
GPU2 SYS SYS X NV1 8-15
|
|
GPU3 SYS SYS NV1 X 8-15
|
|
|
|
Legend:
|
|
X = Self
|
|
NV1 = Connected via NVLink with 1 hop
|
|
SYS = Connected via the PCIe or CPU subsystem""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_nvlink_partially_connected_output(
|
|
mock_run, nvlink_partially_connected_output, mock_cuda_is_available_true, mock_nvidia_device_properties
|
|
):
|
|
mock_run.return_value = nvlink_partially_connected_output
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call(
|
|
"Warning: Not all GPUs are fully connected via NVLink. Some GPUs are connected via slower interfaces. "
|
|
"It is recommended to switch to a different machine with faster GPU connections for optimal multi-GPU training performance."
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def nvlink_not_connected_output():
|
|
return mock.MagicMock(
|
|
stdout=""" GPU0 GPU1 GPU2 GPU3 CPU Affinity NUMA Affinity GPU NUMA ID
|
|
GPU0 X PHB PHB PHB 0-47 0 N/A
|
|
GPU1 PHB X PHB PHB 0-47 0 N/A
|
|
GPU2 PHB PHB X PHB 0-47 0 N/A
|
|
GPU3 PHB PHB PHB X 0-47 0 N/A
|
|
|
|
Legend:
|
|
|
|
X = Self
|
|
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
|
|
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
|
|
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
|
|
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
|
|
PIX = Connection traversing at most a single PCIe bridge
|
|
NV# = Connection traversing a bonded set of # NVLinks""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_nvlink_not_connected_output(
|
|
mock_run, nvlink_not_connected_output, mock_cuda_is_available_true, mock_nvidia_device_properties
|
|
):
|
|
mock_run.return_value = nvlink_not_connected_output
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call(
|
|
"Warning: Not all GPUs are fully connected via NVLink. Some GPUs are connected via slower interfaces. "
|
|
"It is recommended to switch to a different machine with faster GPU connections for optimal multi-GPU training performance."
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def nvlink_all_gpu_connected_but_other_connected_output():
|
|
return mock.MagicMock(
|
|
stdout=""" GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7 NIC0 NIC1 NIC2 NIC3 NIC4 NIC5 NIC6 NIC7 NIC8 NIC9 CPU Affinity NUMA Affinity GPU NUMA ID
|
|
GPU0 X NV12 NV12 NV12 NV12 NV12 NV12 NV12 SYS SYS PXB PXB SYS SYS SYS SYS SYS SYS 0-63,128-191 0 N/A
|
|
GPU1 NV12 X NV12 NV12 NV12 NV12 NV12 NV12 SYS SYS PXB PXB SYS SYS SYS SYS SYS SYS 0-63,128-191 0 N/A
|
|
GPU2 NV12 NV12 X NV12 NV12 NV12 NV12 NV12 PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS 0-63,128-191 0 N/A
|
|
GPU3 NV12 NV12 NV12 X NV12 NV12 NV12 NV12 PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS 0-63,128-191 0 N/A
|
|
GPU4 NV12 NV12 NV12 NV12 X NV12 NV12 NV12 SYS SYS SYS SYS SYS SYS SYS SYS PXB PXB 64-127,192-254 1 N/A
|
|
GPU5 NV12 NV12 NV12 NV12 NV12 X NV12 NV12 SYS SYS SYS SYS SYS SYS SYS SYS PXB PXB 64-127,192-254 1 N/A
|
|
GPU6 NV12 NV12 NV12 NV12 NV12 NV12 X NV12 SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS 64-127,192-254 1 N/A
|
|
GPU7 NV12 NV12 NV12 NV12 NV12 NV12 NV12 X SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS 64-127,192-254 1 N/A
|
|
NIC0 SYS SYS PXB PXB SYS SYS SYS SYS X PIX SYS SYS SYS SYS SYS SYS SYS SYS
|
|
NIC1 SYS SYS PXB PXB SYS SYS SYS SYS PIX X SYS SYS SYS SYS SYS SYS SYS SYS
|
|
NIC2 PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS X PXB SYS SYS SYS SYS SYS SYS
|
|
NIC3 PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS PXB X SYS SYS SYS SYS SYS SYS
|
|
NIC4 SYS SYS SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS X PXB SYS SYS SYS SYS
|
|
NIC5 SYS SYS SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS PXB X SYS SYS SYS SYS
|
|
NIC6 SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS X PIX SYS SYS
|
|
NIC7 SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS PIX X SYS SYS
|
|
NIC8 SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS X PXB
|
|
NIC9 SYS SYS SYS SYS PXB PXB SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS PXB X
|
|
|
|
Legend:
|
|
|
|
X = Self
|
|
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
|
|
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
|
|
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
|
|
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
|
|
PIX = Connection traversing at most a single PCIe bridge
|
|
NV# = Connection traversing a bonded set of # NVLinks
|
|
|
|
NIC Legend:
|
|
|
|
NIC0: mlx5_0
|
|
NIC1: mlx5_1
|
|
NIC2: mlx5_2
|
|
NIC3: mlx5_3
|
|
NIC4: mlx5_4
|
|
NIC5: mlx5_5
|
|
NIC6: mlx5_6
|
|
NIC7: mlx5_7
|
|
NIC8: mlx5_8
|
|
NIC9: mlx5_9
|
|
|
|
""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_nvlink_all_gpu_connected_but_other_connected_output(
|
|
mock_run,
|
|
nvlink_all_gpu_connected_but_other_connected_output,
|
|
mock_cuda_is_available_true,
|
|
mock_nvidia_device_properties,
|
|
):
|
|
mock_run.return_value = nvlink_all_gpu_connected_but_other_connected_output
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("All GPUs are fully connected via NVLink.")
|
|
|
|
|
|
@pytest.fixture
|
|
def nvidia_smi_nvlink_output_dual_gpu_no_numa():
|
|
return mock.MagicMock(
|
|
stdout="""
|
|
GPU0 GPU1 CPU Affinity NUMA Affinity GPU NUMA ID
|
|
GPU0 X NV1 0-15 0 N/A
|
|
GPU1 NV1 X 0-15 0 N/A
|
|
|
|
Legend:
|
|
|
|
X = Self
|
|
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
|
|
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
|
|
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
|
|
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
|
|
PIX = Connection traversing at most a single PCIe bridge
|
|
NV# = Connection traversing a bonded set of # NVLinks
|
|
""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_check_nvlink_connectivity__returns_fully_connected_when_nvidia_all_nvlink_two_gpus(
|
|
mock_run, nvidia_smi_nvlink_output_dual_gpu_no_numa, mock_cuda_is_available_true, mock_nvidia_device_properties
|
|
):
|
|
mock_run.return_value = nvidia_smi_nvlink_output_dual_gpu_no_numa
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("All GPUs are fully connected via NVLink.")
|
|
|
|
|
|
@pytest.fixture
|
|
def rocm_smi_xgmi_output_multi_gpu():
|
|
"""
|
|
rocm-smi --showtopotype on ROCm 6.0.3+
|
|
"""
|
|
return mock.MagicMock(
|
|
stdout="""
|
|
=============================== ROCm System Management Interface ============================
|
|
=============================== Link Type between two GPUs ===============================
|
|
GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7
|
|
GPU0 0 XGMI XGMI XGMI XGMI XGMI XGMI XGMI
|
|
GPU1 XGMI 0 XGMI XGMI XGMI XGMI XGMI XGMI
|
|
GPU2 XGMI XGMI 0 XGMI XGMI XGMI XGMI XGMI
|
|
GPU3 XGMI XGMI XGMI 0 XGMI XGMI XGMI XGMI
|
|
GPU4 XGMI XGMI XGMI XGMI 0 XGMI XGMI XGMI
|
|
GPU5 XGMI XGMI XGMI XGMI XGMI 0 XGMI XGMI
|
|
GPU6 XGMI XGMI XGMI XGMI XGMI XGMI 0 XGMI
|
|
GPU7 XGMI XGMI XGMI XGMI XGMI XGMI XGMI 0
|
|
================================== End of ROCm SMI Log ===================================
|
|
""",
|
|
returncode=0,
|
|
)
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_check_nvlink_connectivity__returns_fully_connected_when_amd_all_xgmi_8_gpus(
|
|
mock_run, rocm_smi_xgmi_output_multi_gpu, mock_cuda_is_available_true, mock_amd_device_properties
|
|
):
|
|
mock_run.return_value = rocm_smi_xgmi_output_multi_gpu
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("All GPUs are fully connected via XGMI.")
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_check_nvlink_connectivity__returns_no_gpus_when_no_gpus(mock_run, monkeypatch):
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("No GPUs available")
|
|
|
|
|
|
@mock.patch("subprocess.run")
|
|
def test_check_nvlink_connectivity__returns_unrecognized_vendor_when_unrecognized_vendor(
|
|
mock_run, monkeypatch, mock_cuda_is_available_true
|
|
):
|
|
mock_device_properties = mock.MagicMock(name="GPU Device", spec=["name"])
|
|
mock_device_properties.name = "GARAGE DIY HYPERSCALER GPU"
|
|
monkeypatch.setattr(torch.cuda, "get_device_properties", lambda idx: mock_device_properties)
|
|
with mock.patch("builtins.print") as mock_print:
|
|
check_nvlink_connectivity()
|
|
mock_print.assert_any_call("Unrecognized GPU vendor: GARAGE DIY HYPERSCALER GPU")
|
|
|
|
|
|
def test_fix_and_load_json():
|
|
# Test 1: Invalid JSON string with a trailing comma
|
|
invalid_json_trailing_comma = """
|
|
{
|
|
"_from_model_config": true,
|
|
"bos_token_id": 128000,
|
|
"eos_token_id": 128001,
|
|
"transformers_version": "4.45.0.dev0",
|
|
"do_sample": true,
|
|
"temperature": 0.6,
|
|
"top_p": 0.9,
|
|
}
|
|
"""
|
|
|
|
expected_output_trailing_comma = {
|
|
"_from_model_config": True,
|
|
"bos_token_id": 128000,
|
|
"eos_token_id": 128001,
|
|
"transformers_version": "4.45.0.dev0",
|
|
"do_sample": True,
|
|
"temperature": 0.6,
|
|
"top_p": 0.9,
|
|
}
|
|
|
|
result_trailing_comma = fix_and_load_json(invalid_json_trailing_comma)
|
|
assert result_trailing_comma == expected_output_trailing_comma
|
|
|
|
# Test 2: Invalid JSON string with missing commas between properties
|
|
invalid_json_missing_commas = """
|
|
{
|
|
"_from_model_config": true,
|
|
"bos_token_id": 128000,
|
|
"eos_token_id": 128001,
|
|
"transformers_version": "4.45.0.dev0"
|
|
"do_sample": true,
|
|
"temperature": 0.6,
|
|
"top_p": 0.9,
|
|
}
|
|
"""
|
|
|
|
expected_output_missing_commas = {
|
|
"_from_model_config": True,
|
|
"bos_token_id": 128000,
|
|
"eos_token_id": 128001,
|
|
"transformers_version": "4.45.0.dev0",
|
|
"do_sample": True,
|
|
"temperature": 0.6,
|
|
"top_p": 0.9,
|
|
}
|
|
|
|
result_missing_commas = fix_and_load_json(invalid_json_missing_commas)
|
|
assert result_missing_commas == expected_output_missing_commas
|
|
|
|
|
|
def test_select_sft_generate_example():
|
|
eval_mock = mock.MagicMock()
|
|
data_mock = mock.MagicMock()
|
|
|
|
test_dataset = {"data": [{"instruction": "Test instruction 1"}, {"instruction": "Test instruction 2"}]}
|
|
train_dataset = {"data": [{"instruction": "Train instruction 1"}, {"instruction": "Train instruction 2"}]}
|
|
|
|
data_mock.test_dataset.data = test_dataset["data"]
|
|
data_mock.train_dataset.data = train_dataset["data"]
|
|
|
|
# Test "first" instruction from test dataset
|
|
eval_mock.evaluate_example = "first"
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Test instruction 1"
|
|
|
|
# Test "first" instruction from train dataset when test dataset is empty
|
|
data_mock.test_dataset.data = []
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Train instruction 1"
|
|
|
|
# Test random selection from test dataset
|
|
eval_mock.evaluate_example = "random"
|
|
data_mock.test_dataset.data = [{"instruction": "Test instruction 1"}, {"instruction": "Test instruction 2"}]
|
|
with mock.patch("random.randint", return_value=1):
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Test instruction 2"
|
|
|
|
# Test random selection from train dataset when test dataset is empty
|
|
data_mock.test_dataset.data = []
|
|
with mock.patch("random.randint", return_value=1):
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Train instruction 2"
|
|
|
|
# Test specific index from test dataset
|
|
eval_mock.evaluate_example = 1
|
|
data_mock.test_dataset.data = [{"instruction": "Test instruction 1"}, {"instruction": "Test instruction 2"}]
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Test instruction 2"
|
|
|
|
# Test specific index from train dataset when test dataset has fewer elements
|
|
data_mock.test_dataset.data = [{"instruction": "Test instruction 1"}]
|
|
instruction = select_sft_generate_example(eval_mock, data_mock)
|
|
assert instruction == "Train instruction 2"
|
|
|
|
# Test out-of-range index
|
|
eval_mock.evaluate_example = 2
|
|
data_mock.test_dataset.data = [{"instruction": "Test instruction 1"}]
|
|
data_mock.train_dataset.data = [{"instruction": "Train instruction 1"}]
|
|
with pytest.raises(IndexError):
|
|
select_sft_generate_example(eval_mock, data_mock)
|
|
|
|
# Test unknown evaluation type
|
|
eval_mock.evaluate_example = "unknown"
|
|
with pytest.raises(ValueError):
|
|
select_sft_generate_example(eval_mock, data_mock)
|