diff --git a/tests/models/gemma4/audio_test.py b/tests/models/gemma4/audio_test.py new file mode 100644 index 000000000..268f69d88 --- /dev/null +++ b/tests/models/gemma4/audio_test.py @@ -0,0 +1,55 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for audio.""" + +from absl.testing import absltest +from absl.testing import parameterized +from flax import nnx +import jax +import jax.numpy as jnp +from tunix.models.gemma4 import audio + + +class AudioTest(parameterized.TestCase): + + def test_audio_encoder_shape(self): + config = audio.AudioEncoderConfig( + num_layers=2, # fewer layers for faster test + model_dims=128, + lm_model_dims=256, + atten_num_heads=4, + ) + + rngs = nnx.Rngs(0) + model = audio.AudioEncoder(config, rngs=rngs) + + batch_size = 2 + num_samples = 16000 # 1 second + + x = jax.random.normal(jax.random.PRNGKey(1), (batch_size, num_samples)) + sequence_lengths = jnp.full((batch_size,), num_samples, dtype=jnp.int32) + + output, mask = model(x, sequence_lengths) + + expected_seq_len = 25 + expected_shape = (batch_size, expected_seq_len, config.lm_model_dims) + + self.assertEqual(output.shape, expected_shape) + self.assertEqual(mask.shape, (batch_size, expected_seq_len)) + self.assertFalse(jnp.any(jnp.isnan(output))) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/models/gemma4/model_test.py b/tests/models/gemma4/model_test.py index 617cdb7a0..d60b7e814 100644 --- a/tests/models/gemma4/model_test.py +++ b/tests/models/gemma4/model_test.py @@ -220,7 +220,7 @@ def test_forward_pass_vision(self): tokens = jax.random.randint( jax.random.PRNGKey(0), (1, 32), 0, config.num_embed ) - tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER) + tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER) positions = jnp.tile( jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1) @@ -277,7 +277,7 @@ def test_forward_pass_vision_bidirectional(self): tokens = jax.random.randint( jax.random.PRNGKey(0), (1, 32), 0, config.num_embed ) - tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER) + tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER) positions = jnp.tile( jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1) @@ -337,9 +337,9 @@ def test_forward_pass_vision_batch(self): jax.random.PRNGKey(0), (batch_size, seq_len), 0, config.num_embed ) # Image placeholders: token shape represents visual soft tokens within sequences. - tokens = tokens.at[0, 10:15].set(model_lib.TOKEN_PLACEHOLDER) - tokens = tokens.at[1, 5:8].set(model_lib.TOKEN_PLACEHOLDER) - tokens = tokens.at[1, 20:25].set(model_lib.TOKEN_PLACEHOLDER) + tokens = tokens.at[0, 10:15].set(model_lib.VISION_TOKEN_PLACEHOLDER) + tokens = tokens.at[1, 5:8].set(model_lib.VISION_TOKEN_PLACEHOLDER) + tokens = tokens.at[1, 20:25].set(model_lib.VISION_TOKEN_PLACEHOLDER) positions = jnp.tile( jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1) @@ -377,6 +377,117 @@ def test_forward_pass_vision_batch(self): ) self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed)) + def test_forward_pass_audio(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 1 + config.embed_dim = 256 + config.hidden_dim = 512 + config.num_heads = 4 + config.head_dim = 64 + config.num_kv_heads = 1 + config.frac_shared_layers = 0.0 + config.audio_encoder = model_lib.audio.AudioEncoderConfig( + num_layers=1, + model_dims=64, + lm_model_dims=128, + atten_num_heads=2, + ) + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs, text_only=False) + + tokens = jax.random.randint( + jax.random.PRNGKey(0), (1, 32), 0, config.num_embed + ) + tokens = tokens.at[0, 10:15].set(model_lib.AUDIO_TOKEN_PLACEHOLDER) + + positions = jnp.tile( + jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1) + ) + attn_mask = jnp.tril( + jnp.ones((tokens.shape[1], tokens.shape[1]), dtype=jnp.bool_) + )[None, ...] + + soft_token_counts = (5,) + audio_data = jnp.zeros((1, 1, 16000), dtype=jnp.float32) + audio_lengths = jnp.full((1, 1), 16000, dtype=jnp.int32) + + audios = model_lib.PreprocessedAudioInput( + audio=audio_data, + audio_lengths=audio_lengths, + soft_token_counts=soft_token_counts, + ) + + logits, _ = model( + tokens, + positions=positions, + attention_mask=attn_mask, + audios=audios, + ) + self.assertEqual(logits.shape, (1, 32, config.num_embed)) + + def test_forward_pass_audio_batch(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 1 + config.embed_dim = 256 + config.hidden_dim = 512 + config.num_heads = 4 + config.head_dim = 64 + config.num_kv_heads = 1 + config.frac_shared_layers = 0.0 + config.audio_encoder = model_lib.audio.AudioEncoderConfig( + num_layers=1, + model_dims=64, + lm_model_dims=128, + atten_num_heads=2, + ) + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs, text_only=False) + + batch_size = 2 + seq_len = 32 + tokens = jax.random.randint( + jax.random.PRNGKey(0), (batch_size, seq_len), 0, config.num_embed + ) + tokens = tokens.at[0, 10:15].set(model_lib.AUDIO_TOKEN_PLACEHOLDER) + tokens = tokens.at[1, 5:8].set(model_lib.AUDIO_TOKEN_PLACEHOLDER) + tokens = tokens.at[1, 20:25].set(model_lib.AUDIO_TOKEN_PLACEHOLDER) + + positions = jnp.tile( + jnp.arange(tokens.shape[1])[None, :], (tokens.shape[0], 1) + ) + attn_mask = jnp.tril( + jnp.ones((tokens.shape[1], tokens.shape[1]), dtype=jnp.bool_) + )[None, ...] + attn_mask = jnp.broadcast_to(attn_mask, (batch_size, seq_len, seq_len)) + + soft_token_counts = ((5,), (3, 5)) + max_n_audios = 2 + max_samples = 16000 + + audio_data = jnp.zeros( + (batch_size, max_n_audios, max_samples), dtype=jnp.float32 + ) + audio_lengths = jnp.full( + (batch_size, max_n_audios), max_samples, dtype=jnp.int32 + ) + audio_lengths = audio_lengths.at[0, 1].set(0) + + audios = model_lib.PreprocessedAudioInput( + audio=audio_data, + audio_lengths=audio_lengths, + soft_token_counts=soft_token_counts, + ) + + logits, _ = model( + tokens, + positions=positions, + attention_mask=attn_mask, + audios=audios, + ) + self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed)) + if __name__ == "__main__": absltest.main() diff --git a/tests/processors/audio_processor_test.py b/tests/processors/audio_processor_test.py new file mode 100644 index 000000000..467b0affe --- /dev/null +++ b/tests/processors/audio_processor_test.py @@ -0,0 +1,98 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for audio_processor.""" + +from absl.testing import absltest +from absl.testing import parameterized +import numpy as np +from tunix.processors import audio_processor + + +class AudioProcessorTest(parameterized.TestCase): + + def test_compute_soft_token_count(self): + # 1 second of audio at 16kHz + self.assertEqual(audio_processor.compute_soft_token_count(16000), 25) + # 2 seconds of audio at 16kHz + self.assertEqual(audio_processor.compute_soft_token_count(32000), 50) + + def test_add_variable_extra_tokens_for_audio(self): + tokens = np.array([ + [1, 2, 258881, 3], + [4, 258881, 258881, 5], + ]) + soft_token_counts = ((2,), (1, 2)) + + expanded = audio_processor.add_variable_extra_tokens_for_audio( + tokens, + soft_token_counts=soft_token_counts, + ) + + expected_row_0 = [1, 2, 256000, -4, -4, 258883, 3] + expected_row_1 = [4, 256000, -4, 258883, 256000, -4, -4, 258883, 5] + + max_len = max(len(expected_row_0), len(expected_row_1)) + padded_expected = np.zeros((2, max_len), dtype=np.int32) + padded_expected[0, : len(expected_row_0)] = expected_row_0 + padded_expected[1, : len(expected_row_1)] = expected_row_1 + + np.testing.assert_array_equal(expanded, padded_expected) + + def test_process_gemma4_audio_inputs_batch(self): + class _DummyAudioConfig: + sample_rate = 16000 + audio_seq_length = 750 + + config = _DummyAudioConfig() + + audio = [[np.zeros(16000), np.zeros(8000)], [np.zeros(32000)]] + tokens = [ + np.array([1, 258881, 3, 258881]), + np.array([258881, 4]), + ] + + processed, new_tokens = audio_processor.process_gemma4_audio_inputs( + audio=audio, + tokens=tokens, + audio_encoder_config=config, + pad_id=0, + ) + + self.assertEqual(processed.audio.shape, (2, 2, 32000)) + self.assertEqual(processed.audio_lengths.shape, (2, 2)) + + np.testing.assert_array_equal(processed.audio_lengths[0], [16000, 8000]) + np.testing.assert_array_equal(processed.audio_lengths[1], [32000, 0]) + + self.assertEqual(processed.soft_token_counts, ((25, 12), (50,))) + + self.assertLen(new_tokens, 2) + expected_new_tokens_0 = np.array( + [1] + + [256000] + + [-4] * 25 + + [258883] + + [3] + + [256000] + + [-4] * 12 + + [258883] + ) + expected_new_tokens_1 = np.array([256000] + [-4] * 50 + [258883] + [4]) + np.testing.assert_array_equal(new_tokens[0], expected_new_tokens_0) + np.testing.assert_array_equal(new_tokens[1], expected_new_tokens_1) + + +if __name__ == "__main__": + absltest.main() diff --git a/tunix/generate/sampler.py b/tunix/generate/sampler.py index 95020e043..12397e087 100644 --- a/tunix/generate/sampler.py +++ b/tunix/generate/sampler.py @@ -36,6 +36,7 @@ from tunix.generate import utils import tunix.generate.beam_search as beam_search_lib import tunix.generate.tokenizer_adapter as tok_adapter +from tunix.processors import audio_processor from tunix.processors import image_processor LayerCache = dict[str, jaxtyping.Array] @@ -562,6 +563,7 @@ def _prefill_fn( params: statelib.State, sampler_state: _SamplingState, images: jnp.ndarray | None = None, + audios: Any = None, echo: bool = True, ) -> _SamplingState: """Performs prefill.""" @@ -600,7 +602,11 @@ def _prefill_fn( ) transformer = nnx.merge(self._transformer_graphdef, params) - kwargs = {} if images is None else {'images': images} + kwargs = {} + if images is not None: + kwargs['images'] = images + if audios is not None: + kwargs['audios'] = audios decode_only_last_token = self._supports_decode_only_last_token and not echo if decode_only_last_token: kwargs['decode_only_last_token'] = True @@ -759,6 +765,13 @@ def __call__( | jnp.ndarray | None ) = None, + audios: ( + str + | np.ndarray + | list[str | np.ndarray | list[str | np.ndarray] | None] + | jnp.ndarray + | None + ) = None, ) -> base_sampler.SamplerOutput: """Samples a completion of the input string. @@ -804,16 +817,16 @@ def __call__( tokens = [self.tokenize(x) for x in input_strings] processed_images = images - is_gemma4_multimodal = ( + is_gemma4_vision = ( hasattr(self.transformer, 'vision_encoder') and self.transformer.vision_encoder is not None ) - if is_gemma4_multimodal and images is not None: + if is_gemma4_vision and images is not None: processed_images, tokens = image_processor.process_gemma4_inputs( images, tokens, - self.transformer.vision_encoder, + getattr(self.transformer, 'vision_encoder'), self.tokenizer.pad_id(), ) diff --git a/tunix/models/gemma4/audio.py b/tunix/models/gemma4/audio.py new file mode 100644 index 000000000..cdac49376 --- /dev/null +++ b/tunix/models/gemma4/audio.py @@ -0,0 +1,853 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Audio encoder implementation for Gemma4.""" + +import dataclasses +from typing import Optional + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np +from tunix.models.gemma4 import vision as vision_lib +from tunix.utils import compat + +ClippedEinsum = vision_lib.ClippedEinsum +RMSNorm = vision_lib.RMSNorm + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class AudioEncoderConfig: + """Configuration for the AudioEncoder.""" + + num_layers: int = 12 + model_dims: int = 1024 + lm_model_dims: int = 1536 + atten_num_heads: int = 8 + atten_left_context: int = 13 + atten_right_context: int = 0 + conv_kernel_size: int = 5 + gradient_clipping: float = 10_000_000_000.0 + conf_reduction_factor: int = 1 + sample_rate: int = 16000 + audio_seq_length: int = 750 + param_dtype: jnp.dtype = jnp.float32 + compute_dtype: Optional[jnp.dtype] = None + + +class GemaxMelFilterbank(nnx.Module): + """Computes Mel-filterbanks from a raw audio waveform.""" + + def __init__( + self, + sample_rate: int = 16000, + win_length: int = 320, + hop_length: int = 160, + subframe_factor: int = 160, + n_mels: int = 128, + f_min: float = 0.0, + f_max: float = 8000.0, + num_mel_bins: float = 128.0, + constant: float = 0.001, + ): + self.sample_rate = sample_rate + self.win_length = win_length + self.hop_length = hop_length + self.subframe_factor = subframe_factor + self.n_mels = n_mels + self.f_min = f_min + self.f_max = f_max + self.num_mel_bins = num_mel_bins + self.constant = constant + + assert self.win_length > self.hop_length + self.n_fft = int(2 ** np.ceil(np.log2(self.win_length))) + + # Pre-compute window and mel basis + self.window = self.hann_window(self.win_length, True, True) + self.mel_basis = self.linear_to_mel_weight_matrix()[ + jnp.newaxis, :, : + ].transpose(0, 2, 1) + + def hertz_to_mel(self, freq): + return 2595.0 * np.log10(1.0 + (freq / 700.0)) + + def mel_to_hertz(self, mels): + return 700.0 * (np.power(10, mels / 2595.0) - 1.0) + + def _create_triangular_filter_bank(self, fft_freqs, filter_freqs): + filter_diff = np.diff(filter_freqs) + slopes = np.expand_dims(filter_freqs, 0) - np.expand_dims(fft_freqs, 1) + down_slopes = -slopes[:, :-2] / filter_diff[:-1] + up_slopes = slopes[:, 2:] / filter_diff[1:] + return np.maximum(0.0, np.minimum(down_slopes, up_slopes)) + + def linear_to_mel_weight_matrix(self) -> jnp.ndarray: + num_spectrogram_bins = int(self.n_fft / 2) + 1 + nyquist_hertz = self.sample_rate / 2.0 + linear_frequencies = np.linspace( + 0.0, nyquist_hertz, num_spectrogram_bins, dtype=np.float64 + ) + + mel_min = self.hertz_to_mel(self.f_min) + mel_max = self.hertz_to_mel(self.f_max) + mel_freqs = np.linspace(mel_min, mel_max, int(self.num_mel_bins) + 2) + filter_freqs = self.mel_to_hertz(mel_freqs) + + mel_weights_matrix = self._create_triangular_filter_bank( + linear_frequencies, filter_freqs + ) + return jnp.array(mel_weights_matrix.T.astype(np.float32)) + + def hann_window( + self, window_length: int, periodic: bool, nonzero: bool = False + ) -> jnp.ndarray: + if nonzero: + arg = jnp.pi * 2.0 / window_length + return 0.5 - ( + 0.5 + * jnp.cos(arg * (jnp.arange(window_length, dtype=jnp.float32) + 0.5)) + ) + a = 0.5 + b = 1 - a + even = 1 - window_length % 2 + n = jnp.asarray(window_length + int(periodic) * even - 1, dtype=jnp.float32) + count = jnp.arange(window_length, dtype=jnp.float32) + cos_arg = 2 * jnp.pi * count / n + hann_values = a - b * jnp.cos(cos_arg) + return hann_values + + def __call__(self, waveform: jax.Array) -> jax.Array: + waveform = waveform.reshape(waveform.shape[0], 1, -1) + assert len(waveform.shape) == 3, "Must be [batch, 1, seq_len]" + assert waveform.shape[1] == 1, "Must be 1" + + frame_size_for_unfold = self.win_length + 1 + seq_len = waveform.shape[-1] + num_frames = (seq_len - frame_size_for_unfold) // self.hop_length + 1 + + start_indices = (jnp.arange(num_frames) * self.hop_length)[:, jnp.newaxis] + window_indices = jnp.arange(frame_size_for_unfold)[jnp.newaxis, :] + indices = start_indices + window_indices + + frames = waveform[:, 0, :][:, indices] + frames = frames[..., :-1] + + windowed_frames = frames * self.window + stft_spectrogram = jnp.fft.rfft(windowed_frames, n=self.n_fft) + spectrogram = jnp.abs(stft_spectrogram) + + batch_size = spectrogram.shape[0] + mel_basis = jnp.repeat(self.mel_basis, batch_size, axis=0) + mel_spectrogram = spectrogram @ mel_basis + mel_spectrogram += self.constant + mel_spectrogram = jnp.log(mel_spectrogram) + return mel_spectrogram + + +class SubSamplingBlock(nnx.Module): + """Subsampling block (Conv+LN+ReLU) reducing temporal dimension.""" + + def __init__( + self, + input_features: int, + output_proj_dim: int, + *, + rngs: nnx.Rngs, + dtype: jnp.dtype = jnp.float32, + ): + # conv0: [B, T, F, 1] -> [B, T/2, F/2, 128] + self.conv0 = nnx.Conv( + in_features=1, + out_features=128, + kernel_size=(3, 3), + strides=(2, 2), + padding=((1, 1), (1, 1)), + use_bias=False, + dtype=dtype, + rngs=rngs, + ) + self.norm0 = nnx.LayerNorm( + num_features=128, use_bias=False, use_scale=True, rngs=rngs + ) + + # conv1: [B, T/2, F/2, 128] -> [B, T/4, F/4, 32] + self.conv1 = nnx.Conv( + in_features=128, + out_features=32, + kernel_size=(3, 3), + strides=(2, 2), + padding=((1, 1), (1, 1)), + use_bias=False, + dtype=dtype, + rngs=rngs, + ) + self.norm1 = nnx.LayerNorm( + num_features=32, use_bias=False, use_scale=True, rngs=rngs + ) + + # Project collapsed features to output_proj_dim + collapsed_features = (input_features // 4) * 32 + self.input_proj = nnx.Linear( + in_features=collapsed_features, + out_features=output_proj_dim, + use_bias=False, + dtype=dtype, + rngs=rngs, + ) + + def __call__( + self, x: jax.Array, mask: jax.Array + ) -> tuple[jax.Array, jax.Array]: + x = jnp.expand_dims(x, -1) # [batch, time, features, 1] + + x = self.conv0(x) + mask = mask[:, ::2][:, : x.shape[1]] + x = self.norm0(x) + x = jax.nn.relu(x) + + x = self.conv1(x) + mask = mask[:, ::2][:, : x.shape[1]] + x = self.norm1(x) + x = jax.nn.relu(x) + + b, t, f, c = x.shape + x = jnp.reshape(x, (b, t, f * c)) + x = self.input_proj(x) + return x, mask + + +class FFNBlock(nnx.Module): + """Residual FFN block with RMSNorm.""" + + def __init__( + self, + config: AudioEncoderConfig, + ffn_residual_weight: float = 0.5, + *, + rngs: nnx.Rngs, + ): + self.config = config + self.ffn_residual_weight = ffn_residual_weight + + self.pre_layer_norm = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.ffn_layer1 = ClippedEinsum( + shape=(config.model_dims, config.model_dims * 4), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.ffn_layer2 = ClippedEinsum( + shape=(config.model_dims * 4, config.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.post_layer_norm = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + residual = x + x = jnp.clip( + x, -self.config.gradient_clipping, self.config.gradient_clipping + ) + + y = self.pre_layer_norm(x) + y = self.ffn_layer1("...D,DF->...F", y) + y = jax.nn.swish(y) + y = self.ffn_layer2("...D,DF->...F", y) + y = jnp.clip( + y, -self.config.gradient_clipping, self.config.gradient_clipping + ) + y = self.post_layer_norm(y) + return residual + y * self.ffn_residual_weight + + +class TransformerXLRelativePositionEmbedding(nnx.Module): + """Relative position embedding logic from Transformer-XL.""" + + def __init__(self, config: AudioEncoderConfig, *, rngs: nnx.Rngs): + self.config = config + self.atten_num_heads = config.atten_num_heads + self.units_per_head = config.model_dims // config.atten_num_heads + self.model_dims = config.model_dims + self.atten_left_context = config.atten_left_context + self.atten_right_context = config.atten_right_context + + assert ( + self.atten_right_context == 0 + ), "Not yet implemented for right context" + + self.pos_proj = nnx.Linear( + self.model_dims, + self.model_dims, + use_bias=False, + kernel_init=nnx.initializers.glorot_uniform(), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + @staticmethod + def _get_timing_signal_1d_pos( + position: jnp.ndarray, + channels: int, + min_timescale: float = 1.0, + max_timescale: float = 1.0e4, + dtype: jnp.dtype = jnp.float32, + ) -> jnp.ndarray: + position = jnp.asarray(position, jnp.float32) + num_timescales = channels // 2 + log_timescale_increment = jnp.log( + float(max_timescale) / float(min_timescale) + ) / max(num_timescales - 1, 1) + inv_timescales = min_timescale * jnp.exp( + jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment + ) + scaled_time = ( + position[:, :, jnp.newaxis] + * inv_timescales[jnp.newaxis, jnp.newaxis, :] + ) + timing_signal = jnp.concatenate( + [jnp.sin(scaled_time), jnp.cos(scaled_time)], axis=2 + ) + timing_signal = jnp.pad(timing_signal, [[0, 0], [0, 0], [0, channels % 2]]) + return timing_signal.astype(dtype) + + def __call__(self, queries: jax.Array, keys: jax.Array) -> jax.Array: + term_ac = jnp.einsum( + "BuwNH,BucNH->BNuwc", + queries, + keys, + precision="highest", + ) + b = queries.shape[0] + u = queries.shape[1] + w = queries.shape[2] + c = keys.shape[2] + n = self.atten_num_heads + l = max(0, self.atten_left_context - 1) + r = self.atten_right_context + assert c == w + l + r + + pos = jnp.arange(l, -r - 1, -1)[jnp.newaxis, :] + sin_emb = self._get_timing_signal_1d_pos( + pos, + self.model_dims, + min_timescale=1, + max_timescale=10000, + dtype=queries.dtype, + ) + sin_emb = self.pos_proj(sin_emb) + sin_emb = sin_emb.reshape( + 1, l + r + 1, self.atten_num_heads, self.units_per_head + ) + sin_emb = jnp.squeeze(sin_emb, 0) + + term_bd = jnp.einsum( + "BuwNH,FNH->BNuwF", + queries, + sin_emb, + precision="float32", + ) + term_bd = jnp.pad( + term_bd, + [ + [0, 0], + [0, 0], + [0, 0], + [0, 0], + [0, (c + 1) - (l + r + 1)], + ], + constant_values=jnp.array(0, dtype=term_bd.dtype), + ) + term_bd = jnp.reshape( + term_bd, + [b, n, u, w * (c + 1)], + ) + term_bd = term_bd[:, :, :, : w * c] + term_bd = jnp.reshape( + term_bd, + [b, n, u, w, c], + ) + return term_ac + term_bd + + +class LocalDotProductAttention(nnx.Module): + """Local dot-product self-attention with relative position embeddings.""" + + block_size: int = 12 + + def __init__( + self, + config: AudioEncoderConfig, + *, + rngs: nnx.Rngs, + ): + self.config = config + self.atten_num_heads = config.atten_num_heads + self.units_per_head = config.model_dims // config.atten_num_heads + self.model_dims = config.model_dims + self.atten_left_context = config.atten_left_context + self.atten_right_context = config.atten_right_context + self.attention_logits_soft_capping = 50.0 + + self.query = ClippedEinsum( + shape=(self.model_dims, self.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.key = ClippedEinsum( + shape=(self.model_dims, self.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.value = ClippedEinsum( + shape=(self.model_dims, self.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + self.per_dim_scale = nnx.Param( + jnp.ones( + (self.units_per_head,), dtype=config.param_dtype or jnp.float32 + ) + ) + self.relative_position_embedding = TransformerXLRelativePositionEmbedding( + config=config, + rngs=rngs, + ) + + @staticmethod + def _extract_block_context( + x: jnp.ndarray, + block_size: int, + left_context: int, + right_context: int, + padding_val: float | jnp.bool_ = 0.0, + ) -> jnp.ndarray: + if block_size < 1: + raise ValueError(f"{block_size=} must be at least 1.") + paddings = [(0, 0)] * len(x.shape) + paddings[1] = (left_context, right_context + block_size - 1) + x = jnp.pad(x, paddings, constant_values=jnp.asarray(padding_val, x.dtype)) + + frame_length = block_size + left_context + right_context + frame_step = block_size + num_frames = (x.shape[1] - frame_length) // frame_step + 1 + + start_indices = jnp.arange(num_frames) * frame_step + relative_indices = jnp.arange(frame_length) + indices = start_indices[:, jnp.newaxis] + relative_indices[jnp.newaxis, :] + return jnp.take(x, indices, axis=1) + + @staticmethod + def _convert_to_block( + x: jnp.ndarray, block_size: int, padding_val: float = 0.0 + ) -> jnp.ndarray: + shape = x.shape + b, t = shape[0], shape[1] + if block_size < 1: + raise ValueError(f"{block_size=} must be at least 1.") + num_blocks = (t + block_size - 1) // block_size + pad_length = num_blocks * block_size - t + + if pad_length > 0: + paddings = [[0, 0]] * len(shape) + paddings[1] = [0, pad_length] + x = jnp.pad(x, paddings, constant_values=jnp.array(padding_val, x.dtype)) + reshaped = jnp.reshape(x, (b, num_blocks, block_size) + shape[2:]) + return reshaped + + @staticmethod + def ones_matrix_band_part( + rows: int, + cols: int, + num_lower: int, + num_upper: int, + out_dtype: jnp.dtype = jnp.float32, + out_shape: Optional[tuple[int, ...]] = None, + ) -> jnp.ndarray: + m = jnp.arange(rows).reshape((rows, 1)) + n = jnp.arange(cols).reshape((1, cols)) + mask_lower = True + if num_lower >= 0: + mask_lower = (m - n) <= num_lower + mask_upper = True + if num_upper >= 0: + mask_upper = (n - m) <= num_upper + band = jnp.logical_and(mask_lower, mask_upper).astype(out_dtype) + if out_shape: + band = jnp.reshape(band, out_shape) + return band + + def __call__( + self, x: jax.Array, mask: jax.Array, causal_valid_mask: jax.Array + ) -> jax.Array: + batch_size, seq_len, _ = x.shape + + q = self.query("...D,DF->...F", x) + k = self.key("...D,DF->...F", x) + v = self.value("...D,DF->...F", x) + + q = q.reshape( + batch_size, seq_len, self.atten_num_heads, self.units_per_head + ).astype("float32") + k = k.reshape( + batch_size, seq_len, self.atten_num_heads, self.units_per_head + ).astype("float32") + v = v.reshape( + batch_size, seq_len, self.atten_num_heads, self.units_per_head + ).astype("float32") + + r_softplus_0 = 1.442695041 + query_scale = jnp.array( + r_softplus_0 / jnp.sqrt(self.units_per_head), dtype=q.dtype + ) + q *= query_scale * jax.nn.softplus(self.per_dim_scale.value.astype(q.dtype)) + + key_scale = jnp.array( + r_softplus_0 * jax.nn.softplus(jnp.ones(())), dtype=k.dtype + ) + k *= key_scale + + q = q.astype("float32") + k = k.astype("float32") + + original_query_time = q.shape[1] + k_context = self._extract_block_context( + k, + self.block_size, + max(0, self.atten_left_context - 1), + self.atten_right_context, + ) + q_blocked = self._convert_to_block(q, block_size=self.block_size) + + logits = self.relative_position_embedding(q_blocked, k_context) + logits = self.attention_logits_soft_capping * jnp.tanh( + logits / self.attention_logits_soft_capping + ) + + num_query_blocks = q_blocked.shape[1] + valid_mask_blocked = self._extract_block_context( + mask, + self.block_size, + max(0, self.atten_left_context - 1), + self.atten_right_context, + padding_val=jnp.bool_(False), + ) + valid_mask_blocked = valid_mask_blocked[:, jnp.newaxis, :, jnp.newaxis, :] + valid_mask_blocked = jnp.logical_and( + valid_mask_blocked, + causal_valid_mask[jnp.newaxis, jnp.newaxis, jnp.newaxis, :, :], + ) + + logits = jnp.where( + valid_mask_blocked, + logits, + jnp.asarray(-1e9, dtype=logits.dtype), + ) + probabilities = jax.nn.softmax(logits, axis=-1).astype("float32") + + values_blocks = self._extract_block_context( + v, + self.block_size, + max(0, self.atten_left_context - 1), + self.atten_right_context, + ) + context_vectors = jnp.einsum( + "BNuwc,BucNH->BuwNH", + probabilities, + values_blocks.astype("float32"), + precision="float32", + ) + context_vectors = jnp.reshape( + context_vectors, + [ + batch_size, + num_query_blocks * self.block_size, + self.atten_num_heads, + self.units_per_head, + ], + ) + return context_vectors[:, :original_query_time] + + +class AttentionBlock(nnx.Module): + """Attention block wrapping the local attention mechanism.""" + + def __init__(self, config: AudioEncoderConfig, *, rngs: nnx.Rngs): + self.config = config + self.pre_norm = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.self_atten = LocalDotProductAttention(config=config, rngs=rngs) + self.post = ClippedEinsum( + shape=( + config.atten_num_heads, + config.model_dims // config.atten_num_heads, + config.model_dims, + ), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.post_norm = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + def __call__( + self, x: jax.Array, mask: jax.Array, causal_valid_mask: jax.Array + ) -> jax.Array: + residual = x + x = jnp.clip( + x, -self.config.gradient_clipping, self.config.gradient_clipping + ) + y = self.pre_norm(x) + y = self.self_atten(y, mask, causal_valid_mask) + y = self.post("...NH,NHD->...D", y) + y = jnp.clip( + y, -self.config.gradient_clipping, self.config.gradient_clipping + ) + y = self.post_norm(y) + return residual + y + + +class LightweightConvBlock(nnx.Module): + """Residual lightweight 1D convolutional block.""" + + def __init__(self, config: AudioEncoderConfig, *, rngs: nnx.Rngs): + self.config = config + self.ln = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.linear_start = ClippedEinsum( + shape=(config.model_dims, 2 * config.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.depthwise_conv1d = nnx.Conv( + in_features=config.model_dims, + out_features=config.model_dims, + kernel_size=(config.conv_kernel_size,), + strides=(1,), + padding="CAUSAL", + feature_group_count=config.model_dims, + use_bias=False, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.conv_norm = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + self.linear_end = ClippedEinsum( + shape=(config.model_dims, config.model_dims), + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + def __call__(self, x: jax.Array) -> jax.Array: + residual = x + y = self.ln(x) + gated_input = self.linear_start("...D,DF->...F", y) + y = jax.nn.glu(gated_input) + y = self.depthwise_conv1d(y) + y = jnp.clip( + y, -self.config.gradient_clipping, self.config.gradient_clipping + ) + y = self.conv_norm(y) + y = jax.nn.swish(y) + y = self.linear_end("...D,DF->...F", y) + return residual + y + + +class ConformerLayer(nnx.Module): + """A single layer of the Conformer model.""" + + def __init__( + self, config: AudioEncoderConfig, layer_idx: int = -1, *, rngs: nnx.Rngs + ): + self.config = config + self.layer_idx = layer_idx + self.fflayer_start = FFNBlock(config, ffn_residual_weight=0.5, rngs=rngs) + self.trans_atten = AttentionBlock(config, rngs=rngs) + self.lconv = LightweightConvBlock(config, rngs=rngs) + self.fflayer_end = FFNBlock(config, ffn_residual_weight=0.5, rngs=rngs) + self.final_ln = RMSNorm( + dim=config.model_dims, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + def __call__( + self, x: jax.Array, mask: jax.Array, causal_valid_mask: jax.Array + ) -> jax.Array: + x = self.fflayer_start(x) + x = self.trans_atten(x, mask, causal_valid_mask) + validity_mask = mask[:, :, jnp.newaxis].astype(x.dtype) + x = x * validity_mask + x = self.lconv(x) + x = self.fflayer_end(x) + x = jnp.clip( + x, -self.config.gradient_clipping, self.config.gradient_clipping + ) + x = self.final_ln(x) + return x + + +class AudioEncoder(nnx.Module): + """Conformer-based Audio Encoder.""" + + def __init__(self, config: AudioEncoderConfig, *, rngs: nnx.Rngs): + self.config = config + self.mel_filterbank = GemaxMelFilterbank( + sample_rate=config.sample_rate, + win_length=320, + hop_length=160, + subframe_factor=160, + n_mels=128, + f_min=0.0, + f_max=8000.0, + num_mel_bins=128.0, + constant=0.001, + ) + self.subsampling = SubSamplingBlock( + input_features=128, + output_proj_dim=config.model_dims, + rngs=rngs, + dtype=config.compute_dtype or jnp.float32, + ) + self.blocks = compat.ModuleList([ + ConformerLayer(config, layer_idx=i, rngs=rngs) + for i in range(config.num_layers) + ]) + self.output_projection = nnx.Linear( + config.model_dims, + config.lm_model_dims, + use_bias=True, + dtype=config.compute_dtype or jnp.float32, + param_dtype=config.param_dtype, + rngs=rngs, + ) + + def to_float32(self, audio_data: jnp.ndarray): + if audio_data.dtype == jnp.int16: + return audio_data.astype(jnp.float32) / 32768.0 + elif audio_data.dtype == jnp.int32: + return audio_data.astype(jnp.float32) / 2147483648.0 + elif audio_data.dtype == jnp.uint8: + return (audio_data.astype(jnp.float32) - 128.0) / 128.0 + elif audio_data.dtype in [jnp.float16, jnp.float32]: + return audio_data.astype(jnp.float32) + else: + raise ValueError(f"Unsupported format: {audio_data.dtype}") + + def infer_mask( + self, x: jnp.ndarray, sequence_lengths: jnp.ndarray, original_seq_len: int + ) -> jnp.ndarray: + compressed_seq_len = x.shape[1] + compression_rate = original_seq_len / compressed_seq_len + new_sequence_lengths = jnp.floor( + sequence_lengths / compression_rate + ).astype(jnp.int32) + indices = jnp.arange(compressed_seq_len)[jnp.newaxis, :] + mask = indices < new_sequence_lengths[:, jnp.newaxis] + return mask + + @staticmethod + def _compute_causal_valid_mask(config: AudioEncoderConfig): + chunk_size = LocalDotProductAttention.block_size + max_future_horizon = config.atten_right_context + max_past_horizon = max(0, config.atten_left_context - 1) + context_size = chunk_size + max_past_horizon + max_future_horizon + upper_diagonal = max_past_horizon + max_future_horizon + + lower_causal_mask = LocalDotProductAttention.ones_matrix_band_part( + context_size, + chunk_size, + num_lower=-1, + num_upper=0, + out_dtype=jnp.bool_, + ).T + upper_causal_mask = LocalDotProductAttention.ones_matrix_band_part( + chunk_size, + context_size, + num_lower=-1, + num_upper=upper_diagonal, + out_dtype=jnp.bool_, + ) + causal_valid_mask = lower_causal_mask & upper_causal_mask + return causal_valid_mask + + def __call__( + self, x: jax.Array, sequence_lengths: jax.Array + ) -> tuple[jax.Array, jax.Array]: + x = self.to_float32(x) + original_seq_len = x.shape[-1] + + # Mel Filterbank + x = self.mel_filterbank(x) + + # Infer mask and apply + mask = self.infer_mask(x, sequence_lengths, original_seq_len) + x = jnp.where(mask[:, :, jnp.newaxis], x, 0.0) + + # Subsampling + x, mask = self.subsampling(x, mask) + + causal_valid_mask = self._compute_causal_valid_mask(self.config) + + # Conformer stack + for block in self.blocks: + x = block(x, mask, causal_valid_mask) + + if self.config.conf_reduction_factor > 1: + x = x[:, :: self.config.conf_reduction_factor] + mask = mask[:, :: self.config.conf_reduction_factor] + + # Final projection + x = self.output_projection(x) + + x = jnp.where((~mask)[:, :, jnp.newaxis], 0.0, x) + + return x, mask diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py index e27800c8e..102148fbc 100644 --- a/tunix/models/gemma4/model.py +++ b/tunix/models/gemma4/model.py @@ -32,13 +32,15 @@ import jaxtyping import numpy as np from tunix.generate.mappings import BackendMappingMixin +from tunix.models.gemma4 import audio from tunix.models.gemma4 import moe from tunix.models.gemma4 import vision from tunix.utils import compat from tunix.utils import env_utils from tunix.utils.sharding_utils import shard -TOKEN_PLACEHOLDER = -2 +VISION_TOKEN_PLACEHOLDER = -2 +AUDIO_TOKEN_PLACEHOLDER = -4 @dataclasses.dataclass(frozen=True) @@ -55,6 +57,20 @@ class PreprocessedVisionInput: ) +@dataclasses.dataclass(frozen=True) +class PreprocessedAudioInput: + audio: Any + audio_lengths: Any + soft_token_counts: tuple[tuple[int, ...], ...] + + +jax.tree_util.register_dataclass( + PreprocessedAudioInput, + data_fields=['audio', 'audio_lengths'], + meta_fields=['soft_token_counts'], +) + + env_utils.setup_sharding_environment() @@ -85,6 +101,7 @@ class ShardingConfig: act_btnh: Tuple[str | None, ...] vision_proj: Tuple[str | None, ...] vision_soft_emb_norm_weight: Tuple[str | None, ...] + audio_proj: Tuple[str | None, ...] # MoE sharding exp_weight_edf: Tuple[str | None, ...] exp_weight_efd: Tuple[str | None, ...] @@ -113,13 +130,16 @@ def get_default_sharding(is_sampling: bool = False): act_btnh=('fsdp', None, 'tp', None), vision_proj=(fsdp, 'tp'), vision_soft_emb_norm_weight=('tp',), + audio_proj=(fsdp, 'tp'), exp_weight_edf=(fsdp, None, None, 'tp'), exp_weight_efd=(fsdp, 'tp', None), per_layer_model_projection=(fsdp, None, 'tp'), per_layer_input_gate=(fsdp, 'tp'), per_layer_projection=('tp', fsdp), per_layer_input_embedding=('tp', None, fsdp), - vision_shd=vision.VisionShardingConfig.get_default_sharding(is_sampling), + vision_shd=vision.VisionShardingConfig.get_default_sharding( + is_sampling + ), ) @@ -170,6 +190,9 @@ class ModelConfig: vision_encoder: vision.VisionEncoderConfig | None = None use_bidirectional_attention: str | None = None + # Audio config + audio_encoder: audio.AudioEncoderConfig | None = None + def __post_init__(self): # TODO(tunix-dev): support flash attention with sliding window KV cache if self.use_sliding_window_kv_cache and self.use_flash_attention: @@ -377,6 +400,24 @@ def __init__( with_scale=False, ) + if config.audio_encoder is not None: + self.mm_audio_input_projection = Einsum( + einsum_str='...tm,md->...td', + shape=(config.audio_encoder.lm_model_dims, self.embed_dim), + sharding=config.shd_config.audio_proj, + rngs=rngs, + dtype=self.config.dtype, + param_dtype=self.param_dtype, + ) + self.mm_audio_pre_projection_norm = RMSNorm( + config.audio_encoder.lm_model_dims, + rngs=rngs, + sharding=config.shd_config, + dtype=self.config.dtype, + param_dtype=self.param_dtype, + with_scale=False, + ) + def encode(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: x = self.input_embedding[(x,)] x *= jnp.sqrt(x.shape[-1]).astype(x.dtype) @@ -389,6 +430,11 @@ def encode_vision(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: x = self.mm_input_projection(x) return x + def encode_audio(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: + x = self.mm_audio_pre_projection_norm(x) + x = self.mm_audio_input_projection(x) + return x + def encode_per_layer_input( self, x: jaxtyping.ArrayLike, t: jaxtyping.ArrayLike ) -> jaxtyping.Array: @@ -1423,8 +1469,10 @@ def __init__( self, config: ModelConfig, *, rngs: nnx.Rngs, text_only: bool = True ): self.text_only = text_only - if text_only and config.vision_encoder is not None: - config = dataclasses.replace(config, vision_encoder=None) + if text_only: + config = dataclasses.replace( + config, vision_encoder=None, audio_encoder=None + ) self.config = config self.embedder = Embedder(config, rngs=rngs) @@ -1436,6 +1484,12 @@ def __init__( shd_config=config.shd_config.vision_shd, ) + if config.audio_encoder is not None: + self.audio_encoder = audio.AudioEncoder( + config=config.audio_encoder, + rngs=rngs, + ) + pattern = ( config.attention_pattern if config.attention_pattern @@ -1491,6 +1545,7 @@ def __call__( segment_ids=None, decode_only_last_token: bool = False, images: PreprocessedVisionInput | None = None, + audios: PreprocessedAudioInput | None = None, skip_lm_head: bool = False, ): if positions is None: @@ -1504,13 +1559,22 @@ def __call__( x = self.embedder.encode(tokens) if self.config.vision_encoder is not None and images is not None: soft_embeddings = self._encode_vision(images) - mask = tokens == TOKEN_PLACEHOLDER + mask = tokens == VISION_TOKEN_PLACEHOLDER x = merge_flat_embeddings( text_embeddings=x, multimodal_embeddings=soft_embeddings, mask=mask, ) + if self.config.audio_encoder is not None and audios is not None: + soft_audio_embeddings = self._encode_audio(audios) + audio_mask = tokens == AUDIO_TOKEN_PLACEHOLDER + x = merge_flat_embeddings( + text_embeddings=x, + multimodal_embeddings=soft_audio_embeddings, + mask=audio_mask, + ) + sliding_attention_mask = None if ( is_prefill @@ -1518,7 +1582,7 @@ def __call__( and images is not None and attention_mask is not None ): - bidirectional_mask = tokens == TOKEN_PLACEHOLDER + bidirectional_mask = tokens == VISION_TOKEN_PLACEHOLDER sliding_attention_mask = _add_bidirectional_mask( attention_mask, bidirectional_mask ) @@ -1655,6 +1719,70 @@ def _encode_vision(self, vision_input: PreprocessedVisionInput): all_tokens = all_tokens[:, 0, :, :] return all_tokens + def _encode_audio(self, audio_input: PreprocessedAudioInput): + """Encode audio into the same space as the text embeddings.""" + assert self.audio_encoder is not None + + batch_size = audio_input.audio.shape[0] + + if len(audio_input.soft_token_counts) > 0 and isinstance( + audio_input.soft_token_counts[0], int + ): + soft_token_counts = (audio_input.soft_token_counts,) + else: + soft_token_counts = audio_input.soft_token_counts + + max_n_audios = max((len(counts) for counts in soft_token_counts), default=0) + if max_n_audios == 0: + return jnp.zeros((batch_size, 0, self.config.embed_dim)) + + audio = audio_input.audio + audio_lengths = audio_input.audio_lengths + + time_steps = audio.shape[2] + audio_flat = jnp.reshape(audio, (batch_size * max_n_audios, time_steps)) + audio_lengths_flat = jnp.reshape( + audio_lengths, (batch_size * max_n_audios,) + ) + + # Run audio encoder + # embeddings: [batch * max_n_audios, reduced_time, lm_model_dims] + # mask: [batch * max_n_audios, reduced_time] (True for valid) + embeddings, mask = self.audio_encoder(audio_flat, audio_lengths_flat) + + batch_tokens = [] + max_tokens_per_batch = 0 + for b in range(batch_size): + per_audio_tokens = [] + counts = soft_token_counts[b] if b < len(soft_token_counts) else () + for i in range(len(counts)): + idx = b * max_n_audios + i + expected_count = counts[i] + if mask is not None: + valid_indices = jnp.nonzero(mask[idx], size=expected_count)[0] + real_tokens = embeddings[idx][valid_indices] + else: + real_tokens = embeddings[idx][:expected_count] + per_audio_tokens.append(real_tokens) + + if per_audio_tokens: + b_tokens = jnp.concatenate(per_audio_tokens, axis=0) + else: + b_tokens = jnp.zeros((0, embeddings.shape[-1])) + batch_tokens.append(b_tokens) + max_tokens_per_batch = max(max_tokens_per_batch, b_tokens.shape[0]) + + padded_batch_tokens = [] + for b_tokens in batch_tokens: + pad_len = max_tokens_per_batch - b_tokens.shape[0] + if pad_len > 0: + b_tokens = jnp.pad(b_tokens, ((0, pad_len), (0, 0))) + padded_batch_tokens.append(b_tokens) + + all_tokens = jnp.stack(padded_batch_tokens, axis=0) + all_tokens = self.embedder.encode_audio(all_tokens) + return all_tokens + def compute_final_logits( self, x: jaxtyping.Array, diff --git a/tunix/models/gemma4/params_safetensors.py b/tunix/models/gemma4/params_safetensors.py index bee68d14c..4a3487ef8 100644 --- a/tunix/models/gemma4/params_safetensors.py +++ b/tunix/models/gemma4/params_safetensors.py @@ -426,6 +426,146 @@ def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig): ), }) + if cfg.audio_encoder is not None: + mapping.update({ + # Audio projector + r"model\.embed_audio\.embedding_projection\.weight": ( + "embedder.mm_audio_input_projection.w", + ((1, 0), None), + ), + # Audio subsampling + r"model\.audio_tower\.subsample_conv_projection\.layer0\.conv\.weight": ( + "audio_encoder.subsampling.conv0.kernel", + ((2, 3, 1, 0), None), + ), + r"model\.audio_tower\.subsample_conv_projection\.layer0\.norm\.weight": ( + "audio_encoder.subsampling.norm0.scale", + None, + ), + r"model\.audio_tower\.subsample_conv_projection\.layer1\.conv\.weight": ( + "audio_encoder.subsampling.conv1.kernel", + ((2, 3, 1, 0), None), + ), + r"model\.audio_tower\.subsample_conv_projection\.layer1\.norm\.weight": ( + "audio_encoder.subsampling.norm1.scale", + None, + ), + r"model\.audio_tower\.subsample_conv_projection\.input_proj_linear\.weight": ( + "audio_encoder.subsampling.input_proj.kernel", + ((1, 0), None), + ), + # Conformer Blocks + # FFN Blocks + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward1\.ffw_layer_1\.linear\.weight": ( + r"audio_encoder.blocks.\1.fflayer_start.ffn_layer1.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward1\.ffw_layer_2\.linear\.weight": ( + r"audio_encoder.blocks.\1.fflayer_start.ffn_layer2.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward1\.pre_layer_norm\.weight": ( + r"audio_encoder.blocks.\1.fflayer_start.pre_layer_norm.scale", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward1\.post_layer_norm\.weight": ( + r"audio_encoder.blocks.\1.fflayer_start.post_layer_norm.scale", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward2\.ffw_layer_1\.linear\.weight": ( + r"audio_encoder.blocks.\1.fflayer_end.ffn_layer1.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward2\.ffw_layer_2\.linear\.weight": ( + r"audio_encoder.blocks.\1.fflayer_end.ffn_layer2.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward2\.pre_layer_norm\.weight": ( + r"audio_encoder.blocks.\1.fflayer_end.pre_layer_norm.scale", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.feed_forward2\.post_layer_norm\.weight": ( + r"audio_encoder.blocks.\1.fflayer_end.post_layer_norm.scale", + None, + ), + # Attention Block + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.q_proj\.linear\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.self_atten.query.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.k_proj\.linear\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.self_atten.key.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.v_proj\.linear\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.self_atten.value.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.per_dim_scale": ( + r"audio_encoder.blocks.\1.trans_atten.self_atten.per_dim_scale.value", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.relative_k_proj\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.self_atten.relative_position_embedding.pos_proj.kernel", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.self_attn\.post\.linear\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.post.w", + ( + (1, 0), + ( + cfg.audio_encoder.atten_num_heads, + cfg.audio_encoder.model_dims + // cfg.audio_encoder.atten_num_heads, + cfg.audio_encoder.model_dims, + ), + ), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.norm_pre_attn\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.pre_norm.scale", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.norm_post_attn\.weight": ( + r"audio_encoder.blocks.\1.trans_atten.post_norm.scale", + None, + ), + # Lightweight Conv Block + r"model\.audio_tower\.layers\.([0-9]+)\.lconv1d\.linear_start\.linear\.weight": ( + r"audio_encoder.blocks.\1.lconv.linear_start.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.lconv1d\.depthwise_conv1d\.weight": ( + r"audio_encoder.blocks.\1.lconv.depthwise_conv1d.kernel", + ((2, 1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.lconv1d\.linear_end\.linear\.weight": ( + r"audio_encoder.blocks.\1.lconv.linear_end.w", + ((1, 0), None), + ), + r"model\.audio_tower\.layers\.([0-9]+)\.lconv1d\.pre_layer_norm\.weight": ( + r"audio_encoder.blocks.\1.lconv.ln.scale", + None, + ), + r"model\.audio_tower\.layers\.([0-9]+)\.lconv1d\.conv_norm\.weight": ( + r"audio_encoder.blocks.\1.lconv.conv_norm.scale", + None, + ), + # Final Norm + r"model\.audio_tower\.layers\.([0-9]+)\.norm_out\.weight": ( + r"audio_encoder.blocks.\1.final_ln.scale", + None, + ), + # Output Projection + r"model\.audio_tower\.output_proj\.weight": ( + "audio_encoder.output_projection.kernel", + ((1, 0), None), + ), + r"model\.audio_tower\.output_proj\.bias": ( + "audio_encoder.output_projection.bias", + None, + ), + }) + return mapping diff --git a/tunix/models/gemma4/sampling_example.ipynb b/tunix/models/gemma4/sampling_example.ipynb index fa0c666ca..3bf802a5b 100644 --- a/tunix/models/gemma4/sampling_example.ipynb +++ b/tunix/models/gemma4/sampling_example.ipynb @@ -1,19 +1,16 @@ { "cells": [ { + "id": "igDbPR6Pb-Q7", "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", - "from transformers import AutoTokenizer, AutoProcessor\n", + "from transformers import AutoProcessor, AutoTokenizer\n", "from tunix.generate import sampler\n", "from tunix.models.gemma4 import model\n", "from tunix.models.gemma4 import params_safetensors\n", "\n", - "\n", "MESH = [(1, 2), (\"fsdp\", \"tp\")]\n", "mesh = jax.make_mesh(\n", " *MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0])\n", @@ -69,10 +66,11 @@ "else:\n", " raise ValueError(f\"Unknown version: {version}\")" ], - "id": "igDbPR6Pb-Q7" + "metadata": {}, + "execution_count": null }, { - "metadata": {}, + "id": "iAtteAFzcEfX", "cell_type": "code", "source": [ "def templatize(prompts):\n", @@ -116,20 +114,19 @@ " print(t)\n", " print(\"*\" * 30)" ], - "outputs": [], - "execution_count": null, - "id": "iAtteAFzcEfX" + "metadata": {}, + "execution_count": null }, { - "metadata": {}, "id": "Sf0zF5FJmS8L", "cell_type": "markdown", "source": [ "## Vision example" - ] + ], + "metadata": {}, + "execution_count": null }, { - "metadata": {}, "id": "AL2VTm40mVzA", "cell_type": "markdown", "source": [ @@ -139,72 +136,157 @@ "image1 = ds[0]['image']\n", "image2 = ds[1]['image']\n", "image3 = ds[2]['image']" - ] + ], + "metadata": {}, + "execution_count": null }, { - "metadata": {}, "id": "24Kyja84mZo5", "cell_type": "code", "source": [ "import matplotlib.pyplot as plt\n", + "\n", "plt.imshow(image1)" ], - "outputs": [], + "metadata": {}, "execution_count": null }, { - "metadata": {}, "id": "qfxgoYV4mbup", "cell_type": "code", "source": [ "plt.imshow(image2)" ], - "outputs": [], + "metadata": {}, "execution_count": null }, { - "metadata": {}, "id": "fLhqwBwlmiz4", "cell_type": "code", "source": [ - "messages = [[\n", - " {\n", + "messages = [\n", + " [{\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\"type\": \"text\", \"text\": \"Describe the images.\"},\n", " {\"type\": \"image\"},\n", " {\"type\": \"image\"},\n", - " ]\n", + " ],\n", " }],\n", " [{\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\"type\": \"text\", \"text\": \"Write a poem for the image.\"},\n", " {\"type\": \"image\"},\n", - " ]\n", - " }]\n", + " ],\n", + " }],\n", "]\n", "\n", "prompt = processor.apply_chat_template(\n", - " messages, \n", - " tokenize=False, \n", - " add_generation_prompt=True\n", + " messages, tokenize=False, add_generation_prompt=True\n", ")\n", "\n", "with mesh:\n", - " output = sampler(\n", - " input_strings=prompt,\n", - " images=[[image1, image2], [image3]],\n", - " max_generation_steps=128,\n", - " temperature=1,\n", - " echo=False,\n", - " eos_tokens=[1, 106, 50],\n", - " )\n", + " output = sampler(\n", + " input_strings=prompt,\n", + " images=[[image1, image2], [image3]],\n", + " max_generation_steps=128,\n", + " temperature=1,\n", + " echo=False,\n", + " eos_tokens=[1, 106, 50],\n", + " )\n", "for o in output.text:\n", - " print(o)\n", - " print(\"####\" * 8)" + " print(o)\n", + " print(\"####\" * 8)" + ], + "metadata": {}, + "execution_count": null + }, + { + "id": "ca0a6bdf", + "cell_type": "markdown", + "source": [ + "## Audio example" + ], + "metadata": {}, + "execution_count": null + }, + { + "id": "acc30e52", + "cell_type": "code", + "source": [ + "import librosa\n", + "import numpy as np\n", + "\n", + "\n", + "def load_audio(path):\n", + " try:\n", + " # librosa automatically resamples to 16kHz and converts to float32 [-1.0, 1.0]\n", + " y, _ = librosa.load(path, sr=16000)\n", + " return y\n", + " except Exception as e:\n", + " print(f\"Failed to load with librosa, trying scipy: {e}\")\n", + " from scipy.io import wavfile\n", + "\n", + " sr, y = wavfile.read(path)\n", + " if y.ndim > 1:\n", + " y = y.mean(axis=-1)\n", + " if y.dtype == np.int16:\n", + " y = y.astype(np.float32) / 32768.0\n", + " elif y.dtype == np.int32:\n", + " y = y.astype(np.float32) / 2147483648.0\n", + " # Note: simple scipy loading does not resample. Ensure your wav is 16kHz.\n", + " return y\n", + "\n", + "\n", + "# Example usage with real files (uncomment and replace with your paths):\n", + "# audio1 = load_audio(\"path/to/audio1.wav\")\n", + "# audio2 = load_audio(\"path/to/audio2.wav\")\n", + "# audio3 = load_audio(\"path/to/audio3.wav\")\n", + "\n", + "# Fallback to mock audio (zeros) for quick verification:\n", + "audio1 = jnp.zeros((16000 * 2,), dtype=jnp.float32) # 2 seconds of silence\n", + "audio2 = jnp.zeros((16000 * 1,), dtype=jnp.float32) # 1 second of silence\n", + "audio3 = jnp.zeros((16000 * 3,), dtype=jnp.float32) # 3 seconds of silence\n", + "\n", + "# Define messages with audio placeholders\n", + "audio_messages = [\n", + " [{\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"Describe these audio clips.\"},\n", + " {\"type\": \"audio\"},\n", + " {\"type\": \"audio\"},\n", + " ],\n", + " }],\n", + " [{\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"What is in this audio?\"},\n", + " {\"type\": \"audio\"},\n", + " ],\n", + " }],\n", + "]\n", + "\n", + "audio_prompt = processor.apply_chat_template(\n", + " audio_messages, tokenize=False, add_generation_prompt=True\n", + ")\n", + "\n", + "with mesh:\n", + " audio_output = sampler(\n", + " input_strings=audio_prompt,\n", + " audios=[[audio1, audio2], [audio3]],\n", + " max_generation_steps=128,\n", + " temperature=1,\n", + " echo=False,\n", + " eos_tokens=[1, 106, 50],\n", + " )\n", + "\n", + "for o in audio_output.text:\n", + " print(o)\n", + " print(\"####\" * 8)" ], - "outputs": [], + "metadata": {}, "execution_count": null } ], @@ -225,6 +307,6 @@ "display_name": "Python 3" } }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat_minor": 5, + "nbformat": 4 } diff --git a/tunix/processors/audio_processor.py b/tunix/processors/audio_processor.py new file mode 100644 index 000000000..9b2b4de60 --- /dev/null +++ b/tunix/processors/audio_processor.py @@ -0,0 +1,193 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Audio processing for Gemma4.""" + +from typing import Any +import jax.numpy as jnp +import numpy as np +from tunix.models.gemma4 import model as model_lib + +POSITIONS_PAD_VALUE = -1 + + +def compute_soft_token_count( + length: int, + sample_rate: int = 16000, + max_audio_seq_len: int = 750, +) -> int: + """Computes the number of soft tokens for a given audio length.""" + frame_length = int(round(sample_rate * 20.0 / 1000.0)) + hop_length = int(round(sample_rate * 10.0 / 1000.0)) + frame_size_for_unfold = frame_length + 1 + num_mel_frames = (length - frame_size_for_unfold) // hop_length + 1 + t = num_mel_frames + for _ in range(2): + t_padded = t + 2 + t = (t_padded - 3) // 2 + 1 + return min(t, max_audio_seq_len) + + +def add_variable_extra_tokens_for_audio( + tokens: np.ndarray, + *, + soft_token_counts: list[int] | tuple[tuple[int, ...], ...], + placeholder_token: int = 258881, # <|audio|> + start_token: int = 256000, # <|audio> + end_token: int = 258883, # +) -> np.ndarray: + """Expand `AUDIO_PLACEHOLDER` with a variable number of placeholders.""" + batch_size = tokens.shape[0] + results = [] + for b in range(batch_size): + row = tokens[b].tolist() + expanded = [] + audio_idx = 0 + + if len(soft_token_counts) > 0 and isinstance(soft_token_counts[0], int): + counts = soft_token_counts + else: + counts = soft_token_counts[b] if b < len(soft_token_counts) else () + + for token in row: + if token == placeholder_token and audio_idx < len(counts): + count = counts[audio_idx] + expanded.append(start_token) + expanded.extend([model_lib.AUDIO_TOKEN_PLACEHOLDER] * count) + expanded.append(end_token) + audio_idx += 1 + else: + expanded.append(token) + results.append(expanded) + + max_len = max(len(r) for r in results) + padded = np.zeros((batch_size, max_len), dtype=np.int32) + for b, row in enumerate(results): + padded[b, : len(row)] = row + + return padded + + +def process_gemma4_audio_inputs( + audio: Any, + tokens: list[np.ndarray], + audio_encoder_config: Any, + pad_id: int, +) -> tuple[model_lib.PreprocessedAudioInput, list[np.ndarray]]: + """Processes audio and tokens for Gemma4 multimodal models.""" + + # Normalize audio input to list[list[np.ndarray]] (batch, clips) + if not isinstance(audio, list): + audio = [[audio]] + elif len(audio) > 0 and not isinstance(audio[0], list): + audio = [[clip] for clip in audio] + + max_n_clips = max((len(batch) for batch in audio), default=0) + + batch_audio = [] + batch_lengths = [] + all_soft_token_counts = [] + max_samples = 0 + + sample_rate = getattr(audio_encoder_config, "sample_rate", 16000) + max_audio_seq_len = getattr(audio_encoder_config, "audio_seq_length", 750) + + for batch in audio: + if not batch: + batch_audio.append([]) + batch_lengths.append([]) + all_soft_token_counts.append(()) + continue + + clip_audios = [] + clip_lengths = [] + clip_soft_token_counts = [] + for clip in batch: + clip = np.asarray(clip, dtype=np.float32) + clip_audios.append(clip) + clip_lengths.append(len(clip)) + clip_soft_token_counts.append( + compute_soft_token_count(len(clip), sample_rate, max_audio_seq_len) + ) + max_samples = max(max_samples, len(clip)) + + batch_audio.append(clip_audios) + batch_lengths.append(clip_lengths) + all_soft_token_counts.append(tuple(clip_soft_token_counts)) + + if max_samples == 0: + max_samples = 16000 + + final_audio = [] + final_lengths = [] + + for b_idx in range(len(audio)): + clips = batch_audio[b_idx] + lengths = batch_lengths[b_idx] + + padded_clips = [] + padded_lengths = [] + + for clip_idx in range(len(clips)): + clip = clips[clip_idx] + length = lengths[clip_idx] + pad_len = max_samples - len(clip) + if pad_len > 0: + clip = np.pad(clip, (0, pad_len)) + padded_clips.append(clip) + padded_lengths.append(length) + + n_pad = max_n_clips - len(clips) + for _ in range(n_pad): + padded_clips.append(np.zeros(max_samples, dtype=np.float32)) + padded_lengths.append(0) + + if padded_clips: + final_audio.append(np.stack(padded_clips, axis=0)) + final_lengths.append(np.array(padded_lengths, dtype=np.int32)) + else: + final_audio.append(np.zeros((max_n_clips, max_samples), dtype=np.float32)) + final_lengths.append(np.zeros(max_n_clips, dtype=np.int32)) + + if final_audio: + audio_tensor = jnp.stack(final_audio, axis=0) + lengths_tensor = jnp.stack(final_lengths, axis=0) + else: + batches = len(audio) + audio_tensor = jnp.zeros( + (batches, max_n_clips, max_samples), dtype=jnp.float32 + ) + lengths_tensor = jnp.zeros((batches, max_n_clips), dtype=jnp.int32) + + processed_audio = model_lib.PreprocessedAudioInput( + audio=audio_tensor, + audio_lengths=lengths_tensor, + soft_token_counts=tuple(all_soft_token_counts), + ) + + if all_soft_token_counts: + max_len = max(len(t) for t in tokens) + padded_tokens = np.array([ + np.pad(x, (0, max_len - len(x)), constant_values=pad_id) for x in tokens + ]) + expanded_tokens = add_variable_extra_tokens_for_audio( + padded_tokens, + soft_token_counts=tuple(all_soft_token_counts), + ) + tokens = [ + np.array([tid for tid in row if tid != pad_id]) + for row in expanded_tokens.tolist() + ] + + return processed_audio, tokens diff --git a/tunix/processors/image_processor.py b/tunix/processors/image_processor.py index ba5529ecd..70452d625 100644 --- a/tunix/processors/image_processor.py +++ b/tunix/processors/image_processor.py @@ -17,7 +17,7 @@ from typing import Any import numpy as np from PIL import Image -from tunix.models.gemma4.model import PreprocessedVisionInput +from tunix.models.gemma4 import model as model_lib class ImageProcessor: @@ -424,7 +424,7 @@ def add_variable_extra_tokens_for_images( ) -> np.ndarray: """Expand placeholder token with a variable number of placeholders.""" double_new_line_token = 108 - soft_token_placeholder = -2 # img soft tokens + soft_token_placeholder = model_lib.VISION_TOKEN_PLACEHOLDER # img soft tokens batch_size = tokens.shape[0] results = [] @@ -553,7 +553,7 @@ def process_gemma4_inputs( dtype=jnp.int32, ) - processed_images = PreprocessedVisionInput( + processed_images = model_lib.PreprocessedVisionInput( patches=patches, positions_xy=positions_xy, soft_token_counts=tuple(all_soft_token_counts),