Padding-free SFT

[garrett361]

This PR adds the ability to perform padding-free SFT, which reduces the memory costs and increases throughput when per_device_train_batch_size>1. The model itself must support padding-free training to properly use this feature (Llama and bamba models support this, instance; see huggingface/transformers#35861 for some typical throughput improvements).

Pass --padding-free True to use.

[vwxyzjn]

Is this packing?

[garrett361]

Yeah, think it's the same @vwxyzjn .

Basically rather than padding out uneven examples which are batched together, you just concatenate them and add additional info to say where example boundaries are.

[garrett361]

Same idea as here

[garrett361]

FYI I'm double checking a possible issue with this at the moment; will report back after I have more info.

[garrett361]

I'm double checking a possible issue

Okay, I was debugging some bad-looking padding-free SFT training curves, but it turned out to be an issue with the mamba kernels which was recently fixed here. So, not an issue on the part of this commit.

[garrett361]

Some plots of padding-free training. All curves use 8xA100s and have --per_device_train_batch_size 2. The only difference between any pair of curves on a plot is whether --padding_free is true or false.

First, tuluv3 training on a BambaForCausalLM model. Loss curves are nearly identical (they will never be precisely the same due to numerics), and the throughput of the --padding_free true run is ~40% higher.

And then the same experiment with LlamaForCausalLM. Similar results: loss curves are similar and --padding_free true througput is ~40% higher.

CC @vwxyzjn @hamishivi

[hamishivi]

Thanks! So it seems like if the concatted sequence goes over the maximum supported length of the model, there might be issues? I guess in the ppo/grpo packing implementations, our code instead tries to fit samples into a given max length, and does use multiple microbatches if it has to.

Generally happy to merge this in since its optional, just also might be nice to name it packing since I think that's a more common term that people will understand.

[hamishivi]

minor nit: could we call this packing instead of padding_free? To make it clearer to end-users what the feature is.

Suggested change

	padding_free: bool = field(
	packing: bool = field(

[garrett361]

yep, that is fine

[garrett361]

done

[hamishivi]

Suggested change

	if args.padding_free:
	if args.packing:

[garrett361]

done

[garrett361]

So it seems like if the concatted sequence goes over the maximum supported length of the model, there might be issues?

Maybe? I'm not sure what the max supported length affects, apart from RoPE embeddings? Or I guess similar mechanisms like alibi or just fixed positional embeddings. For the models I've used (llama, bamba) padding-free/packing with, there's no problem because the explicitly sequence dependent bits adjust accordingly. Certainly this only works for some model classes, which is why the collator raises a warning.

I guess in the ppo/grpo packing implementations, our code instead tries to fit samples into a given max length, and does use multiple microbatches if it has to.

I'm only very familiar with the SFT finetune.py script at the moment, so can't yet speak to how you might use this for ppo/grpo packing.