feat: support tensor parallel & Data loader #5389

Job	Run time
run-tests (latest, test_prod)	1m 37s
run-tests (latest, test_core)	2m 52s
run-tests (latest, test_cli)	2m 36s
run-tests (latest, test_big_modeling)	2m 41s
run-tests (latest, test_deepspeed)	2m 19s
run-tests (latest, test_fsdp)	2m 43s
run-tests (latest, test_example_differences)	2m 46s
run-tests (latest, test_checkpoint_step)	3m 59s
run-tests (latest, test_checkpoint_epoch)	3m 39s
run-tests (latest, test_rest)	11m 22s
run-tests (minimum, test_prod)	2m 25s
run-tests (minimum, test_core)	3m 25s
run-tests (minimum, test_cli)	3m 26s
run-tests (minimum, test_big_modeling)	3m 35s
run-tests (minimum, test_deepspeed)	3m 31s
run-tests (minimum, test_fsdp)	4m 14s
run-tests (minimum, test_example_differences)	3m 27s
run-tests (minimum, test_checkpoint_step)	3m 25s
run-tests (minimum, test_checkpoint_epoch)	3m 17s
run-tests (minimum, test_rest)	3m 31s
	1h 10m 50s

Provide feedback