WIP: Initial FSDP2 support #5445

Job	Run time
run-tests (latest, test_prod)	2m 4s
run-tests (latest, test_core)	2m 57s
run-tests (latest, test_cli)	2m 46s
run-tests (latest, test_big_modeling)	2m 25s
run-tests (latest, test_deepspeed)	2m 16s
run-tests (latest, test_fsdp)	2m 23s
run-tests (latest, test_example_differences)	2m 20s
run-tests (latest, test_checkpoint_step)	3m 48s
run-tests (latest, test_checkpoint_epoch)	3m 38s
run-tests (latest, test_rest)	9m 46s
run-tests (minimum, test_prod)	2m 44s
run-tests (minimum, test_core)	3m 48s
run-tests (minimum, test_cli)	3m 47s
run-tests (minimum, test_big_modeling)	3m 30s
run-tests (minimum, test_deepspeed)	3m 25s
run-tests (minimum, test_fsdp)	3m 31s
run-tests (minimum, test_example_differences)	3m 17s
run-tests (minimum, test_checkpoint_step)	4m 53s
run-tests (minimum, test_checkpoint_epoch)	4m 44s
run-tests (minimum, test_rest)	10m 15s
	1h 18m 17s

Provide feedback