Feat/val split

[TonyChen06]

9826848 — Add held-out validation split for training
Introduces --val_split (train phase only): a value <1 holds out that fraction of the training set, >=1 holds out that many examples. The split is a seeded shuffle of the combined dataset, so every DDP rank holds out the identical set. When set, per-epoch validation loss drives best-checkpoint selection and is logged to wandb as val/loss; on multi-GPU it's reduced across ranks (new all_reduce_sum helper) to cover the full val set. Skipped for RL.

The reduction across ranks is necessary. Without it, val/loss would reflect only rank 0's shard, which wasn't too much of a problem before because we were just logging using train loss for vibes but now were probably stopping runs with val loss so we should make it precise and not waste other ranks' data.

13298d6 — Disable augmentation for the validation split
The held-out set was built with the same train-mode args, so --augment_ecg / --augment_rgb randomly augmented val inputs every epoch, making val/loss non-deterministic (it should be deterministic). The val dataset is now built through a copied args namespace with augmentation forced off (build_data_representation gains an optional args override, defaulting to self.args so the train path is unchanged). May not actually be used all too often, but just letting you know this problem exists in case we want to use augment signal ever. No need to merge if not needed.

d136869 — Tighten early-stopping delta default to 0.0001
Lowers --patience_delta from 0.1 to 0.0001 so early stopping reacts to smaller improvements; --patience stays at 5. This is to prepare for next commit, where 1e-4 delta is standard.

959bf — Gate early stopping on an active validation split
Early stopping now triggers only when a held-out validation loss exists (val_loss is not None); otherwise it would stop on train loss, which we I think we should fully discontinue because early stopping without validation is sort of crazy. Early stopping is now no-op when --val_split is unset or in RL, where val_loss is None. Not necessary to merge but I think it should probably work this way.

[willxxy]

Overall Comments:

I think this is much more complicated than necessary. Meaning, it changes a lot of the fundamental logic, for example i dont like how our DatasetMixer class now returns train or val loader. I want it to be simply a class to return a single dataset, as constructing a single dataset can be quite complicated in it of itself.

I think we should do the splitting in our BuildDataloader class. The main_trainer.py logic seems okay, but should be cleaner.

I also don't understand why the validation runner needs to be so different from the training runner. Shouldn't it just be the same but just with the validation data?

ChatGPT suggested something like this for splitting

from torch.utils.data import DataLoader, random_split, DistributedSampler


class BuildDataLoader:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.dataset_mixer = DatasetMixer(self.args)

    def build_dataloaders(self):
        torch_dataset = self.dataset_mixer.build_torch_dataset()

       if self.args.val_split:
          val_size = int(len(torch_dataset) * self.args.val_split)
          train_size = len(torch_dataset) - val_size
          train_dataset, val_dataset = random_split(
              torch_dataset,
              [train_size, val_size],
              generator=torch.Generator().manual_seed(self.args.seed),
          )
          train_loader = self.build_torch_dataloader(train_dataset, is_train=True)
          val_loader = self.build_torch_dataloader(val_dataset, is_train=False)
          return train_loader, val_loader
      else:
          return torch_dataset

And handling the tuple vs single instance return outside in the main_trainer (isinstance check if output is tuple).

[TonyChen06]

DatasetMixer returns a single dataset again; split moved to BuildDataLoader via seeded random_split.
run_validation folded into run_train (no optimizer → no-grad eval pass).
Simplified the main_trainer loop; RL excluded from the split.

Kept a val_dataloader attribute instead of a (train, val) tuple + isinstance because it avoids the isinstance branch and leaves main_evaluator untouched. Can switch if you prefer.

Shuffling disabling is necessary for determinism, especially for uneven batches. Augment disabling is needed for similar reasons, though I think we run into that case less.

[willxxy]

I think we can do a cleaner implementation.

Lets try this.
It should be possible to get the hf dataset from load_dataset then split that into train and validation splits if args.val_split is defined. Then we can just use the same dataloader creation pipeline we already have. The reason for this is that I don't like how build_dataloader is looking. the self.dataset and self.val_dataloader are kind of weird. It would be better to just use the functions we already have in which we can 1. grab a hf dataset, 2) create a torch dataset, and 3) create a torch dataloader. We should try to use these functionalities that exist. If these functionalities are just fundamentally bad, then we should just rethink how we create the dataloder from a hf dataset.

as for the validation runner, i want it to be a separate file called validator.py, then just copy paste the train runner function and modify it as needed (no optimizer, no backward). we already get the average loss in the train runner function so in the validator runner we can just log that with the all reduce logic as well.

For determining augment, we should put it in the init.py of base.py dataset, where we just check if its train or not then augment based on that. we can also put the indication of train or not in the class args as well if thats easier. let me know if your confused on anything.

Also in the future, make sure you either answer or address via code all comments I make. This makes it helpful for me to align what has been addressed and what has not.

[TonyChen06]

Ok I gave it another shot; I think I handled all your comments

The one thing is that we split in DatasetMixer again due to "get the hf dataset from load_dataset then split that into train and validation splits" and contrary to "dont like how our DatasetMixer class now returns train or val loader. I want it to be simply a class to return a single dataset" from the first comment.

There's surely ways to make it so both are simultaneously true but I feel like that is more complicated. Let me know if I did something wrong here.