Fix model memory issue #3758

Job	Run time
run-tests (latest, test_prod)	3m 10s
run-tests (latest, test_core)	4m 26s
run-tests (latest, test_cli)	5m 19s
run-tests (latest, test_big_modeling)	3m 54s
run-tests (latest, test_deepspeed)	4m 20s
run-tests (latest, test_fsdp)	4m 2s
run-tests (latest, test_example_differences)	5m 16s
run-tests (latest, test_checkpoint_step)	6m 20s
run-tests (latest, test_checkpoint_epoch)	5m 53s
run-tests (latest, test_rest)	15m 52s
run-tests (minimum, test_prod)	3m 16s
run-tests (minimum, test_core)	4m 37s
run-tests (minimum, test_cli)	5m 12s
run-tests (minimum, test_big_modeling)	4m 33s
run-tests (minimum, test_deepspeed)	4m 6s
run-tests (minimum, test_fsdp)	4m 13s
run-tests (minimum, test_example_differences)	4m 32s
run-tests (minimum, test_checkpoint_step)	5m 53s
run-tests (minimum, test_checkpoint_epoch)	8m 12s
run-tests (minimum, test_rest)	16m 48s
	1h 59m 54s

Provide feedback