add device_mesh support for FSDP #4918

Job	Run time
run-tests (latest, test_prod)	1m 55s
run-tests (latest, test_core)	4m 33s
run-tests (latest, test_cli)	3m 21s
run-tests (latest, test_big_modeling)	2m 13s
run-tests (latest, test_deepspeed)	3m 47s
run-tests (latest, test_fsdp)	2m 53s
run-tests (latest, test_example_differences)	2m 15s
run-tests (latest, test_checkpoint_step)	3m 44s
run-tests (latest, test_checkpoint_epoch)	3m 29s
run-tests (latest, test_rest)	8m 34s
run-tests (minimum, test_prod)	2m 30s
run-tests (minimum, test_core)	3m 37s
run-tests (minimum, test_cli)	3m 51s
run-tests (minimum, test_big_modeling)	3m 7s
run-tests (minimum, test_deepspeed)	2m 59s
run-tests (minimum, test_fsdp)	4m 7s
run-tests (minimum, test_example_differences)	2m 55s
run-tests (minimum, test_checkpoint_step)	4m 24s
run-tests (minimum, test_checkpoint_epoch)	4m 18s
run-tests (minimum, test_rest)	9m 26s
	1h 17m 58s

Provide feedback