Change max_steps=500 and checkpoint every 100 steps

Author: camiloCienet

axlearn/experiments/text/gpt/common.py

@@ -417,7 +417,8 @@ def adamw_decoupled_learner_config(
     peak_lr: float,
     max_step: int,
     weight_decay: float,
-    lr_warmup_steps: int = 2000,
+    # lr_warmup_steps: int = 2000,
+    lr_warmup_steps: int = 50,
     alpha: float = 0.1,
     b1: float = 0.9,
     b2: float = 0.95,
@@ -455,7 +456,8 @@ def adastar_learner_config(
     *,
     peak_lr: float,
     max_step: int,
-    lr_warmup_steps: int = 2000,
+    # lr_warmup_steps: int = 2000,
+    lr_warmup_steps: int = 50,
     alpha: float = 0.005,
     weight_decay: float = 3.16e-4,
     b1: float = 0.95,

axlearn/experiments/text/gpt/fuji.py

@@ -292,6 +292,7 @@ def get_trainer_kwargs(
     rope_theta = ROPE_THETA[version]
 
     trn2_config = _generate_trn2_custom_configs(model_size, version=version)
+    max_step=500
 
     # dict() is more readable here.
     # pylint: disable=use-dict-literal
@@ -403,6 +404,7 @@ def get_trainer_kwargs(
             max_sequence_length=max_sequence_length,
             train_batch_size=train_batch_size,
             max_step=max_step,
+            save_every_n_steps=100,
             mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8),
             mesh_rules=(
                 # Step time: