Bug description

Hi,
I tried running llama3_8b with rowwise fp8 quantization and get NaN loss. Same command worked ~2 weeks ago.

Repro:

CONFIG_FILE="./torchtitan/models/llama/train_configs/llama3_8b.toml" ./run_train.sh --training.steps=100 --training.batch_size=2 --training.compile --model.converters="float8" --float8.recipe_name rowwise

After discussion with @lessw2020: AC is 'selective' via default.

Versions

torch Apr 4 nightly
torchtitan commit 82ef16a Fri Apr 4 14:27:25
torchao commit 6922733cbd06aa04a0db4b3b3f29606662f6fe75 Fri Apr 4 10:56:28
8xNVIDIA H100