Ensure compatibility between models and datasets

[jlamypoirier]

✨ Description

When specifying a dataset and model in Fast-LLM, it's not currently possible to tell whether the two are compatible. This PR mitigates the issue by adding runtime checks. The dataset preparator saves the relevant preprocessing options with the dataset, and the memmap dataset compares it against the options required by the model, forwarded through the new preprocessing field of SamplingData. Additional benefits:

• Added flexibility in the dataset content, i.e. the dataset will produce trivial loss masking spans and/or images patches if missing instead of crashing, see Text-only training of multimodal models #403. Added a warning in case that's unintended.
• The dataset reads only what it needs, ex. a dataset with images won't read them for text-only models.

Should be backward compatible, the older dataset will just warn that they can't check compatibility. (Haven't tested though.)

So far this checks vocab size, use_loss_masking_spans, use_preference_spans, use_image_patches and patch shape. Missing the actual tokenizer, max image shape and image special tokens as it would require extra work and additional fields in the training config.

Also address multiple test failures from recent PRs

• (from Add stochastic mixer for supernet training #373) Fix checkpoint import/export for pipeline-parallel configurations with tied weights. The issue caused a failure in apriel 2 checkpoint conversion tests, and was previously unknown because we never tried conversion with tied weights.
• (From GDN mixer #392) Mark conversion as broken for hybrid_gdn since it's crashing. @oleksost
• (From activation-level disillation #388, GDN mixer #392) Fix skip_test for hybrid_gdn (was not excluding all tp configs), mistral_distill_activations (missing "df", "bf", "fp16")

And some maintenance.

• Cleanup some model configs (especially skip_tests), mark hybrid_mamba model config as unimportant
• Remove some outdated todos

That leaves 5 failing tests:

FAILED tests/models/test_checkpoint.py::test_huggingface_model[hybrid_mamba_2]@dependency_group_2 - AttributeError: 'NoneType' object has no attribute 'ssm_states'
FAILED tests/layers/test_lm_head.py::test_lm_head[config_dict12-distributed_config_dict12-True-1-auto] - AssertionError: Rms diff too big (4.942e-01 > 1.000e-05) between tensors 0.9908757209777832 and 0.4966764450073242
FAILED tests/layers/test_lm_head.py::test_lm_head[config_dict12-distributed_config_dict12-True-1-fused] - AssertionError: Rms diff too big (4.942e-01 > 1.000e-05) between tensors 0.9908757209777832 and 0.4966764450073242
FAILED tests/layers/test_lm_head.py::test_lm_head[config_dict12-distributed_config_dict12-True-1-triton] - AssertionError: Rms diff too big (4.942e-01 > 1.000e-05) between tensors 0.9908757209777832 and 0.4966764450073242
FAILED tests/layers/test_lm_head.py::test_lm_head[config_dict12-distributed_config_dict12-True-1-torch] - AssertionError: Rms diff too big (4.942e-01 > 1.000e-05) between tensors 0.9908757209777832 and 0.4966764450073242

The first one has been there for a long time, the remaining 4 concern reverse KL (#400) and are being investigated by @oleksost

[tscholak]

LGTM

[oleksost]

Looks good!
The 4 failing tests concerning reverse KL were fixed in #405.