Feat (brevitas_examples/common): support for mse scale for weights quantized to float formats

[JP-Amboage]

Reason for this PR

It was not possible to run the LLM Examples with a configuration using weight_quant_format: float_eXmY and weight_param_method: mse. That is, if a "regular" (no OCP) float was being used to quantize the weights, then it was not possible to use mean squared error to choose the scale. I need to use this configuration for some of my experiments.

Changes Made in this PR

• In src/brevitas_examples/common/generative/quantizers.py, I added the class Fp8e4m3WeightPerChannelFloatMSE that inherits from MSESymmetricScale and Fp8e4m3WeightPerChannelFloat. This was done following the same pattern used to support mse scale for Float OCP per-channel weights which was already available.
• In src/brevitas_examples/common/generative/quantize.py, I added the entry

'mse': {
                'per_channel': {
                    'sym': Fp8e4m3WeightPerChannelFloatMSE}}}},

for float, float_scale in the dictionary WEIGHT_QUANT_MAP. This is how it is currently done for other data types (i.e. Float OCP).

Testing Summary

• Added a new test in tests/brevitas_examples/test_llm_cases.py to ensure that the combination of weight_quant_format: float_e2m1 with weight_param_method: mse runs.
  • In the second commit a small change was done to the tests to skip the new added test when JIT is enabled as this does not support weight_param_method: mse.
• Tested locally

[pablomlago]

LGTM