Feature Request (Pre-run sanity checks via Trackio)

[reed-meyerson]

The Speculators training script supports Trackio logging.

Currently, there are many sanity checks that are useful to do before investing substantial resources into a large speculators training run. These include

• Verify responses are sane
• Verify the loss mask
• Verify the presence of thinking (or absence) when expected (or unexpected)
• When fine-tuning or continuing a prior run, verify the initial acceptance rates are reasonable

Trackio supports plain-text and markdown logging. It would be amazing to generate a pre-run sanity check report before training begins.

[fynnsu]

Yeah, I think we could do something to improve this, including a small "report" of sorts.
A few comments:

1. We do some of this already during the "preprocessing" stage, when we print out a sample with loss mask coloring. I get the value of also having this during training / adding it to the trackio logger but there is some added complexity in doing so.
2. The added complexity comes from the fact that the training process only has raw token ids and loss mask since it doesn't do any tokenization itself. We could probably instantiate a tokenizer object and detokenize the inputs to get the input text, but we definitely can't un-chat template it. So thinking tokens/loss mask verification/etc would all have to be just manually reviewing samples.
3. "When fine-tuning or continuing a prior run, verify the initial acceptance rates are reasonable", we can probably achieve this by just running a val epoch at the start when resuming training and pushing the results.
4. I think the other thing to consider is that we don't just support trackio, but also wandb and tensorboard. It would be good to do this in a way where all backends get the report (if they support it), but that shouldn't be too difficult to implement.