chore: revert accidental changes to optimizer.py, state_builder_test.py, trainer.py and fuji-70B txt during rebase

Author: judyzhaoxinwu

axlearn/common/optimizers.py

@@ -37,6 +37,8 @@
 import typing_extensions
 from absl import logging
 from jax import numpy as jnp
+
+# from jax._src.sharding_impls import TransferToMemoryKind
 from optax._src import numerics
 
 from axlearn.common import flax_struct, schedule

axlearn/common/state_builder_test.py

@@ -293,7 +293,6 @@ def _mock_bert_trainer_config(
                 embedding_cfg=bert.bert_embedding_config(max_position_embeddings=max_len),
                 stack_cfg=bert.bert_transformer_config(num_layers=3, num_heads=2),
             ),
-            mesh_axis_names=("fsdp", "model"),
         )
 
     def _mock_bert_trainer_config_and_state(

axlearn/common/trainer.py

@@ -646,7 +646,7 @@ def run(
                         )
                         self.vlog(3, "Done step %s", self.step)
                         num_steps += 1
-                        if num_steps % 1 == 0:
+                        if num_steps % 100 == 0:
                             now = time.perf_counter()
                             average_step_time = (now - start_time) / num_steps
                             self._step_log("Average step time: %s seconds", average_step_time)
@@ -1372,7 +1372,6 @@ def _maybe_stop_or_start_tracing(
                 cfg.start_trace_process_indices == "all"
                 or jax.process_index() in cfg.start_trace_process_indices
             )
-            self._step_log(f"{should_start_tracing=}")
         if should_start_tracing:
             self._step_log("Start profiler tracing")
             profiler_options = jax.profiler.ProfileOptions()

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-70B-v1-flash-fp8.txt

@@ -248,7 +248,7 @@ mesh_rules[7][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[7][1].config_modifiers[1].fp8_amax_history_length: 128
 mesh_rules[7][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[7][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
-mesh_rules[8][0]: 'gpu-(a4-highgpu-8g|a4x-highgpu-4g)-(256|512|1024)'
+mesh_rules[8][0]: 'gpu-(a4-highgpu-8g)-(256|512|1024)'
 mesh_rules[8][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'
 mesh_rules[8][1].config_modifiers[0].mesh_shape[0]: 1
 mesh_rules[8][1].config_modifiers[0].mesh_shape[1]: -1
@@ -259,15 +259,8 @@ mesh_rules[8][1].config_modifiers[0].mesh_shape[5]: 1
 mesh_rules[8][1].config_modifiers[1].gpu_block_size: 64
 mesh_rules[8][1].config_modifiers[1].klass: 'axlearn.common.flash_attention.layer.FlashBlockSizeModifier'
 mesh_rules[8][1].config_modifiers[1].tpu_block_size: 512
-mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: '.*input'
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*[kqv]_proj'
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
-mesh_rules[8][1].config_modifiers[2].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
-mesh_rules[8][1].config_modifiers[3].fp8_amax_history_length: 128
-mesh_rules[8][1].config_modifiers[3].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
+mesh_rules[8][1].config_modifiers[2].fp8_amax_history_length: 128
+mesh_rules[8][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.FP8ConfigModifier'
 mesh_rules[8][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
 mesh_rules[9][0]: 'neuron-(trn2|trn2n).48xlarge-64'
 mesh_rules[9][1].config_modifiers[0].klass: 'axlearn.common.trainer_config_modifier.MeshShapeModifier'