Adding a model config for 7B model and v5litepod-32-2

Changing the batch size to match the chip count

Author: lukebaumann

axlearn/experiments/text/gpt/fuji.py

@@ -15,6 +15,7 @@
 import itertools
 from typing import Any, List, NamedTuple, Optional, Union
 
+import jax
 from jax.ad_checkpoint import checkpoint_policies as jax_remat_policies
 
 from axlearn.common import causal_lm, config
@@ -248,7 +249,8 @@ 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 = len(jax.devices())
 
     # Whether to use grouped query attention.
     num_kv_heads = None
@@ -392,6 +394,24 @@ def get_trainer_kwargs(
                 # tpu-v4-(1024|2048).
                 ("tpu-v4-(1024|2048)", mesh_shape_from_axes(data=-1, fsdp=16)),
                 # tpu-v5e.
+                (
+                    "tpu-v5litepod-32-2",
+                    ChainConfigModifier.default_config().set(
+                        config_modifiers=[
+                            MeshShapeModifier.default_config().set(
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8)
+                            ),
+                            RematSpecModifier.default_config().set(
+                                remat_policies={
+                                    "model.decoder.transformer.layer": RematSpec(
+                                        prevent_cse=False,
+                                        policy=offload_dots_saveable_policy,
+                                    ),
+                                }
+                            ),
+                        ],
+                    ),
+                ),
                 (
                     "tpu-v5litepod-256",
                     ChainConfigModifier.default_config().set(