ViperEkura
738cb8f128
- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks - Fix create_ref_model returning None on non-rank-0 under FSDP - Fix sync_old_model only updating old_model on rank-0 under FSDP - Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest - Add distributed tests for broadcast_state_dict and create_ref_model
94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
|
|
import pytest
|
|
import torch
|
|
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
|
|
|
|
from astrai.extension import KERNEL_NAMES, is_available
|
|
from astrai.model.transformer import AutoRegressiveLM
|
|
from astrai.tokenize import AutoTokenizer
|
|
from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
|
|
|
|
CUDA_AVAIL = torch.cuda.is_available()
|
|
KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES)
|
|
skip_no_cuda = pytest.mark.skipif(not CUDA_AVAIL, reason="CUDA not available")
|
|
skip_no_kernel = pytest.mark.skipif(not KERNEL_AVAIL, reason="CUDA kernels not built")
|
|
|
|
|
|
def pytest_configure(config):
|
|
config.addinivalue_line("markers", "slow: marks tests as slow")
|
|
config.addinivalue_line("markers", "integration: integration tests")
|
|
config.addinivalue_line("markers", "unit: fast unit tests")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def device():
|
|
"""Session-scoped device string (``"cuda"`` if available, else ``"cpu"``)."""
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
|
|
"""Create a simple tokenizer for testing purposes."""
|
|
tokenizer = Tokenizer(models.BPE())
|
|
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
|
|
trainer = trainers.BpeTrainer(
|
|
vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
|
|
)
|
|
tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
|
|
auto_tokenizer = AutoTokenizer()
|
|
auto_tokenizer._tokenizer = tokenizer
|
|
auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
|
|
return auto_tokenizer
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_tokenizer():
|
|
"""Session-scoped tokenizer, created once for the entire test run."""
|
|
return create_test_tokenizer()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_model(device):
|
|
"""Session-scoped small AutoRegressiveLM model, created once."""
|
|
config = make_tiny_config()
|
|
model = AutoRegressiveLM(config).to(device=device)
|
|
return {"model": model, "device": device, "config": config}
|
|
|
|
|
|
@pytest.fixture
|
|
def base_test_env(test_model, test_tokenizer):
|
|
"""Function-scoped test environment with isolated temp directory."""
|
|
test_dir = tempfile.mkdtemp()
|
|
config_path = os.path.join(test_dir, "config.json")
|
|
with open(config_path, "w") as f:
|
|
json.dump(TINY_CONFIG, f)
|
|
|
|
yield {
|
|
"device": test_model["device"],
|
|
"test_dir": str(test_dir),
|
|
"config_path": config_path,
|
|
"transformer_config": test_model["config"],
|
|
"model": test_model["model"],
|
|
"tokenizer": test_tokenizer,
|
|
}
|
|
|
|
shutil.rmtree(test_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def random_dataset():
|
|
return RandomTokenDataset(length=None)
|
|
|
|
|
|
@pytest.fixture
|
|
def multi_turn_dataset():
|
|
return RandomTokenDataset(length=None, with_loss_mask=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def early_stopping_dataset():
|
|
return RandomTokenDataset(length=10, stop_after=5)
|