Multi-Dataset Validation (LM-Loss/Perplexity)

[bigximik]

✨ Description

Closes #65

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

📝 Changes

List the key changes introduced in this PR:

1. Added support for multipole validation datasets with loss and other existing metrics tracking. This introduces breaking change as now training.validation config field is a dictionary of validation dataset names and their application parameters.
2. Changed existing tests to accommodate new config format
3. Changed and extended documentation to describe new config format

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

• 📜 I have read and followed the contributing guidelines.
• 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
• 🎉 The functionality is complete, and I have tested the changes.
• 📝 I have updated the documentation if needed.
• ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• 🧩 I have commented my code, especially in hard-to-understand areas.

Testing

• 🧪 I have added or updated tests to cover my changes.
• ✔️ New and existing tests pass locally with my changes.
• 🚦 I have tested these changes on GPUs and verified training stability.
• 🏋️ I have tested the changes on realistic training workloads, if applicable.

[bigximik]

@jlamypoirier What do you think of such approach?

[jlamypoirier]

@bigximik That's really close to the approach I had in mind. I would suggest moving some of the changes from the data to the trainer though, it would simplify things and we'll want it anyway in the next step of #65.

[bigximik]

I have implemented it so that the config will look something like this:

data:
  datasets:
    Training: # training dataset, hardcoded name
      type: memmap
      path: some_path1
    Test:  # test dataset, hardcoded name
      type: memmap
      path: some_path2
    validation_dataset_name1:  # validation dataset, any name
      type: memmap
      path: some_path3
    validation_dataset_name2:  # validation dataset, any name
      type: memmap
      path: some_path4

training:
  training_iters: 2
  test_iters: 2
  validation:
    validation_dataset_name1:
      interval: 2
      iterations: 1
    validation_dataset_name2:
      interval: 2
      iterations: 1

@jlamypoirier, have I got it right?

[bigximik]

Maybe to add training_dataset_name and test_dataset_name into TrainingConfig with default values Training and Test so people can override it if they like?

[jlamypoirier]

I have implemented it so that the config will look something like this:
@jlamypoirier, have I got it right?

Yes

Maybe to add training_dataset_name and test_dataset_name into TrainingConfig with default values Training and Test so people can override it if they like?

I don't think it's worth it, but feel free to fix the capitalization if you find a way to do it without breaking backward comparibility

[bigximik]

in wandb metrics for different validation datasets look like this:

[bigximik]

I have changed validation to evaluations and ValidationConfig to EvaluationConfig, but kept the phase as Validation.
So, in the documentation, it is still referred to as validation datasets, which need to be defined in data.datasets, while their usage should be specified in training.evaluations.

[jlamypoirier]

Almost ready to merge, some minor issues and suggestions

[tscholak]

Guys can we get to a conclusion here please? I’d like us to merge this by tomorrow end of day, at the latest. Thanks

[jlamypoirier]

@bigximik Adjusted the code, ready to merge if it works for you.

[tscholak]

thank you both, looks good to me!

[bigximik]

I have updated the instruction fine-tuning documentation that came from the merge but created a separate issue for moving the check of used dataset definitions to _validate in TrainerConfig (#213).