Add max-prefill-length argument in distillation dataset generation script

[SurbhiJainUSC]

Description

This PR introduces max-prefill-length argument to the script that is used to generate dataset for distillation. This argument will be used to filter out prompt sequences that are larger than max-prefill-length before running inference.

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

• Tested with generate_distillation_dataset.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[hengtaoguo]

Thank you Surbhi!

Question out of curiosity: which model are you experimenting with this SFT, and what is roughly the maximum prefill/max length supported on the chip you are using?

Asking because of the multimodal work, even the shortest prompt Describe image <start_of_image> will require 272 tokens in the prefill length. Mentioning this so that the group can be aware of, thanks!

[SurbhiJainUSC]

Thank you Surbhi!

Question out of curiosity: which model are you experimenting with this SFT, and what is roughly the maximum prefill/max length supported on the chip you are using?

Asking because of the multimodal work, even the shortest prompt Describe image <start_of_image> will require 272 tokens in the prefill length. Mentioning this so that the group can be aware of, thanks!

I am using the base configurations for both: https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/configs/base.yml#L470

[shralex]

Thanks Surbhi. Added a suggestion