Training Improvements: MultipackV2, Statistics, Mock Data

[RobotSail]

Adds a number of enhancements to provide better training performance, have clarity on raining times, and experiment more robustly.

Multipack V2

Multipack V2 has been tested as a batch sampler and found to improve training throughput by an amount that benefits training long-context models. However; it does not support models non padding-free models, therefore these would need to use Multipack V1.

Experimental setup

Constants:

GPUs: 8xA100
MBL: 52k
Distributed Backend: FSDP
MSL: 50k
Liger: on

Independent variables:

• Time per epoch: The bias-corrected moving average of the elapsed time per-step multiplied by the total number of batches with $\alpha = 0.9$
• Throughput: Zero-corrected mving average of the throughput with $\alpha = 0.9$
• Total tokens: Load the processed dataset and sum up all lens
• Average token len: Mean of lens

Distributed Sampler	Dataset	Num samples	Time per epoch	Throughput	Total tokens	Average Token Len
Multipack V1	1.4 skills-v1 (baseline short-context)	346432	4.5 hours	21.0	406470537	1173.31
Multipack V2	1.4 skills-v1 (baseline short-context)	346432	4.5 hours	21.0	406470537	1173.31
Multipack V1	1.5 skills-v2	374821	14.2 hours	7.495938060254868	669853984	1801.76
Multipack V2	1.5 skills-v2	374821	13.1 hours	8.44455839070264	669853984	1801.76
Multipack V1	1.5 skill-v2-reduced	163561	5.8 hours	7.8696600699200845	295960957	1809.483660530322
Multipack V2	1.5 skill-v2-reduced	163561	5 hours	9.556907986534208	295960957	1809.483660530322

Todos:

• Clean up API
• Add MIT license & attribution for MultipackDistributedSamplerV2
• Add logic for checking that model will be using non-padding training with MultipackDistributedSamplerV2

[cdoern]

I know Fynn has been working on "SDK-ifying" the sampler specifically @RobotSail , maybe we should sync on this with the training team

[github-actions]

This pull request has been automatically marked as stale because it has not had activity within 90 days. It will be automatically closed if no further activity occurs within 30 days.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. @RobotSail please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork