Fix convert_model_to_fp8_ao converting the first and last linear layers

[vineethsaivs]

What does this PR do?

convert_model_to_fp8_ao documents that it converts every nn.Linear "except the first and last", and find_first_last_linear_layers exists precisely because quantizing those two destabilises training:

For stability reasons, we skip the first and last linear layers
Otherwise can lead to the model not training or converging properly

It does not skip them. On a three-linear model, with real torchao 0.17.0 on CPU:

default                  {'embed_proj': True, 'block.0': True, 'lm_head': True}
module_filter_func=None  {'embed_proj': False, 'block.0': True, 'lm_head': False}

True means the layer was swapped for Float8Linear. The first line is what a direct caller of convert_model_to_fp8_ao(model) gets today, which is the call its own docstring example shows.

Root cause

The default module_filter_func is filter_first_and_last_linear_layers, which does:

first_linear, last_linear = find_first_last_linear_layers(module)
return filter_linear_layers(module, fqn, layers_to_filter=[first_linear, last_linear])

module there is not the model. torchao's swap_linear_layers calls module_filter_fn(module, cur_fqn) once per candidate layer, so find_first_last_linear_layers runs on a single nn.Linear, whose named_modules() yields only ("", itself). It therefore returns ("", ""), and fqn in ["", ""] is false for every real FQN, so the filter approves everything:

model's real first/last linear: ('embed_proj', 'lm_head')

what the default filter decides for each linear (True = convert to FP8):
  embed_proj         convert=True  <-- should have been skipped
  block.0            convert=True
  lm_head            convert=True  <-- should have been skipped

find_first_last_linear_layers(a single Linear) -> ('', '')

The fix

The correct implementation is already in the function, one line below, but unreachable by default:

first_linear, last_linear = find_first_last_linear_layers(model)
if module_filter_func is None:
    module_filter_func = partial(filter_linear_layers, layers_to_filter=[first_linear, last_linear])

So the change is to default module_filter_func to None. That also makes the signature agree with its own docstring, which has said defaults to filter_linear_layers since #3348, and it restores the branch #3450 was fixing ("I didn't actually update the call in the case of the default being used"), which has been dead for a direct caller since the signature default was set.

Blast radius, stated rather than implied

Accelerator is not affected. AORecipeKwargs.module_filter_func defaults to None and is passed straight through to convert_model_to_fp8_ao, and an explicit None argument beats the signature default, so the FP8 path through Accelerator was already taking the correct branch. What this fixes is the direct call, which is exported from accelerate.utils and is what the docstring example demonstrates.

One thing I did not change, and would like your call on

After this change nothing references filter_first_and_last_linear_layers, and I do not think it can be made to work with its current signature: torchao's contract is (module, fqn) -> bool and the first and last linear cannot be derived from one candidate layer. The options I see are to turn it into a factory that takes the model and returns a bound filter, or to drop it. Both change or remove a symbol exported from accelerate.utils, so I left it alone rather than decide that in a bug-fix PR. Happy to do either here or in a follow-up, whichever you prefer.

Testing

Added test_convert_model_to_fp8_ao_skips_the_first_and_last_linear_layers to tests/test_utils.py, guarded with @require_torchao. It builds a three-linear toy model, calls convert_model_to_fp8_ao(model) with no filter argument, and asserts the first and last stay nn.Linear while the middle one becomes Float8Linear. It is CPU-only; the layer swap needs no GPU, which is why it can live in tests/test_utils.py rather than the launcher-based tests/test_fp8.py.

Three runs against accelerate 1.14.0, whose src/accelerate/utils/ao.py is byte-identical to main at 16cb6eb, with torch 2.13.0 and torchao 0.17.0:

control   upstream ao.py + upstream tests/test_utils.py   45 passed, 4 skipped
before    upstream ao.py + this test                      1 failed
after     this branch                                     46 passed, 4 skipped

ruff check and ruff format --check report the same result on the changed files as on the unmodified ones (two pre-existing preview-rule findings in ao.py, collapsible-if and needless-bool, both present before this change and untouched by it), so the lint result is a real comparison rather than a config that checks nothing.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

On the documentation box: no docs change was needed, because the docstring already describes the fixed behaviour; it was the signature that disagreed with it.

Who can review?

@SunMarc @BenjaminBossan