How to try with FP8 for LLM finetuning?

[kailashg26]

What would be the steps to use FP8 from TE (transformer engine) in TorchTune to support instead of BF16 for full finetuning and LoRA?

[ebsmothers]

Hi @kailashg26 I can't speak for TE, but we are working to enable fp8 training for full finetunes in #2546

[kailashg26]

I see, thanks! @ebsmothers Can I try the fp8 for full finetuning on the FP8 branch of torchtune? Would that work on torch-2.8?
Also, could you please let me know what are the steps normally to enable the fp8 support? and do you know if there will be any differences in the performance form FP8 (the one that is integrated) and the TE version of FP8?

[ebsmothers]

@kailashg26 yes that should work. The one caveat is that we do not yet support tensor parallelism (but hopefully coming soon). Regarding performance comparisons against TE, I think @vkuzo may have the most knowledge here.

[kailashg26]

@vkuzo @felipemello1 any suggestions on how to proceed integrating FP8 form TE into torchtune?

[joecummings]

@kailashg26 FP8 support has landed in our nightlies if you wanna check it out!

[kailashg26]

I see, thanks @joecummings. Do you know if torchtune supports FP8 from TE? I guess the current one is from torchao and uses triton backend. Please correct me if I'm wrong.

[joecummings]

I see, thanks @joecummings. Do you know if torchtune supports FP8 from TE? I guess the current one is from torchao and uses triton backend. Please correct me if I'm wrong.

Nope, at the moment this is a version from torchao, but if you prefer working with TE, we'd be happy to showcase a recipe you write using the engine! There shouldn't be an large blockers to doing so.

[kailashg26]

Sure, I'd appreciate that if you could showcase a recipe. @joecummings

[kailashg26]

@joecummings any help/suggestions regarding this issue?