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

Author: andersensam

.github/workflows/stale.yml

@@ -15,3 +15,7 @@ jobs:
           close-pr-message: "This pull request was closed because it has been inactive for more than 7 days since being marked as stale. Please feel free to reopen it if you would like to continue."
           exempt-pr-labels: "ready-to-merge"
           stale-pr-label: "stale"
+permissions:
+  actions: write
+  issues: write
+  pull-requests: write

Dockerfile

@@ -45,7 +45,7 @@ RUN pip install -qq --upgrade pip && \
 FROM base AS ci
 
 # TODO(markblee): Remove gcp,vertexai_tensorboard from CI.
-RUN uv pip install -qq .[core,audio,orbax,dev,gcp,vertexai_tensorboard,open_api] && \
+RUN uv pip install -qq .[core,audio,orbax,dev,gcp,vertexai_tensorboard] && \
     uv cache clean
 COPY . .
 

README.md

@@ -1,5 +1,7 @@
 # The AXLearn Library for Deep Learning
 
+[![build-and-test](https://github.com/apple/axlearn/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/apple/axlearn/actions/workflows/build.yml)
+
 **This library is under active development and the API is subject to change.**
 
 ## Table of Contents

axlearn/audio/encoder_asr.py

@@ -16,7 +16,7 @@
 from axlearn.audio.frontend import LogMelFrontend
 from axlearn.audio.spectrum_augmenter import SpectrumAugmenter
 from axlearn.audio.subsamplers import ConvSubSampler
-from axlearn.common.attention import SinusoidalPositionalEmbedding
+from axlearn.common.attention import RepeatedTransformerLayer, SinusoidalPositionalEmbedding
 from axlearn.common.base_layer import BaseLayer
 from axlearn.common.config import REQUIRED, Required, config_class
 from axlearn.common.conformer import RepeatedConformerLayer
@@ -165,11 +165,23 @@ def forward(self, inputs: Tensor, *, paddings: Tensor) -> dict[str, Tensor]:
             - outputs: A Tensor of shape [batch_size, seq_len, output_dim].
             - output_paddings: A 0/1 Tensor of shape [batch_size, seq_len].
         """
+        cfg = self.config
         # [batch, seq_len, input_dim].
         x = self.input_linear(inputs)
         x = self.dropout(x)
-        x = x + self.pos_emb(jnp.arange(x.shape[1]))
-        x = self.context(inputs=x, paddings=paddings)
+
+        if isinstance(cfg.context, RepeatedConformerLayer.Config):
+            x = x + self.pos_emb(jnp.arange(x.shape[1]))
+            x = self.context(inputs=x, paddings=paddings)
+        elif isinstance(cfg.context, RepeatedTransformerLayer.Config):
+            # We don't need to do add pos_emb for transformer block
+            x = self.context(data=x)
+            x = x.data
+        else:
+            raise ValueError(
+                f"The type of `self.context` ({cfg.context.klass}) "
+                "is not supported by SpeechContextNetwork."
+            )
         self._add_activation_summary(
             name="speech_context",
             activations=x,

axlearn/audio/encoder_asr_test.py

@@ -10,6 +10,8 @@
 from axlearn.audio import frontend_utils
 from axlearn.audio.encoder_asr import ASREncoder, SpeechContextNetwork, SpeechFeatureLayer
 from axlearn.audio.test_utils import fake_audio
+from axlearn.common.attention import RepeatedTransformerLayer
+from axlearn.common.kv_cache.sliding_window_kv_cache import enable_sliding_window_attention
 from axlearn.common.module import functional as F
 from axlearn.common.test_utils import TestCase
 from axlearn.common.utils import Tensor, shapes
@@ -173,6 +175,88 @@ def test_speech_context_network(self, is_training: bool):
             output_collections.summaries["activations/speech_context_norm"].weight, weights
         )
 
+    @parameterized.parameters([True, False])
+    def test_transformer(self, is_training: bool) -> None:
+        """Test the code branch with RepeatedTransformerLayer as context layer.
+
+        Args:
+            is_training: Whether the is_training code path is tested.
+        """
+        input_dim, output_dim, dropout_rate, num_layers = 32, 16, 0.2, 2
+        num_heads = 8
+        hidden_dim = 4 * input_dim
+
+        cfg = SpeechContextNetwork.default_config().set(
+            input_dim=input_dim, output_dim=output_dim, dtype=jnp.float64
+        )
+        cfg.dropout.rate = dropout_rate
+        cfg.context = RepeatedTransformerLayer.default_config().set(num_layers=num_layers)
+        attention = cfg.context.layer.self_attention.attention
+        attention.num_heads = num_heads
+        attention = enable_sliding_window_attention(attention, sliding_window_size=3)
+        cfg.context.layer.self_attention.attention = attention
+        # Dropout in transformer
+        cfg.context.layer.self_attention.dropout.rate = dropout_rate
+        cfg.context.layer.feed_forward.set(
+            hidden_dim=hidden_dim,
+        )
+
+        # Initialize layer parameters.
+        prng_key = jax.random.PRNGKey(123)
+        prng_key, init_key, input_key, length_key = jax.random.split(prng_key, num=4)
+        layer = cfg.set(name="test").instantiate(parent=None)
+        layer_params = layer.initialize_parameters_recursively(init_key)
+
+        # Generate inputs.
+        batch_size, seq_len = 4, 10
+        inputs = jnp.tile(
+            jax.random.normal(input_key, [batch_size // 2, seq_len, input_dim]), [2, 1, 1]
+        )
+        lengths = jnp.tile(
+            jax.random.randint(length_key, shape=[batch_size // 2, 1], minval=0, maxval=seq_len),
+            [2, 1],
+        )
+        paddings = jnp.arange(seq_len)[None, :] >= lengths
+        padding_data = jax.random.normal(jax.random.PRNGKey(135), inputs.shape)
+        inputs = jnp.where(paddings[..., None], padding_data, inputs)
+
+        # Compute outputs.
+        output_batch, output_collections = F(
+            layer,
+            inputs=dict(inputs=inputs, paddings=paddings),
+            is_training=is_training,
+            prng_key=prng_key,
+            state=layer_params,
+        )
+        outputs, output_paddings = output_batch["outputs"], output_batch["paddings"]
+        self.assertSequenceEqual(outputs.shape, (batch_size, seq_len, output_dim))
+        self.assertTrue(jnp.all(output_paddings == paddings))
+
+        # If is_training, outputs should always be different due to augmentation.
+        # Otherwise, outputs should be the same despite differences in padding.
+        self.assertEqual(not is_training, bool(jnp.allclose(outputs[:2], outputs[2:])))
+
+        outputs = outputs * (1 - output_paddings[:, :, None])
+        weights = jnp.sum(1 - output_paddings)
+        output_norms = jnp.sqrt(jnp.sum(outputs**2, axis=2)) / jnp.sqrt(output_dim)
+        expected_outputs_mean = jnp.sum(outputs) / weights / output_dim
+        expected_outputs_norm = jnp.sum(output_norms) / weights
+
+        self.assertNestedAllClose(
+            output_collections.summaries["activations/speech_context_mean"].mean,
+            expected_outputs_mean,
+        )
+        self.assertNestedAllClose(
+            output_collections.summaries["activations/speech_context_norm"].mean,
+            expected_outputs_norm,
+        )
+        self.assertNestedAllClose(
+            output_collections.summaries["activations/speech_context_mean"].weight, weights
+        )
+        self.assertNestedAllClose(
+            output_collections.summaries["activations/speech_context_norm"].weight, weights
+        )
+
 
 class ASREncoderTest(TestCase):
     """Tests ASREncoder."""

axlearn/cloud/gcp/jobset_utils.py

@@ -502,6 +502,29 @@ def _build_container(self) -> Nested[Any]:
         k8s_env_vars.append(
             {"name": "NODE_NAME", "valueFrom": {"fieldRef": {"fieldPath": "spec.nodeName"}}}
         )
+        # pylint: disable=line-too-long
+        k8s_env_vars.append(
+            {
+                "name": "NUM_REPLICAS",
+                "valueFrom": {
+                    "fieldRef": {
+                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']"
+                    }
+                },
+            }
+        )
+        # pylint: enable=line-too-long
+
+        k8s_env_vars.append(
+            {
+                "name": "REPLICA_ID",
+                "valueFrom": {
+                    "fieldRef": {
+                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
+                    }
+                },
+            }
+        )
 
         return dict(
             name=cfg.name,

axlearn/cloud/gcp/jobset_utils_test.py

@@ -359,6 +359,18 @@ def test_build_pod(
                 container_env["NODE_IP"]["valueFrom"]["fieldRef"]["fieldPath"],
             )
 
+            # Verify NUM_REPLICAS in container env.
+            self.assertEqual(
+                "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']",
+                container_env["NUM_REPLICAS"]["valueFrom"]["fieldRef"]["fieldPath"],
+            )
+
+            # Verify REPLICA_ID in container env.
+            self.assertEqual(
+                "metadata.annotations['jobset.sigs.k8s.io/job-index']",
+                container_env["REPLICA_ID"]["valueFrom"]["fieldRef"]["fieldPath"],
+            )
+
             # Verify uploader container specs
             self.assertEqual(len(pod_spec["initContainers"]), 1)
 

axlearn/common/array_serialization_test.py

@@ -466,6 +466,7 @@ def test_shard_info_partially_replicated(
         devices = mesh_utils.create_device_mesh((8,))
         mesh = jax.sharding.Mesh(devices.reshape((4, 2)), ("x", "y"))
         sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "y"))
+
         arr = jax.device_put(single_device_arr, sharding)
 
         replica_count = _num_replicas_per_shard(arr)
@@ -486,6 +487,7 @@ def test_shard_info_fully_sharded(self, max_data_shard_degree: int, shard_thresh
         devices = mesh_utils.create_device_mesh((8,))
         mesh = jax.sharding.Mesh(devices.reshape((4, 2)), ("x", "y"))
         sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("x", "y"))
+
         arr = jax.device_put(single_device_arr, sharding)
 
         replica_count = _num_replicas_per_shard(arr)
@@ -509,6 +511,7 @@ def test_shard_info_fully_replicated(
         devices = mesh_utils.create_device_mesh((8,))
         mesh = jax.sharding.Mesh(devices, "x")
         sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
+
         arr = jax.device_put(single_device_arr, sharding)
 
         replica_count = _num_replicas_per_shard(arr)

axlearn/common/causal_lm_test.py

@@ -8,13 +8,12 @@
 import jax
 import jax.random
 import numpy as np
-import pytest
 from absl.testing import absltest, parameterized
 from jax import numpy as jnp
 from jax.experimental.pjit import pjit
 from transformers.models.gpt2 import modeling_gpt2 as hf_gpt2
 
-from axlearn.common import causal_lm, utils
+from axlearn.common import causal_lm
 from axlearn.common.attention import (
     BaseStackedTransformerLayer,
     CausalAttentionLogitBiasLayer,
@@ -476,14 +475,12 @@ def test_no_conflict(
         )
 
     # TODO(markblee): Add a pytest marker for multi-device tests.
-    @pytest.mark.skipif(
-        jax.device_count() != 4 or jax.process_count() != 1,
-        reason=(
-            "Incorrect device & process count for mesh.\n"
-            "Use XLA_FLAGS=--xla_force_host_platform_device_count=4 to run locally."
-        ),
-    )
     def test_constrain_input_batch(self):
+        if jax.device_count() != 4 or jax.process_count() != 1:
+            self.skipTest(
+                "Incorrect device & process count for mesh.\n"
+                "Use XLA_FLAGS=--xla_force_host_platform_device_count=4 to run locally."
+            )
         model = (
             self._model_config(vocab_size=10, seq_len=10)
             .set(
@@ -817,5 +814,4 @@ def loss_fn(model_params, inputs):
 
 
 if __name__ == "__main__":
-    with utils.numeric_checks(True):
-        absltest.main()
+    absltest.main()

axlearn/common/checkpointer_test.py

@@ -21,6 +21,7 @@
 
 import jax
 import orbax.checkpoint as ocp
+import pytest
 import tensorflow as tf
 from absl import logging
 from absl.testing import absltest, parameterized
@@ -171,6 +172,7 @@ def test_save_and_restore(self, checkpointer_cls: Type[BaseCheckpointer]):
             ckpt.stop()
 
     @parameterized.parameters(Checkpointer, OrbaxCheckpointer)
+    @pytest.mark.for_8_devices
     def test_save_and_restore_mesh(self, checkpointer_cls: Type[BaseCheckpointer]):
         """Tests that we can save with one sharding and restore with a different sharding."""
         mesh_shape = (4, 2)
@@ -245,6 +247,7 @@ def state_specs(state, partition_spec):
             num_files=6,  # 1 array 4 shards (2 model, 2 data) + 1 array 2 shards (small array).
         ),
     )
+    @pytest.mark.for_8_devices
     def test_save_restore_files_count(
         self, max_data_shard_degree: int, shard_threshold_bytes: int, num_files: int
     ):
@@ -668,9 +671,11 @@ def patch_tf_io_behavior(*args):
         # pylint: disable=line-too-long
         with (
             _mesh(mesh_shape),
-            mock.patch("axlearn.common.file_system.listdir", patch_tf_io_behavior)
-            if listdir_add_trailing_slash
-            else nullcontext(),
+            (
+                mock.patch("axlearn.common.file_system.listdir", patch_tf_io_behavior)
+                if listdir_add_trailing_slash
+                else nullcontext()
+            ),
             tempfile.TemporaryDirectory() as temp_dir,
         ):
             cfg = Checkpointer.default_config().set(