Revert updates to fuji model config

apple · Jan 14, 2025 · a0bf9df · a0bf9df
1 parent c7bc5df
commit a0bf9df
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
@@ -434,13 +434,13 @@ def get_trainer_kwargs(
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=128)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=256)
                             ),
                             RematSpecModifier.default_config().set(
                                 remat_policies={
                                     "model.decoder.transformer.layer": RematSpec(
                                         prevent_cse=True,
-                                        policy=jax_remat_policies.dots_saveable,
+                                        policy=offload_dots_saveable_policy,
                                     ),
                                 }
                             ),