Merge pull request #2863 from AI-Hypercomputer:chengnuojin-revert-sft

PiperOrigin-RevId: 847796543

Author: Google-ML-Automation

src/MaxText/rl/train_rl.py

@@ -94,7 +94,7 @@ def get_maxtext_model(config, devices=None):
   # load_parameters_path=/path/to/your/output/directory/0/items
   """
   model, mesh = model_creation_utils.create_nnx_model(config, devices=devices)
-  with jax.set_mesh(mesh):
+  with mesh:
     use_no_op_mappings = "maxtext_config" in config.vllm_additional_config
     tunix_model = TunixMaxTextAdapter(base_model=model, use_no_op_mappings=use_no_op_mappings)
     tunix_model.config = None

src/MaxText/sft/sft_trainer.py

@@ -160,7 +160,7 @@ def setup_trainer_state(mt_config, goodput_recorder=None):
 
 def train_model(mt_config, trainer, mesh):
   """Runs the SFT training loop in Tunix."""
-  with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules):
+  with mesh, nn_partitioning.axis_rules(mt_config.logical_axis_rules):
     trainer.train(trainer.data_hooks.train_data_iterator, trainer.data_hooks.eval_data_iterator)
   return trainer