the error when save optimizer state

[zyuanzi]

Hi,

I have finished reading your excellent paper and am trying to run and reproduce your code. I would like to ask a few questions:

1. I noticed that the saved weights are still stored in FP32 format. Why are they not saved in FP8?
2. Apart from the optimizer states, where should I look for the code implementation of layers and operators that use FP8? Is it possible for me to observe gradients, activations, weights, etc., in FP8 format in the model architecture?
3. I noticed that in the paper, you mentioned training on the MAmmoth dataset with llama2-7B for 3 epochs and then testing its performance on four datasets, including NumGLUE. May I ask which testing toolkit you used for evaluation?
4. How should I evaluate the memory footprint it occupies?

Thank you very much for your time!

Best Regards!

[haochengxi]

1. Hi, saved weights are stored in FP32 because they are "master weights". Master weights need to be stored with high precision, to make sure that the update of each step using gradient can be precisely reflected in the weights. Usually, the gradients of each step are relatively small, so if the master weights are stored in FP8, the weight update might loss critical informations.

2. I think you can look into https://github.com/NVlabs/COAT/blob/main/coat/models/coat_olmo.py. You can print observe the gradients, activations, weight there. An example:

   COAT/coat/models/coat_olmo.py

   Line 324 in 80ec99f

   fc2_x = fp8_linear_forward(flash_qx, flash_s, weight2, weight2_s, False, group_size) #

   .

3. We use the official MAmmoth training recipe to do this.

4. You can refer to the code here to see how to evaluate the memory footprint it occupies:

   COAT/examples/OLMo/olmo/train.py

   Line 1262 in 80ec99f

   print(f"\n\nPeak Memory Usage : {peak_memory:.1f} MB\n"

[zyuanzi]

Hi ,
Thank you for your reply！

I still have one more question about the third point. I noticed that the MAmmoth dataset contains data from four downstream task datasets such as NumGlue. Did you train on the MAmmoth dataset for three epochs and then test them respectively in the test files of datasets like numGlue? I tried the model's inference, and even if it was correct, It's also very difficult to be exactly the same as labels. May I ask how this metric is evaluated~

Best Regards~

[haochengxi]

Hi, I remember that I just follow their official repo's guide to do the evaluation. Maybe you can begin with fine-tuning using their official repo using BF16 to see if everything works fine, then migrate to FP8.