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