change mesh shape to fit the model in memory

apple · Dec 25, 2024 · 8926ffb · 8926ffb
1 parent 0cc22e0
commit 8926ffb
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
@@ -341,7 +341,7 @@ def get_trainer_kwargs(
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=256)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=512)
                             ),
                             RematSpecModifier.default_config().set(
                                 remat_policies={
@@ -359,7 +359,7 @@ def get_trainer_kwargs(
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=256)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=512)
                             ),
                             RematSpecModifier.default_config().set(
                                 remat_policies={