Checking in

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -321,7 +321,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
             f"--server_port={_PATHWAYS_PROXY_PORT}",
             f"--gcs_scratch_location={staging_location}",
             # This should be made configurable
-            f"--num_elastic_slices={cfg.accelerator.num_replicas}"
+            f"--num_elastic_slices={cfg.accelerator.num_replicas}",
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())
 

axlearn/experiments/text/gpt/fuji.py

@@ -384,7 +384,7 @@ def get_trainer_kwargs(
             max_sequence_length=max_sequence_length,
             train_batch_size=len(jax.devices()),
             max_step=max_step,
-            save_every_n_steps=1000000,
+            save_every_n_steps=100,
             mesh_shape=mesh_shape_from_axes(data=-1, fsdp=8),
             mesh_rules=(
                 # Step time: