Merge branch 'main' into jax_0.8.0_py3.12_v7x

Author: andersensam

Dockerfile

@@ -25,8 +25,6 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.
 RUN mkdir -p /root
 WORKDIR /root
 # Introduce the minimum set of files for install.
-COPY README.md README.md
-COPY pyproject.toml pyproject.toml
 RUN mkdir axlearn && touch axlearn/__init__.py
 # Setup venv to suppress pip warnings.
 ENV VIRTUAL_ENV=/opt/venv
@@ -44,6 +42,7 @@ RUN pip install -qq --upgrade pip && \
 # Leverage multi-stage build for unit tests.
 FROM base AS ci
 
+COPY pyproject.toml README.md /root/
 # TODO(markblee): Remove gcp,vertexai_tensorboard from CI.
 RUN uv pip install -qq .[core,audio,orbax,dev,gcp,vertexai_tensorboard] && \
     uv cache clean
@@ -76,6 +75,7 @@ FROM base AS dataflow
 # Beam workers default to creating a new virtual environment on startup. Instead, we want them to
 # pickup the venv setup above. An alternative is to install into the global environment.
 ENV RUN_PYTHON_SDK_IN_DEFAULT_ENVIRONMENT=1
+COPY pyproject.toml README.md /root/
 RUN uv pip install -qq .[core,gcp,dataflow] && uv cache clean
 COPY . .
 
@@ -98,6 +98,7 @@ ARG INSTALL_PATHWAYS_JAXLIB=false
 
 # Ensure we install the TPU version, even if building locally.
 # Jax will fallback to CPU when run on a machine without TPU.
+COPY pyproject.toml README.md /root/
 RUN uv pip install -qq --prerelease=allow .[core,tpu] && uv cache clean
 RUN if [ -n "$EXTRAS" ]; then uv pip install -qq .[$EXTRAS] && uv cache clean; fi
 RUN if [ "$INSTALL_PATHWAYS_JAXLIB" = "true" ]; then \
@@ -119,6 +120,7 @@ RUN curl -o cuda-keyring_1.1-1_all.deb https://developer.download.nvidia.com/com
     dpkg -i cuda-keyring_1.1-1_all.deb && \
     apt-get update && apt-get install -y cuda-libraries-dev-12-9 ibverbs-utils && \
     apt clean -y
+COPY pyproject.toml README.md /root/
 RUN uv pip install --prerelease=allow .[core,gpu] && uv cache clean
 COPY . .
 

axlearn/cloud/gcp/lws_utils.py

@@ -7,7 +7,7 @@
 from absl import flags
 
 from axlearn.cloud.common.bundler import Bundler
-from axlearn.cloud.common.utils import AcceleratorConfig, FlagConfigurable, accelerator_flags
+from axlearn.cloud.common.utils import FlagConfigurable
 from axlearn.cloud.gcp.config import gcp_settings
 from axlearn.cloud.gcp.jobset_utils import TPUJobBuilder
 from axlearn.cloud.gcp.system_characteristics import USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS
@@ -27,7 +27,6 @@ class Config(FlagConfigurable.Config):
         Attributes:
             name: Name of the LeaderWorkerSet
             command: Command to be executed.
-            accelerator: Accelerator configuration.
             env_vars: Optional env vars to set.
             service_account: Optional service account to execute the job as.
             output_dir: An optional GCS path to upload LWS outputs to.
@@ -37,7 +36,6 @@ class Config(FlagConfigurable.Config):
         # TODO: Change this to be a list of str[], to support different commands
         # between leader and workers
         command: Required[str] = REQUIRED
-        accelerator: AcceleratorConfig = AcceleratorConfig()
         env_vars: dict[str, str] = {}
         service_account: Optional[str] = None
         output_dir: Optional[str] = None
@@ -47,7 +45,6 @@ class Config(FlagConfigurable.Config):
     def define_flags(cls, fv):
         super().define_flags(fv)
         common_kwargs = dict(flag_values=fv, allow_override=True)
-        accelerator_flags(**common_kwargs)
         # NOTE: the parent typically sets these flags, so we leave them as None.
         flags.DEFINE_string("name", None, "Name of the LWS.", **common_kwargs)
         flags.DEFINE_string("command", None, "Command to execute.", **common_kwargs)
@@ -77,7 +74,6 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs):
         cfg.service_account = cfg.service_account or gcp_settings(
             "k8s_service_account", default="default", fv=fv
         )
-        cfg.accelerator.set(instance_type=fv.instance_type, num_replicas=fv.num_replicas)
         return cfg
 
     def __init__(self, cfg: Config, *, bundler: Bundler):

axlearn/cloud/gcp/measurement.py

@@ -30,16 +30,16 @@
 from ml_goodput_measurement import monitoring as goodput_monitoring
 
 from axlearn.cloud.common.utils import parse_kv_flags, to_bool
-from axlearn.common import measurement
+from axlearn.common import measurement_base
 from axlearn.common.config import REQUIRED, Required, config_class, maybe_set_config
 
 
-@measurement.register_recorder("goodput")
-class GoodputRecorder(measurement.Recorder):
+@measurement_base.register_recorder("goodput")
+class GoodputRecorder(measurement_base.Recorder):
     """Records overall training goodput."""
 
     @config_class
-    class Config(measurement.Recorder.Config):
+    class Config(measurement_base.Recorder.Config):
         """Configures GoodputRecorder.
 
         Attributes:
@@ -77,7 +77,7 @@ def from_flags(cls, fv: flags.FlagValues) -> "GoodputRecorder":
         - jax_backend: The type of jax backend.
         - enable_monitoring: Boolean to enable/disable goodput monitoring (default: true).
         """
-        cfg: measurement.Recorder.Config = cls.default_config()
+        cfg: measurement_base.Recorder.Config = cls.default_config()
         parsed_flags = parse_kv_flags(fv.recorder_spec, delimiter="=")
         if "upload_interval" in parsed_flags:
             parsed_flags["upload_interval"] = int(parsed_flags["upload_interval"])
@@ -100,7 +100,7 @@ def __init__(self, cfg):
         self._logger_name = f"goodput_logger_{cfg.name}"
 
     @contextlib.contextmanager
-    def record_event(self, event: measurement.EventType, *args, **kwargs):
+    def record_event(self, event: measurement_base.EventType, *args, **kwargs):
         """Records a goodput event using a context manager."""
         # Lazily instantiate the recorder if it hasn't been already.
         if self._recorder is None:
@@ -225,7 +225,7 @@ def monitor_goodput():
 
         return monitor_goodput()
 
-    def record(self, event: measurement.Event, *args, **kwargs):
+    def record(self, event: measurement_base.Event, *args, **kwargs):
         """Deprecated: `record()` is not used in GoodputRecorder.
         Use the record_event context manager instead.
         """
@@ -238,27 +238,27 @@ def record(self, event: measurement.Event, *args, **kwargs):
                 logging_enabled=(jax.process_index() == 0),
             )
 
-        if event == measurement.Event.START_JOB:
+        if event == measurement_base.Event.START_JOB:
             self._recorder.record_job_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_JOB:
+        elif event == measurement_base.Event.END_JOB:
             self._recorder.record_job_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_STEP:
+        elif event == measurement_base.Event.START_STEP:
             self._recorder.record_step_start_time(*args, **kwargs)
-        elif event == measurement.Event.START_ACCELERATOR_INIT:
+        elif event == measurement_base.Event.START_ACCELERATOR_INIT:
             self._recorder.record_tpu_init_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_ACCELERATOR_INIT:
+        elif event == measurement_base.Event.END_ACCELERATOR_INIT:
             self._recorder.record_tpu_init_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_TRAINING_PREPARATION:
+        elif event == measurement_base.Event.START_TRAINING_PREPARATION:
             self._recorder.record_training_preparation_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_TRAINING_PREPARATION:
+        elif event == measurement_base.Event.END_TRAINING_PREPARATION:
             self._recorder.record_training_preparation_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_DATA_LOADING:
+        elif event == measurement_base.Event.START_DATA_LOADING:
             self._recorder.record_data_loading_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_DATA_LOADING:
+        elif event == measurement_base.Event.END_DATA_LOADING:
             self._recorder.record_data_loading_end_time(*args, **kwargs)
-        elif event == measurement.Event.START_CUSTOM_BADPUT_EVENT:
+        elif event == measurement_base.Event.START_CUSTOM_BADPUT_EVENT:
             self._recorder.record_custom_badput_event_start_time(*args, **kwargs)
-        elif event == measurement.Event.END_CUSTOM_BADPUT_EVENT:
+        elif event == measurement_base.Event.END_CUSTOM_BADPUT_EVENT:
             self._recorder.record_custom_badput_event_end_time(*args, **kwargs)
         else:
             logging.log_first_n(

axlearn/cloud/gcp/pathways_utils.py

@@ -311,30 +311,6 @@ def _build_pathways_head_container(self) -> dict:
             }
         )
 
-        # pylint: disable=line-too-long
-        env_list.append(
-            {
-                "name": "NUM_REPLICAS",
-                "valueFrom": {
-                    "fieldRef": {
-                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']"
-                    }
-                },
-            }
-        )
-        # pylint: enable=line-too-long
-
-        env_list.append(
-            {
-                "name": "REPLICA_ID",
-                "valueFrom": {
-                    "fieldRef": {
-                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
-                    }
-                },
-            }
-        )
-
         head_container["env"] = env_list
 
         cpu_req = f"{float(self.config.pathways_head_cpu) * 1000}m"
@@ -841,11 +817,13 @@ def default_config(cls):
         cfg = super().default_config()
         return cfg.set(inner=TPULeaderWorkerTemplate.default_config())
 
-    def __init__(self, cfg, *, bundler):
+    def __init__(self, cfg: BaseLeaderWorkerTemplate.Config, *, bundler):
         super().__init__(cfg, bundler=bundler)
+        cfg: PathwaysLeaderWorkerTemplate.Config = self.config
+
         self._bundler = bundler
         self._inner: TPULeaderWorkerTemplate = cfg.inner.instantiate(bundler=self._bundler)
-        self._tpu_type = infer_tpu_type(cfg.accelerator.instance_type)
+        self._tpu_type = infer_tpu_type(cfg.inner.accelerator.instance_type)
         if self._tpu_type not in USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS:
             raise NotImplementedError(f"Missing system characteristics for {self._tpu_type}")
 

axlearn/cloud/gcp/pathways_utils_test.py

@@ -120,6 +120,15 @@ def test_build_pathways_head_pod(self, instance_type):
                             }
                         },
                     )
+                if env_pair["name"] == "REPLICA_ID":
+                    self.assertEqual(
+                        env_pair["valueFrom"],
+                        {
+                            "fieldRef": {
+                                "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
+                            }
+                        },
+                    )
                 if env_pair["name"] == "IFRT_PROXY_LARGE_TRANSFER_THRESHOLD":
                     self.assertEqual(env_pair["value"], "1")
                 if env_pair["name"] == "IFRT_PROXY_LARGE_TRANSFER_OPTIMIZATION_DIRECTORY":

axlearn/common/flash_attention/tpu_attention.py

@@ -441,14 +441,12 @@ def kv_segment_ids_index_map(batch_index, head_index, q_seq_index, kv_seq_index)
         out_shape=out_shape,
         debug=debug,
         interpret=interpret,
-        compiler_params=dict(
-            mosaic=dict(
-                dimension_semantics=(
-                    "parallel",
-                    "parallel",
-                    "parallel",
-                    "arbitrary",
-                )
+        compiler_params=pltpu.CompilerParams(
+            dimension_semantics=(
+                "parallel",
+                "parallel",
+                "parallel",
+                "arbitrary",
             )
         ),
     )(q, k, v, ab, q_segment_ids, kv_segment_ids)
@@ -649,14 +647,12 @@ def dkv_index_map(batch_index, head_index, kv_seq_index, _):
             out_shape=out_shapes,
             debug=debug,
             interpret=interpret,
-            compiler_params=dict(
-                mosaic=dict(
-                    dimension_semantics=(
-                        "parallel",
-                        "parallel",
-                        "parallel",
-                        "arbitrary",
-                    )
+            compiler_params=pltpu.CompilerParams(
+                dimension_semantics=(
+                    "parallel",
+                    "parallel",
+                    "parallel",
+                    "arbitrary",
                 )
             ),
         )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di)
@@ -842,14 +838,12 @@ def kv_segment_ids_index_map(batch_index, head_index, q_seq_index, kv_seq_index)
             out_shape=out_shapes,
             debug=debug,
             interpret=interpret,
-            compiler_params=dict(
-                mosaic=dict(
-                    dimension_semantics=(
-                        "parallel",
-                        "parallel",
-                        "parallel",
-                        "arbitrary",
-                    )
+            compiler_params=pltpu.CompilerParams(
+                dimension_semantics=(
+                    "parallel",
+                    "parallel",
+                    "parallel",
+                    "arbitrary",
                 )
             ),
         )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di)

axlearn/common/input_dispatch.py

@@ -117,7 +117,9 @@ class Config(BaseInputDispatcher.Config):
     def __init__(self, cfg: Config, *, parent: Optional[Module]):
         cfg = cfg.clone()
         cfg.num_physical_feeds = cfg.num_physical_feeds or jax.process_count()
-        cfg.physical_feed_index = cfg.physical_feed_index or jax.process_index()
+        cfg.physical_feed_index = (
+            cfg.physical_feed_index if cfg.physical_feed_index is not None else jax.process_index()
+        )
         if cfg.logical_feed_indices is None:
             num_logical_feeds = min(cfg.global_logical_batch_size, cfg.num_physical_feeds)
             cfg.logical_feed_indices = list(range(num_logical_feeds))

axlearn/common/input_tf_data_gcs_test.py

@@ -76,7 +76,11 @@ def test_maybe_shard_examples(
                 dataset_name=dataset_name,
             )
             if expected == "even split":
-                shard_index = read_config.shard_index or jax.process_index()
+                shard_index = (
+                    read_config.shard_index
+                    if read_config.shard_index is not None
+                    else jax.process_index()
+                )
                 expected_split = tfds.even_splits(split, n=required_shards, drop_remainder=False)[
                     shard_index
                 ]