Migrate Gpt3 to NNX. #2537

Job	Run time
Build and upload maxtext package / Build and upload maxtext package (v4-8)	34s
maxtext_tpu_integration_tests (py312) / run	13m 1s
maxtext_gpu_integration_tests (py312, cuda12) / run	23m 49s
maxtext_gpu_unit_tests (py312, cuda12) / run	14m 4s
maxtext_tpu_unit_tests (py312) / run	18m 19s
maxtext_cpu_unit_tests (py312, 1) / run	3m 39s
maxtext_cpu_unit_tests (py312, 3) / run	8m 40s
maxtext_cpu_unit_tests (py312, 2) / run	10m 11s
maxtext_cpu_unit_tests (py312, 4) / run	5m 58s
Notify failed build	5s
	1h 38m 20s

Provide feedback