add use_all_gather for option #5144

Job	Run time
run-tests (latest, test_prod)	1m 23s
run-tests (latest, test_core)	2m 17s
run-tests (latest, test_cli)	2m 24s
run-tests (latest, test_big_modeling)	7m 11s
run-tests (latest, test_deepspeed)	7m 57s
run-tests (latest, test_fsdp)	2m 22s
run-tests (latest, test_example_differences)	6m 45s
run-tests (latest, test_checkpoint_step)	2m 23s
run-tests (latest, test_checkpoint_epoch)	2m 21s
run-tests (latest, test_rest)	2m 27s
run-tests (minimum, test_prod)	2m 14s
run-tests (minimum, test_core)	3m 9s
run-tests (minimum, test_cli)	2m 59s
run-tests (minimum, test_big_modeling)	3m 4s
run-tests (minimum, test_deepspeed)	10m 11s
run-tests (minimum, test_fsdp)	3m 0s
run-tests (minimum, test_example_differences)	2m 59s
run-tests (minimum, test_checkpoint_step)	8m 12s
run-tests (minimum, test_checkpoint_epoch)	2m 56s
run-tests (minimum, test_rest)	3m 1s
	1h 19m 15s

Provide feedback