Fix base.yaml config to run on 1 TPU

AI-Hypercomputer · Jan 30, 2025 · dac03f8 · dac03f8
1 parent bc7c087
commit dac03f8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchprime/torch_xla_models/configs/base.yaml b/torchprime/torch_xla_models/configs/base.yaml
@@ -23,6 +23,6 @@ lr_scheduler:
   warmup_steps: 0
 mesh:
   dcn: 1
-  fsdp: 2
+  fsdp: 4
   tensor: 1
   expert: 1