Ref: Changes to make it work Python 0.5.3. Topology v5p-128. Model 7b-fuji.

Author: camiloCienet

axlearn/common/trainer.py

@@ -1116,12 +1116,11 @@ def _run_step(
             self._trainer_state, outputs = compiled_train_step_fn(self.trainer_state, input_batch)
 
         n = self._config.log_every_n_steps or 100
-        if self.step % n == 0 or 0 <= self.step <= 5:
-            self._step_log(
-                "loss=%s aux=%s",
-                outputs["loss"],
-                jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]),
-            )
+        self._step_log(
+            "loss=%s aux=%s",
+            outputs["loss"],
+            jax.tree.map(lambda x: x.item() if x.ndim == 0 else f"T{x.shape}", outputs["aux"]),
+        )
 
         self.summary_writer(self.step, {"loss": outputs["loss"], **outputs["summaries"]})
         # Aggregate summaries across evalers.

axlearn/experiments/text/gpt/common.py

@@ -413,7 +413,7 @@ def adamw_decoupled_learner_config(
     peak_lr: float,
     max_step: int,
     weight_decay: float,
-    lr_warmup_steps: int = 2000,
+    lr_warmup_steps: int = 50,
     alpha: float = 0.1,
     b1: float = 0.9,
     b2: float = 0.95,
@@ -451,7 +451,7 @@ def adastar_learner_config(
     *,
     peak_lr: float,
     max_step: int,
-    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

@@ -291,10 +291,12 @@ def get_trainer_kwargs(
     tokens_per_batch = TOKENS_PER_BATCH[version]
     max_step = TOTAL_TOKENS[version][model_size] // tokens_per_batch
     max_sequence_length = MAX_SEQUENCE_LENGTH[version]
-    train_batch_size = tokens_per_batch // max_sequence_length
+    # train_batch_size = tokens_per_batch // max_sequence_length
+    train_batch_size = 128
 
     # Whether to use grouped query attention.
     num_kv_heads = None
+    max_step = 300
     if version in (Version.V3, Version.V3_TIKTOKEN):
         num_kv_heads = 8
 
@@ -412,6 +414,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:
@@ -504,7 +507,7 @@ def get_trainer_kwargs(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
                                 # fsdp=8 is also ok, only 2% slower step time.
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=64)
+                                mesh_shape=mesh_shape_from_axes(data=1, fsdp=128)
                             ),
                             RematSpecModifier.default_config().set(
                                 remat_policies={