Repo: github.com/AmmarHassona/trainsafe

I was working on fine-tuning an open-source small language model (SLM) on Arabic using DPO. I had the data, the pipeline, and everything set up for training. I was fairly confident that this training run would improve the model and align it further to what I wanted. I started the training and let it run until it finished. When I came back to test the checkpoint, it was speaking Chinese.

Loss only tells you the model is learning something — not what it's actually learning. By the time training finished, I had wasted my time and my compute with nothing useful to show for it. If only there was something to tell me if training was actually going well before it was too late.

This is when I began looking at tools that could help me solve this issue. Nothing existed that did exactly what I needed, so I built it myself. I built trainsafe to plug into any HuggingFace or TRL training pipeline with two lines of code. It runs alongside your training and checks whether the model's outputs are still behaving correctly at every eval checkpoint — catching issues like language drift, output collapse, and repetition loops before the run finishes.

Getting Started

Install with pip:

pip install trainsafe

# with language drift detection
pip install "trainsafe[language]"

Then add it to your training script with no other changes needed:

from trainsafe import TrainSafeCallback

trainer = SFTTrainer(
    model=model,
    ...
    callbacks=[TrainSafeCallback()]
)
trainer.train()

What it checks

At every eval checkpoint, trainsafe generates a small sample of outputs and runs five checks automatically:

Language — detects if the model switches output language mid-training. This is exactly what would have caught my situation.

Length — catches output collapse (model suddenly generating much shorter text) or runaway growth. Compares against a rolling baseline so legitimate learning doesn't trigger false alarms.

Repetition — flags n-gram loops inside individual outputs, the classic "the the the the" failure mode.

Echo — flags outputs that are mostly a copy of the prompt rather than an actual response.

Format — detects if a model trained to output JSON starts producing plain text, or vice versa.

All five run with zero configuration. If the overall health score drops below the warning threshold, you get a warning. If it drops below the stop threshold, training stops and trainsafe points you at the last healthy checkpoint.

What it looks like

Healthy run

[TrainSafe @ step 5] ✅ Language consistent (en)
[TrainSafe @ step 5] ✅ Output length normal (avg 62 words)
[TrainSafe @ step 5] ✅ No repetition detected
[TrainSafe @ step 5] ✅ No prompt echoing
[TrainSafe @ step 5] ✅ Format consistent (plain)
[TrainSafe @ step 5] Overall health: 1.00

When something goes wrong

[TrainSafe @ step 600] 🚨 Language drift — expected ar, got zh
[TrainSafe @ step 600] 🚨 Output length collapsed (avg 3 words vs baseline 87)
[TrainSafe @ step 600] ⚠️  Repetition detected in 3/5 outputs
[TrainSafe @ step 600] Overall health: 0.20
>>> TrainSafe stopped training. Recommended checkpoint: step 400.

Custom probes

If you have a specific capability you can't afford to lose, you can define fixed prompts and expected behaviors in a YAML file:

probes:
  - prompt: "مرحبا، كيف يمكنني مساعدتك؟"
    checks:
      - language: ar
      - min_length: 10
      - not_contains: ["<|im_start|>", "###"]

These run at every checkpoint alongside the automatic checks.

trainsafe is MIT licensed, early stage, and feedback is very welcome. If you've hit a similar problem during fine-tuning I'd love to hear about it in the comments.

Repo: github.com/AmmarHassona/trainsafe