Migrate MaxText/train.py to maxtext/trainers/pre_train/train.py

Author: bvandermoon

tests/unit/pipeline_parallelism_test.py

@@ -26,6 +26,7 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from MaxText import pyconfig
+from maxtext.common.gcloud_stub import is_decoupled
 from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.globals import MAXTEXT_ASSETS_ROOT
 from maxtext.common.gcloud_stub import is_decoupled