works for fp8 with deepspeed #5431

Job	Run time
run-tests (latest, test_prod)	3m 27s
run-tests (latest, test_core)	2m 49s
run-tests (latest, test_cli)	2m 31s
run-tests (latest, test_big_modeling)	2m 31s
run-tests (latest, test_deepspeed)	2m 24s
run-tests (latest, test_fsdp)	2m 20s
run-tests (latest, test_example_differences)	2m 22s
run-tests (latest, test_checkpoint_step)	3m 53s
run-tests (latest, test_checkpoint_epoch)	3m 43s
run-tests (latest, test_rest)	9m 56s
run-tests (minimum, test_prod)	2m 44s
run-tests (minimum, test_core)	7m 18s
run-tests (minimum, test_cli)	3m 33s
run-tests (minimum, test_big_modeling)	4m 26s
run-tests (minimum, test_deepspeed)	3m 30s
run-tests (minimum, test_fsdp)	3m 26s
run-tests (minimum, test_example_differences)	3m 14s
run-tests (minimum, test_checkpoint_step)	5m 36s
run-tests (minimum, test_checkpoint_epoch)	4m 58s
run-tests (minimum, test_rest)	12m 56s
	1h 27m 37s

Provide feedback