Sequential sample packing

[DreamGenX]

Description

This change adds support for next-fit bin packing in MultipackBatchSampler. What this means is that we can now use sample packing while preserving the order of the examples.

Motivation and Context

Current sample packing affects order of examples in ways that can affect training results. This PR adds a new flag:
--sample_pack_sequentially which uses a simple greedy / sequential next-fit bin packing.

If you use --sample_pack_sequentially alone, the order of the examples will be determined by the underlying RandomSampler. If you use --sample_pack_sequentially with --curriculum_sampling, the order of the examples will be the same as in the training data. Both options make sense and differ from the current settings.

For example, previously it was not possible to do proper curriculum learning since the bin packing algorithm would reorder the examples anyway.

How has this been tested?

I tested it by running the new config: examples/llama-3/lora-1b-sample-packing-sequentially.yml
On this dataset, packing efficiency is high (>97%).

---

I would also like to bring attention to a potential bug in the multipack code:

• Here we set group_size and bin_size
• But it's never used.

[winglian]

@DreamGenX can you make this PR editable my maintainers? I'm happy to help get the linting and tests passing for this

[DreamGenX]

@winglian done, thank you

[winglian]

would it make sense to have pydantic validation to warn if curriculum sampling is not enabled when sample_packing_sequentially is enabled?

[DreamGenX]

There's a LOG.warn in the dataset class:

        if self.sequential and not isinstance(sampler, SequentialSampler):
            LOG.warn(
                "using sequential sample packing with non-sequential sampler, did you want to also enable curriculum_sampling?"
            )

Though it aslo makes sense to use the sequential packing without curriculum sampling. That way you get the same order as if you did not do packing. Can be useful way to get the benefits of packing with minimal diff.

[winglian]

I rebased again against the latest main since there were still some conflicts and linting issues. Hopefully I didn't screw up the implementation.

[NanoCode012]

Is this double dataset intentional?

[DreamGenX]

that's how it was in the original config, it just doubles the data

[NanoCode012]

Oh, I think you were basing off a config meant for deduplication examples/llama-3/lora-1b-deduplicate-sft.yml. We won't need this for your case.

[NanoCode012]

Maybe forgot to remove this?

[DreamGenX]

yep, seems unnecessary -- it's a leftover from the config i forked

[NanoCode012]

Could we add both of these to documentation docs/config.qmd to explain them?

[NanoCode012]

Similar to the earlier discussion, could we move this validation to the schema src/axolotl/utils/schemas/config.py and not here?

[NanoCode012]

You can change the check to be: if sample_pack_sequential and not curriculum_sampling instead

[DreamGenX]

@winglian @NanoCode012 BTW, when i run train, I see both "using non-sequential sample packing" and "using sequential sample packing" in the logs (I think due to eval, which is strange), but during pre-process I see only using sequential sample packing as expected. I don't know why, because all occurences of the MultiPackBatchSampler get the sequential flag (even the instance used for eval), except for the one used for pre-train... but maybe I missed something.