Update params to work for number of devices

Author: SujeethJinesh

axlearn/cloud/gcp/jobset_utils.py

@@ -452,14 +452,14 @@ def _build_container(self) -> Nested[Any]:
             env_vars["ENABLE_ICI_RESILIENCY"] = str(cfg.enable_tpu_ici_resiliency).lower()
 
         resources = {"limits": {"google.com/tpu": system.chips_per_vm}}
-        # Set request memory by host machine type.
-        machine_memory_gi = GCE_MACHINE_TYPE_TO_MEMORY_CHARACTERISTICS.get(
-            system.gce_machine_type, None
-        )
-        if machine_memory_gi is not None:
-            request_memory_gi = machine_memory_gi * _MEMORY_REQUEST_PERCENTAGE
-            resources["limits"]["memory"] = f"{machine_memory_gi}Gi"
-            resources["requests"] = {"memory": f"{math.floor(request_memory_gi)}Gi"}
+        # # Set request memory by host machine type.
+        # machine_memory_gi = GCE_MACHINE_TYPE_TO_MEMORY_CHARACTERISTICS.get(
+        #     system.gce_machine_type, None
+        # )
+        # if machine_memory_gi is not None:
+        #     request_memory_gi = machine_memory_gi * _MEMORY_REQUEST_PERCENTAGE
+        #     resources["limits"]["memory"] = f"{machine_memory_gi}Gi"
+        #     resources["requests"] = {"memory": f"{math.floor(request_memory_gi)}Gi"}
 
         k8s_env_vars = [dict(name=k, value=str(v)) for k, v in env_vars.items()]
         k8s_env_vars.append(
@@ -509,8 +509,8 @@ def _build_uploader_container(
         interval_s = 60
         sync_command = f"while true; do gsutil -m rsync -r {src} {dst}; sleep {interval_s}; done"
         resources = {
-            "requests": {"cpu": "100m", "memory": "128Mi"},
-            "limits": {"cpu": "500m", "memory": "256Mi"},
+            # "requests": {"cpu": "100m", "memory": "128Mi"},
+            # "limits": {"cpu": "500m", "memory": "256Mi"},
         }
         return dict(
             name="output-uploader",

axlearn/common/trainer.py

@@ -624,24 +624,10 @@ def run(
 
                         self._step = self._step + 1
                         self.vlog(3, "Start step %s", self.step)
-<<<<<<< HEAD
-                        self._maybe_record_event(measurement.Event.START_STEP, self._step)
-                        output = self._run_step(
-                            utils.host_to_global_array(
-                                input_batch,
-                                partition=self._train_step_input_partition_specs(),
-                            ),
-                            force_run_evals=(
-                                force_run_eval_sets_at_max_step
-                                if self.step >= cfg.max_step
-                                else None
-                            ),
-=======
                         step_events_manager = (
                             self._recorder.record_event(measurement.Event.STEP, self.step)
                             if self._recorder
                             else contextlib.nullcontext()
->>>>>>> 3755939 (Add workload hang monitoring & rolling window goodput support)
                         )
                         with step_events_manager:
                             output = self._run_step(

axlearn/experiments/text/gpt/fuji.py

@@ -13,7 +13,9 @@
 import enum
 import functools
 import itertools
+import jax
 from typing import Any, List, NamedTuple, Optional, Union
+from absl import logging
 
 from jax.ad_checkpoint import checkpoint_policies as jax_remat_policies
 
@@ -813,6 +815,7 @@ def get_trainer_kwargs(
             ),
         )
     elif model_size == "150B":
+        logging.info("******* debugging number of devices: %s", len(jax.devices()))
         trainer_kwargs = dict(
             model_kwargs=dict(
                 num_layers=80,
@@ -827,8 +830,9 @@ def get_trainer_kwargs(
             ),
             learner_kwargs=dict(peak_lr=1.5e-4, weight_decay=0.1),
             max_sequence_length=max_sequence_length,
-            train_batch_size=train_batch_size,
-            max_step=max_step,
+            train_batch_size=len(jax.devices()), # train_batch_size,
+            max_step=10_000, # max_step,
+            save_every_n_steps=100,
             mesh_shape=mesh_shape_from_axes(data=-1, fsdp=64, model=4),
             mesh_rules=(
                 (
@@ -839,7 +843,7 @@ def get_trainer_kwargs(
                     ChainConfigModifier.default_config().set(
                         config_modifiers=[
                             MeshShapeModifier.default_config().set(
-                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=64, model=4)
+                                mesh_shape=mesh_shape_from_axes(data=-1, fsdp=256)
                             ),
                             RematSpecModifier.default_config().set(
                                 remat_policies={