[Bug] Pretraining on sample-data throws StopIteration in initialization

Nobody has claimed this yet.

Hi, I followed your instructions precisely and set up the environment and data but because of the dataset sizes in the sample dataset an error occurs.
the arvix data set (and others) is small (8 bins after pre-processing), and therefore when it is sharded to 4, each one is of size 2 which is smaller then the _n_chunks is set to 4 therefore the linked line throws an error : https://github.com/Lightning-AI/lit-gpt/blob/c81f7a9ff379a0453c83a4fd79a764698ed1fea9/lit_gpt/packed_dataset.py#L170

for my use case I'm using the sample dataset only to test that the system is set-up correctly and therefore I implemented a hacky-fix:
https://github.com/Lightning-AI/lit-gpt/blob/c81f7a9ff379a0453c83a4fd79a764698ed1fea9/lit_gpt/packed_dataset.py#L48
remove the indexing from this line, meaning replace it with:
`python filenames = self._filenames`

Start in lit_gpt/packed_dataset.py at the linked lines around filename selection and chunk initialization, then reproduce pretraining with the sample dataset described in the issue. Confirm that initialization handles shards smaller than _n_chunks without StopIteration, and verify that sample-data pretraining completes initialization successfully.

Written by the indexing model from the issue text.