Add support for drop_last=False

[odelalleau]

Is your feature request related to a problem? Please describe.

The current code generally doesn't support drop_last=False.

Describe the solution you'd like

Proper support for drop_last=False. This could be done in two phases:

1. Support it for validation only
2. Support it for training as well

[trias702]

I see that the current SFT code essentially has drop_last=False for validation. But in a DP setting, how can that work, you won't have equal shards on each DP?

[odelalleau]

I see that the current SFT code essentially has drop_last=False for validation. But in a DP setting, how can that work, you won't have equal shards on each DP?

Yeah, you may need to pad with dummy samples whose loss would be masked out.

[trias702]

I see that the current SFT code essentially has drop_last=False for validation. But in a DP setting, how can that work, you won't have equal shards on each DP?

Yeah, you may need to pad with dummy samples whose loss would be masked out.

That sounds like a fairly complex solution which could be cumbersome to maintain. I think the easier, and generally acceptable fix here is to just make all validation use drop_last=True in Aligner. We already do that for DPO and RM.

[odelalleau]

But there may be cases where people don't want to drop samples. For instance if you're reporting metrics over an "official" validation set on a common benchmark.

Edit: but I agree the main priority is to avoid crashing by default, so if our current SFT code actually uses drop_last=False by default and crashes because of that, we should instead make it use drop_last=True.