File size: 1,586 Bytes
014b14c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """Smoke #6: the shared train.py wires up the REAL trl 1.7.1 SFT API.
Trains a tiny random Qwen2 for 1 step on a handful of dummy-teacher rows, no
4-bit, on CPU. Proves SFTConfig(max_length, completion_only_loss) +
SFTTrainer(processing_class, peft_config) + conversational prompt/completion
data all fit together. Requires network (downloads the tiny test model).
"""
import argparse
from pathlib import Path
import pytest
from mathcompose.common.io import write_jsonl
from mathcompose.datagen.prm800k_loader import iter_prm800k
from mathcompose.datagen.gen_verifier_data import generate_verifier_dataset
from mathcompose.teachers import get_teacher
TINY = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
FIX = Path(__file__).parent / "fixtures" / "prm800k_sample.jsonl"
@pytest.mark.slow
@pytest.mark.network
def test_train_one_step(tmp_path):
pytest.importorskip("torch")
from mathcompose.train.train import build_and_train
rows = list(generate_verifier_dataset(iter_prm800k(FIX), get_teacher("dummy"), banned=set()))
train_file = tmp_path / "train.jsonl"
write_jsonl(rows, train_file)
out_dir = tmp_path / "vsmoke"
args = argparse.Namespace(
task="v", config="configs/verifier_v.yaml", base_id=TINY,
train_file=str(train_file), val_file=None, output_dir=str(out_dir),
no_4bit=True, max_steps=1, wandb=False, push=False, hub_model_id=None,
)
build_and_train(args)
# a LoRA adapter should have been saved
assert (out_dir / "adapter_config.json").exists() or (out_dir / "adapter_model.safetensors").exists()
|