add mxfp8 and nvfp4 to Llama eval scripts

[vkuzo]

Summary:

Adds mxfp8 and nvfp4 to llama eval scripts.

• baseline: wikitext 7.55, winogrande 0.743
• mxfp8_floor: wikitext 7.61, winogrande 0.729
• mxfp8_rceil: wikitext 7.60, winogrande 0.739
• nvfp4: wikitext 8.44, winogrande 0.718
• float8 rowwise: wikitext 7.62, winogrande 0.737

Results:

// bf16 baseline
with-proxy time python torchao/_models/llama/eval.py --checkpoint_path
checkpoints/meta-llama/Meta-Llama-3.1-8B/model.pth --print_model --tasks
wikitext winogrande
wikitext: {'alias': 'wikitext', 'word_perplexity,none':
7.5472105433748435, 'word_perplexity_stderr,none': 'N/A',
'byte_perplexity,none': 1.459319739134015,
'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none':
0.5452960145272896, 'bits_per_byte_stderr,none': 'N/A'}
winogrande: {'alias': 'winogrande', 'acc,none': 0.7426992896606156,
'acc_stderr,none': 0.012285989618865697}

// mxfp8 with floor scaling, turned off compile as it seemed stuck in coordinate descent
tuning
with-proxy time python torchao/_models/llama/eval.py --checkpoint_path
checkpoints/meta-llama/Meta-Llama-3.1-8B/model.pth --print_model --tasks
wikitext winogrande --quantization mxfp8
wikitext: {'alias': 'wikitext', 'word_perplexity,none':
7.609070006132819, 'word_perplexity_stderr,none': 'N/A',
'byte_perplexity,none': 1.4615491037668933,
'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none':
0.5474983002838458, 'bits_per_byte_stderr,none': 'N/A'}
winogrande: {'alias': 'winogrande', 'acc,none': 0.7292817679558011,
'acc_stderr,none': 0.012487904760626407}

// mxfp8 with rceil scaling
wikitext: {'alias': 'wikitext', 'word_perplexity,none':
7.605445025927753, 'word_perplexity_stderr,none': 'N/A',
'byte_perplexity,none': 1.4614188696390065,
'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none':
0.5473697404554175, 'bits_per_byte_stderr,none': 'N/A'}
winogrande: {'alias': 'winogrande', 'acc,none': 0.7387529597474349,
'acc_stderr,none': 0.012346914863415201}

// nvfp4
wikitext: {'alias': 'wikitext', 'word_perplexity,none':
8.44478255417328, 'word_perplexity_stderr,none': 'N/A',
'byte_perplexity,none': 1.4903102070118779,
'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none':
0.5756126578938119, 'bits_per_byte_stderr,none': 'N/A'}
winogrande: {'alias': 'winogrande', 'acc,none': 0.7182320441988951,
'acc_stderr,none': 0.012643326011853038}

// float8 rowwise (for comparison to existing technique)
wikitext: {'alias': 'wikitext', 'word_perplexity,none':
7.618818730886612, 'word_perplexity_stderr,none': 'N/A',
'byte_perplexity,none': 1.4618990946965715,
'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none':
0.5478437349532752, 'bits_per_byte_stderr,none': 'N/A'}
winogrande: {'alias': 'winogrande', 'acc,none': 0.7371744277821626,
'acc_stderr,none': 0.01237092252726192}

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• flip mx inference scaling setting to RCEIL #3428
• -> add mxfp8 and nvfp4 to Llama eval scripts #3394

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3394

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit b4cf67b with merge base 16aad7c ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests (aarch64) / test-cpu-ops (linux.arm64.2xlarge) (gh) (trunk failure)
  ##[error]The operation was canceled.
• Run Regression Tests (aarch64) / test-cpu-ops (macos-14) (gh) (trunk failure)
  Process completed with exit code 139.

This comment was automatically generated by Dr. CI and updates every 15 minutes.