From afbaad37f14b5b6d8d052eb2effaa3279e2fbd45 Mon Sep 17 00:00:00 2001
From: Camilo Quinones <camilo.quinones@cienet.com>
Date: Tue, 25 Nov 2025 03:52:19 +0000
Subject: [PATCH] Ref: Changes to make it work Python 0.5.3. Topology v5p-128.
 Model 7b-fuji.

---
 axlearn/common/trainer.py              | 11 +++++------
 axlearn/experiments/text/gpt/common.py |  4 ++--
 axlearn/experiments/text/gpt/fuji.py   |  7 +++++--
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/axlearn/common/trainer.py b/axlearn/common/trainer.py
index 56a9b83df..65c4d3989 100644
--- a/axlearn/common/trainer.py
+++ b/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.
diff --git a/axlearn/experiments/text/gpt/common.py b/axlearn/experiments/text/gpt/common.py
index da86a87e9..edf0698d4 100644
--- a/axlearn/experiments/text/gpt/common.py
+++ b/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,
diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
index 96b45a070..274f46f67 100644
--- a/axlearn/experiments/text/gpt/fuji.py
+++ b/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={