[Quantization] MXFP8LinearConverter should offer filter_fqns instead of fqns

[aditvenk]

Float8LinearConverter offers a convenient filter_fqns method to filter out modules to not quantize ("opt out")
MXFP8LinearConverter, on the other hand, offers a fqns method to "opt in" modules to quantize.

For MXFP8LinearConverter, it is common to want to quantize all linear except the router, lm head, and attention's wk/wv (as pointed out in the official blog as well).

For such usages, the filter_fqns approach is more natural and less error prone.

[acisseJZhong]

@danielvegamyhre wondering could you share more insights on why float8 offer filter_fqns but mxfp8 offers fqns? I also found this a little bit confusing when refactoring.

[danielvegamyhre]

• Float8 training currently uses module swaps (nn.Linear swapped to Float8Linear). For nn.Linear module swaps, the conversion automatically filters to target only linear layers, and so accidentally converting other layers or parameters is a non-issue.
• For MXFP8 training, the conversion instead wraps any arbitrary nn.Parameter data tensor in a special tensor subclass, which implements operator overrides for linear and grouped_mm. So if you just passed "filter_fqns" it would convert all parameters in the model to the subclass except the filtered fqns, which would not be the desired behavior. Therefore, you need to explicitly convert given fqns, like ffn,experts
  • We are planning to move away from this approach and directly instantiate Float8/MXFP8 implementations in TorchTitan itself, which should resolve the annoying-ness of the API.

[aditvenk]

So if you just passed "filter_fqns" it would convert all parameters in the model to the subclass except the filtered fqns, which would not be the desired behavior. Therefore, you need to explicitly convert given fqns, like ffn,experts

Given it is MXFP8LinearConverter -- shouldn't it only convert Linear layers?

[danielvegamyhre]

So if you just passed "filter_fqns" it would convert all parameters in the model to the subclass except the filtered fqns, which would not be the desired behavior. Therefore, you need to explicitly convert given fqns, like ffn,experts

Given it is MXFP8LinearConverter -- shouldn't it only convert Linear layers?

There is no MXFP8LinearConverter anymore - "MXFP8LinearConverter" and "MXFP8GroupedMMConverter" were deleted and replaced with one unified model converter here: #2520 Was it added back...?

Edit: Ok I looked at the current code and see it was split back into 2 separate converters in a recent refactor: 6348d93

This is a bit confusing to me, there is only one unified MXFP8TrainingOpConfig which is used to convert any linear and grouped_mm ops to mxfp8 in a single quantize_() call, I'm not sure why it is split into 2 converters. I'm sure the titan folks had reasons for this though, @acisseJZhong can you share context here?

Edit 2: Reading the code further, I understand now. The refactor traverses the config and converts the Linear.Config or GroupedExperts.Config separately, so there are separate converters to handle each type of config conversion. What I'm seeing here is that the torchao design of generically converting the operators in any standard precision model definition is fundamentally different than the torchtitan refactor design, which has separate configs for linears and grouped experts and wants to be able to convert them separately.

The current code should work for now, but this will be cleaner if we wrap the torchao autograd functions that implement "differentiable, dynamically quantized mxfp8 linear/grouped_mm" in nn.Modules that can be initialized directly in the config converter build() method. Does this sound reasonable @acisseJZhong @tianyu-l ?

[tianyu-l]

Edit 2: Reading the code further, I understand now. The refactor traverses the config and converts the Linear.Config or GroupedExperts.Config separately, so there are separate converters to handle each type of config conversion. What I'm seeing here is that the torchao design of generically converting the operators in any standard precision model definition is fundamentally different than the torchtitan refactor design, which has separate configs for linears and grouped experts and wants to be able to convert them separately.

sounds like we don't have to differentiate linear and grouped mm.

The current code should work for now, but this will be cleaner if we wrap the torchao autograd functions that implement "differentiable, dynamically quantized mxfp8 linear/grouped_mm" in nn.Modules that can be initialized directly in the config converter build() method.

sounds good, as long as we make sure it works with torch.compile and AC (via custom ops)

Any blockers or can we start doing this refactor today?

[danielvegamyhre]

Any blockers or can we start doing this refactor today?

The only blocker is annoying infra issue (internal b200 devgpu broken, and external cloud one i've only been able to reserve no more than 1 b200 at a time lately, which makes it difficult to test torchtitan changes). I will try to figure something out