bionemo-framework/bionemo-recipes/recipes/esm2_native_te/fp8_debugging_stats.yaml at main · polinabinder1/bionemo-framework · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
example_fp8_tensor_stat_collection:
    enabled: True
    layers:
        # Match the actual linear layers within attention that support FP8 stats
        layer_types: [layernorm_qkv, proj, fc1, fc2]
    transformer_engine:
        LogFp8TensorStats:
            enabled: True
            tensors_struct:
            - tensor: activation
              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
              freq: 10
            - tensor: gradient
              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
              freq: 10
            - tensor: weight
              stats: [underflows%, scale_inv_min, scale_inv_max, mse]
              freq: 10
        LogTensorStats:
          enabled: True
          stats: [max, min, mean, std, l1_norm]
          tensors: [dgrad, wgrad]
          freq: 1