Add FP8 support for SALMAutomodel

[pzelasko]

Summary

• add SpeechLM2 FP8 helpers for TorchAO config creation, TE FP8 autocast/patching, FSDP scale precompute, and TE FP8 padding/alignment
• wire SALMAutomodel through those helpers without cached FP8 state and keep the forward/backward flow minimal
• support TE FP8 padding for BSHD and THD packed inputs, including context-parallel alignment metadata
• document TE FP8 and TorchAO FP8 config examples and add focused unit/runtime coverage

Testing

• git diff --cached --check
• pytest tests/collections/speechlm2/test_fp8.py tests/collections/speechlm2/test_salm_packed_sequences.py -q
• in nemo-speech:cu13-h100plus: pytest tests/collections/speechlm2/test_salm_automodel.py -q
• in nemo-speech:cu13-h100plus: 2-GPU LibriSpeech training with Nemotron 3 Nano + Canary-1B-Flash, FSDP2 dp=2 ep=2, TE FP8 recipe=block, completed 2000/2000 steps in 1:18:02 with final logged loss 0.0586

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[pzelasko]

/ok to test 614e4b4

[pzelasko]

/ok to test 23761ed

[pzelasko]

/ok to test 329f23a

[KunalDhawan]

Great work @pzelasko, added minor comments below, other than that LGTM

[KunalDhawan]

This only ensures B*T % 8 == 0 and ignores tp_size, which is inconsistent with the THD helper (8 * cp_size * tp_size). Under BSHD + TP + TE-FP8 I think that breaks in two ways:

• prepare_inputs truncates the seq dim to a multiple of tp_size so sequence parallelism doesn't silently reshape the input (salm_automodel.py ~L269), but then forward appends pad = (-T) % seq_multiple tokens, so the padded length is no longer guaranteed divisible by tp_size → SP shape break.
• With SP the local TE Linear sees M = B*T/tp_size, so FP8 actually needs B*T % (8*tp_size) == 0, not just % 8.

Could we either thread tp_size through here (note 8*tp_size alone isn't enough — e.g. B=16, tp=4 → multiple of 2, still not divisible by 4 — so probably needs an explicit lcm(tp_size, ...)), or add a validate_fp8_config rejection for BSHD + TP + TE-FP8 pointing folks at the THD packed path? A BSHD analogue of test_maybe_pad_thd_..._accounts_for_cp_and_tp would lock it down. This combo wasn't in the 2-GPU run (dp=2 ep=2, no TP), so it's currently untested.

[pzelasko]

nice catch

[KunalDhawan]

Quick question on wrapping backward in te_fp8_context too: standard TE usage only wraps the forward, and the backward consumes the FP8 metadata captured during the forward's fp8_autocast. Re-entering fp8_autocast here can, for history/delayed-scaling recipes, trigger an extra amax/scale update (and a second amax all-reduce) at context exit. Probably harmless for block/current (which is what the run used), but it's an easy source of subtle scale drift on other recipes. Is it deliberate / needed for something specific? If not, dropping it from backward seems safer.

[pzelasko]

nice catch too

[pzelasko]

/ok to test 17e203f

[MahmoudAshraf97]

Have you gotten any speedups from this? I tried FP8 with Fastconformer and it was slower than BF16

[pzelasko]

I'm still debugging a few issues before this is ready, I'll report when I have the numbers.

The expected speedup from TransformerEngine's FP8 is about 2x on Hopper, but you have to make sure the training is compute bound to get that speedup. If your matmuls are on very small problem sizes, the overhead of scaling etc can be greater than the speedup of an already very tiny kernel. Two easiest ways to get the speedup are using larger models or larger batch sizes.