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

[SunMarc]

Thanks ! just a nit

[SunMarc]

we have a tests/quantization/torchao file for test related to torhcao, can you move that there ?

[SunMarc]

can you check if this was also the case with prior version of torchao ? we force the user to install a min version of torchao so you can test this one

[vineethsaivs]

Thanks @SunMarc, both done.

On the test location. There is no tests/quantization/torchao in accelerate; the tree has tests/test_quantization.py, which is bitsandbytes-only and gated on require_cuda_or_xpu plus require_bnb, and tests/test_fp8.py, which holds the existing TestTorchAO cases. I moved it to tests/test_fp8.py as a new TestTorchAOFilters class, next to TestTorchAO. It is a plain unittest.TestCase rather than a launcher one on purpose: the layer swap itself needs no accelerator, so routing it through get_launch_command would only make it slower and GPU-bound for no gain. tests/test_utils.py is back to its upstream content. Happy to create a tests/quantization/ package instead if that is the layout you want.

On prior torchao versions: yes, it has always been broken, and I checked rather than assumed. The relevant thing is torchao's calling convention, since the bug is that our filter is handed one candidate layer rather than the model. swap_linear_layers calls module_filter_fn(module, cur_fqn) identically in v0.6.1, v0.7.0, v0.9.0 and main, so it is not a regression from a torchao change.

I also ran it on the exact minimum is_torchao_available() enforces, torchao==0.6.1:

torchao 0.6.1, upstream ao.py -> {'embed_proj': True, 'block.0': True, 'lm_head': True}
torchao 0.6.1, with the fix   -> {'embed_proj': False, 'block.0': True, 'lm_head': False}

True means the layer was swapped for Float8Linear. Same result as on 0.17.0, so every torchao version accelerate supports converts the first and last linear, and the one-token fix repairs all of them.

The test still fails on the previous commit and passes on this one, and ruff check and ruff format --check are clean on the moved file.

[SunMarc]

we have a test file for torchao, don't put it here

[vineethsaivs]

Moved it out of tests/test_fp8.py into a new tests/test_torchao.py in 8179880, so test_fp8.py is back to its upstream content.

I could not find an existing torchao test file to move it into, and I looked twice, so here is the whole picture in case I am missing something. Every path in huggingface/accelerate@main matching torchao:

benchmarks/fp8/torchao/Dockerfile
benchmarks/fp8/torchao/README.md
benchmarks/fp8/torchao/ddp.py
benchmarks/fp8/torchao/distrib_deepspeed.py
benchmarks/fp8/torchao/fp8_utils.py
benchmarks/fp8/torchao/fsdp.py
benchmarks/fp8/torchao/non_distributed.py

and every path matching quantization:

docs/source/usage_guides/quantization.md
tests/test_quantization.py

tests/test_quantization.py is bitsandbytes-only and gated on require_cuda_or_xpu plus require_bnb, and the only torchao tests in tests/ are the TestTorchAO launcher cases inside test_fp8.py. There is no tests/quantization/ package. You may be thinking of transformers, which does have tests/quantization/torchao_integration/test_torchao.py.

So I created tests/test_torchao.py, named to sit beside tests/test_quantization.py. If you would rather have tests/quantization/torchao.py mirroring the transformers layout, or somewhere else entirely, say the word and I will move it; it is one file with one test.

It is still a plain unittest.TestCase rather than a launcher one, since the layer swap needs no accelerator. require_torchao guards it, it fails on the previous commit and passes on this one, and ruff check and ruff format --check are clean.

[SunMarc]

there tests/quantization/torchao_integration ....

[vineethsaivs]

Done in ee69aa0c, it now lives at tests/quantization/torchao_integration/test_torchao.py.

One thing to flag so it is not a surprise: that directory did not exist in accelerate, so this commit creates it. I think you may have had transformers in mind, which does have tests/quantization/torchao_integration/test_torchao.py. Rather than keep going back and forth about it, I have just created the path you asked for.

It still runs: make test collects ./tests/ recursively and only ignores deepspeed, fsdp and tp, so the new directory is picked up with no Makefile change, and I confirmed collection.

If you would rather this became the home for the bitsandbytes tests in tests/test_quantization.py too, say the word and I will do that move in a follow-up rather than widening this PR.