Spaces:
Paused
Paused
| import json | |
| import socket | |
| import subprocess | |
| import sys | |
| from pathlib import Path | |
| import pytest | |
| import torch | |
| from mouh.training import PackedByteDataset, SFTDataset, TrainArgs, lr_at, train | |
| from mouh.verification.correction_store import CorrectionStore | |
| from mouh.jev.types import JevDiagnostic, JevErrorType, JevSeverity | |
| ROOT = Path(__file__).resolve().parents[1] | |
| def corpus(tmp_path): | |
| f = tmp_path / "c.txt" | |
| f.write_text(("def add(a, b):\n return a + b\n\n" * 40)) | |
| return str(f) | |
| def test_packed_dataset_blocks_are_disjoint_between_ranks(tmp_path): | |
| texts = ["x" * 500] # un seul document : le partage par rang doit quand même fonctionner | |
| r0 = list(PackedByteDataset(texts, 64, rank=0, world_size=2, loop=False)) | |
| r1 = list(PackedByteDataset(texts, 64, rank=1, world_size=2, loop=False)) | |
| assert len(r0) > 0 and len(r1) > 0 and abs(len(r0) - len(r1)) <= 1 | |
| assert all(b.shape == (64,) for b in r0 + r1) | |
| def test_lr_schedule(): | |
| a = TrainArgs(lr=1.0, warmup_steps=10, max_steps=110, min_lr_ratio=0.1) | |
| assert lr_at(0, a) == pytest.approx(0.1) and lr_at(9, a) == pytest.approx(1.0) | |
| assert lr_at(110, a) == pytest.approx(0.1, abs=1e-6) | |
| assert lr_at(60, a) < lr_at(20, a) | |
| def test_single_process_training_reduces_loss_and_saves_loadable_checkpoint(tmp_path): | |
| a = TrainArgs(preset="mouh-tiny", data=corpus(tmp_path), seq_len=96, micro_batch_size=4, grad_accum=1, | |
| max_steps=25, warmup_steps=3, lr=3e-3, log_every=100, save_every=0, | |
| output_dir=str(tmp_path / "out"), device="cpu") | |
| res = train(a) | |
| assert res["last_loss"] < res["first_loss"] * 0.7 | |
| from mouh.model import MouhForCausalLM | |
| m = MouhForCausalLM.from_pretrained(tmp_path / "out" / "final") | |
| assert m(torch.randint(0, 256, (1, 8))).logits.shape[-1] == 512 | |
| assert json.loads((tmp_path / "out" / "final" / "trainer_state.json").read_text())["step"] == 25 | |
| def test_sft_from_correction_store(tmp_path): | |
| store = CorrectionStore(tmp_path / "c.jsonl") | |
| d = JevDiagnostic(severity=JevSeverity.ERROR, error_type=JevErrorType.TEST_FAILURE, message="assert 1 == 2") | |
| store.add(task="add", failing_files={"a.py": "x"}, diagnostics=(d,), fixed_files={"a.py": "y"}, attempts=2) | |
| store.export_sft(tmp_path / "sft.jsonl") | |
| ds = SFTDataset(str(tmp_path / "sft.jsonl"), seq_len=512) | |
| ids, labels = next(iter(ds)) | |
| assert ids.shape == labels.shape == (512,) | |
| assert (labels == -100).sum() > 0 and (labels != -100).sum() > 0 # perte sur la complétion seulement | |
| def _free_port(): | |
| with socket.socket() as s: | |
| s.bind(("127.0.0.1", 0)) | |
| return s.getsockname()[1] | |
| def test_ddp_two_processes_cpu(tmp_path): | |
| cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", f"--master_port={_free_port()}", | |
| str(ROOT / "scripts" / "train.py"), "--preset", "mouh-tiny", "--data", corpus(tmp_path), "--seq_len", "96", | |
| "--micro_batch_size", "2", "--grad_accum", "2", "--max_steps", "6", "--warmup_steps", "2", "--lr", "3e-3", | |
| "--log_every", "3", "--save_every", "0", "--parallel", "ddp", "--device", "cpu", | |
| "--output_dir", str(tmp_path / "ddp")] | |
| p = subprocess.run(cmd, capture_output=True, text=True, timeout=300) | |
| assert p.returncode == 0, p.stderr[-2000:] | |
| assert (tmp_path / "ddp" / "final" / "model.safetensors").exists() | |
| def test_text_source_is_reiterable_so_epochs_loop(tmp_path): | |
| from mouh.training import TextSource | |
| ds = PackedByteDataset(TextSource(corpus(tmp_path)), 96, loop=True) | |
| it = iter(ds) | |
| blocks = [next(it) for _ in range(60)] # bien plus qu'une époque (~15 blocs) | |
| assert len(blocks) == 60 | |