diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 4151a3824dfd..8adbf8fa20e7 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -147,7 +147,11 @@ def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu waveform += self.dither * torch.randn(waveform.shape, dtype=waveform.dtype, device=waveform.device) stft = torch.stft(waveform, self.n_fft, self.hop_length, window=window, return_complex=True) - magnitudes = stft[..., :-1].abs() ** 2 + # `stft[..., :-1]` is a non-contiguous view; on some CPU backends the + # downstream `mel_filters.T @ magnitudes` matmul hits a slow strided + # path (observed ~8x slower). Forcing contiguity here is cheap and + # keeps the matmul on the fast path. + magnitudes = (stft[..., :-1].abs() ** 2).contiguous() mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32) mel_spec = mel_filters.T @ magnitudes diff --git a/tests/models/whisper/test_feature_extraction_whisper.py b/tests/models/whisper/test_feature_extraction_whisper.py index 65f150cc1295..f7a9d590619e 100644 --- a/tests/models/whisper/test_feature_extraction_whisper.py +++ b/tests/models/whisper/test_feature_extraction_whisper.py @@ -268,6 +268,45 @@ def test_double_precision_pad(self): pt_processed = feature_extractor.pad([{"input_features": inputs}], return_tensors="pt") self.assertTrue(pt_processed.input_features.dtype == torch.float32) + @require_torch + def test_torch_extract_fbank_features_contiguous_magnitudes(self): + """Guards the mel-magnitude contiguity fix. + + `stft[..., :-1]` is a non-contiguous view, and squaring it keeps that + layout. On some CPU backends the downstream `mel_filters.T @ magnitudes` + matmul then falls onto a slow strided-GEMM path (observed ~8x slower on + a ROCm PyTorch build, ~43 ms vs ~3 ms for a 30 s window). Forcing + `magnitudes` contiguous restores the fast path with no change in output. + + Correctness (and the non-contiguity of the raw view) is asserted so this + stays stable in CI. See https://github.com/huggingface/transformers/pull/47351 + for the benchmark script and the measured (backend-dependent) speedups. + """ + torch.manual_seed(0) + feature_extractor = WhisperFeatureExtractor() + n_fft = feature_extractor.n_fft + hop_length = feature_extractor.hop_length + # Mirror `_torch_extract_fbank_features`: the window is built from n_fft. + window = torch.hann_window(n_fft) + # One full Whisper window of audio (chunk_length * sampling_rate samples). + waveform = torch.randn(feature_extractor.n_samples, dtype=torch.float32) + + stft = torch.stft(waveform, n_fft, hop_length, window=window, return_complex=True) + magnitudes_non_contiguous = stft[..., :-1].abs() ** 2 + magnitudes_contiguous = magnitudes_non_contiguous.contiguous() + + # Root cause of the slow path: the sliced view is non-contiguous. + self.assertFalse(magnitudes_non_contiguous.is_contiguous()) + self.assertTrue(magnitudes_contiguous.is_contiguous()) + + mel_filters = torch.from_numpy(feature_extractor.mel_filters).to(torch.float32) + + # The fix must not change the result, only the memory layout. + torch.testing.assert_close( + mel_filters.T @ magnitudes_non_contiguous, + mel_filters.T @ magnitudes_contiguous, + ) + def _load_datasamples(self, num_samples): ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") # automatic decoding with librispeech