[mamf-finder] Adding FP4 support

[aflah02]

Hi @stas00
Based on our discussion on X, I looked around and someone on the EleutherAI discord pointed out the scaled_mm API - https://github.com/pytorch/pytorch/blob/88c77db9c862573f9d7a8eda58ae735415bc740d/torch/nn/functional.py#L6752 which is used under the hood with TorchAO. They also pointed to these benchmark functions in the TorchAO repo which seem to imply that post Torch 2.10.0 you can use MXPF4 with scaled_mm - https://github.com/pytorch/ao/blob/fe986580eaafc87f532534a8f222c7d11af18702/benchmarks/float8/bench_matmul.py#L166

I'm thinking of bringing both MXFP4 and NVP4 support in using the function params in the second link as reference.

[stas00]

Thank you for researching, Aflah.

mamf-finder already uses scaled_mm for fp8, so I suppose it might not be too difficult to expand it to use the fp4 formats, if you'd like to try it out.

If torchao is required for this functionality, this should be a soft-dependency (i.e. assert requiring it only if these formats are requested), so that for fp8/bf16/etc the script requires nothing else but pytorch.

[aflah02]

Thanks! I don't think torchao should be a requirement but will only be sure once I finish up testing

[aflah02]

Hi @stas00
In the current script the dtype supplied via the CLI arg is expected to be a valid torch dtype. In the torchao example there are 2 recipes, mxfp4_cutlass and nvfp4 (See here). For mxfp4_cutlass the code path uses the to_mx function to prepare data while in the nvfp4 branch it uses nvfp4_quantize. So I plan to add 2 dtypes for fp4 (mxfp4 and nvfp4) which will use the above-described recipes however will not undergo the dtype parsing present in the script as they don't have a dtype persay. Is that fine? An alternative could be the dtype is float4_e2m1fn_x2 and I create a separate arg called fp4_recipe which takes one of mxfp4_cutlass or nvfp4. Happy to go with whatever you prefer!

[stas00]

so torch has no dtype reserved for these formats? That would be odd, no?

But otherwise for sure, we can try/except on argparse torch dtype and if it doesn't exist fallback to a custom lookup and fail only if that fails. I think a single arg is better.

[aflah02]

so torch has no dtype reserved for these formats? That would be odd, no?

Atleast as far as I can see in the latest release, there is only one fp4 dtype - torch.float4_e2m1fn_x2 (https://docs.pytorch.org/docs/stable/tensor_attributes.html) and inside torchao the other two are simply called recipes (https://github.com/pytorch/ao/blob/fe986580eaafc87f532534a8f222c7d11af18702/benchmarks/float8/bench_matmul.py#L48)

But otherwise for sure, we can try/except on argparse torch dtype and if it doesn't exist fallback to a custom lookup and fail only if that fails. I think a single arg is better.

Makes sense!