Support configuration of num_workers and max_samples_per_sequence in llava_next_pretrain

[bernardhan33]

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

The llava_next_pretrain has been the go-to pretraining script for my LLaVA model pretraining because of #11931 .

It would be great to support modifying the num_workers instead of the hardcoded 32 for Energon data loader. ref.

Additionally, we noticed that it would be important to pass in the max_samples_per_sequence parameter (ref). Otherwise, there are no limits to allocating shards to the workers and there seems to be an issue with Energon (that I will file a bug to Energon repo and share here) that there are duplicate shard ranges allocated to the same worker.

Describe the solution you'd like

Support configuration of num_workers and max_samples_per_sequence through flags.

Describe alternatives you've considered

Modify / Patch the source code directly.

Additional context

Creating this issue and I can follow up with a CL.

Related Energon issue: NVIDIA/Megatron-Energon#70