How to reduce cuda memory usage of fp8?

[XLzed]

The fp8 implementation of TransformerEngine seems to cache the transposed fp8 inputs to get better perfomance (#1261), which bring same memory usage as bf16. Is there any way to disable this behavior? I want to use the dualPipe/dualPipeV mentioned by deepseek, but its high video memory overhead will lead to OOM. So I wanted to use TransformerEngine fp8 to reduce the activation memory usage, but I didn't obtain the expected benefits.

[zhongbozhu]

If you use MegatronLM/Nemo, there is a --fp8-param-gather configs that directly cast fp32 master weight in distributed optimizer into FP8 weight, so each layer in TE doesn't store the high precision bf16 copy.

TE provides this api to Megatron to cast master weight into fp8:

TransformerEngine/transformer_engine/pytorch/tensor/utils.py

Line 47 in d2973cb

def cast_master_weights_to_fp8(

[dxqbYD]

I think what you mean is this: #1880
It is currently not supported I think.

[timmoon10]

Are you are sure the high memory usage is from activations? After the linear layer forward pass, we need to cache the input tensor so that we can use it in the backward pass. Currently we cache the FP8 data (on Blackwell) or the FP8 transpose data (on Hopper), so I expect we'd see a 50% reduction in the activation memory usage compared to PyTorch. If you are not seeing much reduction in activation memory usage, could it be because you have many non-FP8 layers (layernorm, activations, etc) or because you have a very small number of layers?

As @dxqbYD mentioned, I think it's more likely the high memory usage is from weights. Instead of storing weights in BF16, you may want to store in FP8 by initializing the model within an fp8_model_init context (see #1880 (comment)). However, unless you're doing something like LoRA, keeping weights in FP8 has implications for convergence. Currently, the distributed optimizer in Megatron-LM is your best bet for avoiding these numerical issues (it applies Adam to sharded FP32 weights, and then does an all-gather in FP8).

[XLzed]

Are you are sure the high memory usage is from activations? After the linear layer forward pass, we need to cache the input tensor so that we can use it in the backward pass. Currently we cache the FP8 data (on Blackwell) or the FP8 transpose data (on Hopper), so I expect we'd see a 50% reduction in the activation memory usage compared to PyTorch. If you are not seeing much reduction in activation memory usage, could it be because you have many non-FP8 layers (layernorm, activations, etc) or because you have a very small number of layers?

As @dxqbYD mentioned, I think it's more likely the high memory usage is from weights. Instead of storing weights in BF16, you may want to store in FP8 by initializing the model within an fp8_model_init context (see #1880 (comment)). However, unless you're doing something like LoRA, keeping weights in FP8 has implications for convergence. Currently, the distributed optimizer in Megatron-LM is your best bet for avoiding these numerical issues (it applies Adam to sharded FP32 weights, and then does an all-gather in FP8).

Thanks for your reply! In the previous test, I ignored the memory optimization of bf16 master weight in transformerEngine (FP32 master weight can share the first 16 bits with the model weight), this optimization leads to the memory overhead of the non-activations of fp8 being higher than bf16. And I enabled selective recomputation, further diluting the benefits of the fp8 activation memory reduction.
After setting store_param_remainders=False and removing all recomputations, the relative values of cuda memory of fp8 and bf16 expected.