Upgrade Lightning from 2.5.1 → latest

[aditya0by0]

Title:

Description:
Currently the project is pinned to PyTorch Lightning 2.5.1 in f74c8ac because our LightningCLI-based integration fails on 2.5.2 and newer.

Problem

• In Lightning ≥2.5.2, LightningCLI and its underlying jsonargparse validation became stricter.

• Fields that were previously injected dynamically through parser.link_arguments (e.g. num_labels) now must either exist in the config file or be linked at an earlier stage.

• This causes errors like:

  Does not validate against any of the Union subtypes
  Subtypes: [<class 'NoneType'>, <class 'lightning.pytorch.core.module.LightningModule'>]
  Problem with given class_path ...
  Expected key train_metrics.init_args.metrics.macro-f1.init_args.num_labels to be in data
  

• In 2.5.1 this worked fine, but 2.5.2+ fails during config parsing.

[sfluegel05]

not directly related, but also concerning the lightning version:

On feature/new-ensemble-models, I added a custom Fit loop that does not run with lightning 2.1, but does run with 2.5.1 (I haven't tried other versions).