Add bfloat16 support for Parakeet

[larryliu0820]

Enable both export script and runner to support bfloat16. Changing CI to run on bfloat16 by default.

Also added a parakeet-cuda-debug mode

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16821

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 39 Pending, 1 Unrelated Failure

As of commit 578aea9 with merge base b928496 ():

NEW FAILURE - The following job has failed:

• trunk / test-arm-backend-zephyr / linux-job (gh)
  RuntimeError: Command docker exec -t 413bcb840885a733c7a66441b08f682ceefa29a16ec94ec303842fbd75b4f818 /exec failed with exit code 127

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-models-linux (linear, portable, linux.2xlarge) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Gasoonjia]

Thx man for such efficient support!

[Gasoonjia]

im a little bit lost in the logic here: based on my understanding, greedy_decode_executorch should only for joint and decoder, but not encoder. and encoder takes preprocessor's output as input, which is always fp32 now. Since we exported encoder as target dtype (e.g. bf16), how do we makes bf16 encoder consume fp32 preprocessor's outpu?

[larryliu0820]

and encoder takes preprocessor's output as input, which is always fp32 now

Not true, if we pass dtype == bfloat16, preprocessor will take in a float and gives bfloat16 result

[Gasoonjia]

hmm not sure if i misunderstood something, but export script says:
# Preprocessor always uses float32 - runner converts output to encoder's dtype
but i didn't find the code in runner for the type conversion.
if preprocessor will give bfloat16 output directly:

1. perhapes updating the export doc?
2. how does the preprocessor know that? Seems the export configs for preprocessor are the same across different dtype.

[larryliu0820]

Let me dive a bit deeper into the model. I can take a look at the graph to see if it changes the dtype of the output

[larryliu0820]

Graph:

graph():
    %b_preprocessor_dtype_sentinel_tensor : [num_users=0] = placeholder[target=b_preprocessor_dtype_sentinel_tensor]
    %b_preprocessor_featurizer_window : [num_users=1] = placeholder[target=b_preprocessor_featurizer_window]
    %b_preprocessor_featurizer_fb : [num_users=1] = placeholder[target=b_preprocessor_featurizer_fb]
    %audio : [num_users=2] = placeholder[target=audio]
    %length : [num_users=1] = placeholder[target=length]
    %sym_size_int_4 : [num_users=1] = call_function[target=torch.ops.aten.sym_size.int](args = (%audio, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%audio, 0), kwargs = {})
    %submod_5 : [num_users=1] = get_attr[target=submod_1]
    %to : [num_users=1] = call_function[target=torch.ops.higher_order.wrap_with_set_grad_enabled](args = (False, %submod_5, %unsqueeze), kwargs = {})
    %getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%to, 0), kwargs = {})
    %submod_6 : [num_users=1] = get_attr[target=submod_2]
    %wrap_with_set_grad_enabled : [num_users=2] = call_function[target=torch.ops.higher_order.wrap_with_set_grad_enabled](args = (False, %submod_6, %length, %sym_size_int_4, %getitem_4, %b_preprocessor_featurizer_window, %b_preprocessor_featurizer_fb), kwargs = {})
    %masked_fill_2 : [num_users=1] = call_function[target=operator.getitem](args = (%wrap_with_set_grad_enabled, 0), kwargs = {})
    %where : [num_users=1] = call_function[target=operator.getitem](args = (%wrap_with_set_grad_enabled, 1), kwargs = {})
    %submod_7 : [num_users=1] = get_attr[target=submod_3]
    %to_6 : [num_users=1] = call_function[target=torch.ops.higher_order.wrap_with_set_grad_enabled](args = (False, %submod_7, %masked_fill_2), kwargs = {})
    %getitem_5 : [num_users=1] = call_function[target=operator.getitem](args = (%to_6, 0), kwargs = {})
    return (getitem_5, where)

Looking at the graph, the key is in those placeholder buffers:

• b_preprocessor_featurizer_window - the window function (e.g., Hann window)
• b_preprocessor_featurizer_fb - the mel filterbank matrix

When we call model.to(torch.bfloat16), it converts all parameters and buffers in the model, including the preprocessor's window and filterbank tensors.

The mel spectrogram computation does something like:

# Simplified view of what happens inside:
windowed = audio_frame * window          # window is bf16 → result is bf16
spectrum = torch.fft.rfft(windowed)         # stays bf16
mel = spectrum @ filterbank                      # filterbank is bf16 → result is bf16

PyTorch's type promotion rules mean that when our bf16 audio interacts with bf16 buffers, the output stays bf16.

[mergennachin]

@manuelcandales @larryliu0820 it would be good to test bf16 for metal too

[Gasoonjia]

Thanks for detailed update! LGTM now!
nit: mind changing PR title to remove float16?

[Gasoonjia]

Maybe remove this comment? looks like torch.export handle the dtype changing instead of runner,.