Enable loading pre-quantized INT4 weights in Llama4

[jiawenliu64]

Generate INT4 MP8 checkpoint:

torchrun --nproc-per-node=8 -m models.llama4.scripts.quantize --ckpt_dir ../checkpoints/Llama-4-Scout-17B-16E-Instruct --output_dir ../checkpoints/Llama-4-Scout-17B-16E-Instruct-INT4 --quantization_mode int4_mixed --world_size 8

Verify generated INT4 MP8 checkpoint with int4_mixed on single GPU (output):

PYTHONPATH=$(git rev-parse --show-toplevel) torchrun --nproc_per_node=1 -m models.llama4.scripts.chat_completion ../checkpoints/Llama-4-Scout-17B-16E-Instruct-INT4 --world_size 1 --quantization-mode int4_mixed

Generate FP8 MP8 checkpoint:

torchrun --nproc-per-node=8 -m models.llama4.scripts.quantize --ckpt_dir ../checkpoints/Llama-4-Scout-17B-16E-Instruct --output_dir ../checkpoints/Llama-4-Scout-17B-16E-Instruct-FP8 --quantization_mode fp8_mixed --world_size 8

Verify generated FP8 MP8 checkpoint with fp8_mixed (output):

PYTHONPATH=$(git rev-parse --show-toplevel) torchrun --nproc_per_node=8 -m models.llama4.scripts.chat_completion ../checkpoints/Llama-4-Scout-17B-16E-Instruct-FP8 --world_size 8 --quantization-mode fp8_mixed

Verify BF16 MP8 checkpoint (output):

PYTHONPATH=$(git rev-parse --show-toplevel) torchrun --nproc_per_node=8 -m models.llama4.scripts.chat_completion ../checkpoints/Llama-4-Scout-17B-16E-Instruct --world_size 8

Verify BF16 MP8 checkpoint with fp8_mixed (output):

PYTHONPATH=$(git rev-parse --show-toplevel) torchrun --nproc_per_node=8 -m models.llama4.scripts.chat_completion ../checkpoints/Llama-4-Scout-17B-16E-Instruct --world_size 8 --quantization-mode fp8_mixed

Verify BF16 MP8 checkpoint with int4_mixed on single GPU (output):

PYTHONPATH=$(git rev-parse --show-toplevel) torchrun --nproc_per_node=1 -m models.llama4.scripts.chat_completion ../checkpoints/Llama-4-Scout-17B-16E-Instruct --world_size 1 --quantization-mode int4_mixed

[ashwinb]

this feels like complexity that truly doesn't belong at this layer. can we please keep it outside into quantization code somehow?

[ashwinb]

we don't want llama-models to become torchao or vllm or whatever really. it is not a full fledged all powerful inference engine.

[ashwinb]

if you move the model.load_state_dict() to convert_to_quantized_model() then you can do the following:

• change the structure of the Transformer from the outside in this code path (whatever you are doing with Experts)
• move all this scale ckpt paths complexity into quantization land

nobody reading generation.py should know about quantization unless they want to dig into it.