Megatron FP8 training is compatible with recompute?

[yanchenmochen]

Your question
Ask a clear and concise question about Megatron-LM.
when I tried run the training using fp8 parameters, OOM occurs. the model is 7B model. same parameter is valid for --bf16.

    --bf16 \
    --fp8-format hybrid \
    --fp8-amax-compute-algo max \
    --fp8-amax-history-len 16 \
    --transformer-impl transformer_engine

How to set he correct fp8 parameter? I tried to reduce memory usage to avoid OOM by using--recalculate-granularity selective, but failed, still OOM

[shanmugamr1992]

Hi yanchenmochen,
one option is recalculate-granularity (full) , but beware that this will recompute all activations and is extremely slow (~30% slower to train atleast). A better option would just be to leave it at bf16, or if you really want to go fp8 route, maybe try model parallelism. That would be a better bet.