Replaces #16334, which GitHub closed automatically after my DCO sign-off push accidentally rewrote the commit without its parent. Same change, now with the sign-off; no code changes.

Summary

PositionalEmbedding in nemo/collections/tts/modules/transformer.py (used by FFTransformerEncoder/FFTransformerDecoder, e.g. via nemo/collections/tts/modules/adapters.py and nemo/collections/tts/models/fastpitch_ssl.py) builds its inv_freq buffer with:

inv_freq = 1 / (10000 ** (torch.arange(0.0, demb, 2.0) / demb))

with no explicit dtype. PyTorch Lightning's bf16-true precision setting calls torch.set_default_dtype(torch.bfloat16) for the duration of model construction — this is the exact hazard already documented in this repo's own nemo/collections/speechlm2/parts/precision.py (fp32_precision() docstring: "In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models.").

Under that precision mode, torch.arange(0.0, demb, 2.0) (and the division/power that follows) runs and gets stored in bfloat16 instead of float32, silently losing precision in the sinusoidal/RoPE-style positional frequencies — no crash, no warning, just a quieter/less precise positional signal.

I reproduced this directly against the unmodified module: constructing PositionalEmbedding under torch.set_default_dtype(torch.bfloat16) (matching what PTL's bf16-true does) produces an inv_freq buffer in bfloat16 with up to ~0.67% relative error vs. the float32 reference.

Fix

• inv_freq is now always built with dtype=torch.float32, independent of the ambient default dtype.
• forward() upcasts pos_seq to float32 for the sinusoid computation and casts the result back to the caller's original dtype before returning — this keeps the output dtype unchanged for both call sites in the same file (both of which already do pos_seq.to(inp.dtype) before calling in), and avoids a dtype-mismatch in the matmul that would otherwise occur between a half-precision pos_seq and a now-always-float32 inv_freq.

Test plan

• Reproduced the bug against the real, unmodified PositionalEmbedding class (buffer dtype + numeric drift vs. float32 reference).
• Applied the fix, confirmed inv_freq is float32 with zero drift under the same bf16-default-dtype construction.
• Exercised forward() end-to-end with both bfloat16 and float32 pos_seq inputs, confirming matching output dtype and no crash.
• Added tests/collections/tts/modules/test_transformer_posemb_dtype.py covering both the buffer dtype and the numeric match against a float32 reference under a bf16 default dtype.

🤖 Generated with Claude Code