Adding torchao apis to gpt-fast

[HDCharles]

Summary:

adding torchao apis to gpt-fast and some minor tweaks

Test Plan:

(in progress)
export MODEL_REPO=meta-llama/Meta-Llama-3-8B

python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode torchao-int8 python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int8.pth --compile python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int8.pth python eval.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int8.pth --tasks wikitext --compile

For model checkpoints/meta-llama/Meta-Llama-3-8B/model_torchao-int8.pth
wikitext: {'word_perplexity,none': 7.900496793735154, 'word_perplexity_stderr,none': 'N/A', 'byte_perplexity,none': 1.4718578218273202, 'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none': 0.5576383170121927, 'bits_per_byte_stderr,none': 'N/A', 'alias': 'wikitext'}

python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode torchao-int4-hqq python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4-hqq.pth --compile python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4-hqq.pth python eval.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4-hqq.pth --tasks wikitext --compile

For model checkpoints/meta-llama/Meta-Llama-3-8B/model_torchao-int4-hqq.pth
wikitext: {'word_perplexity,none': 8.44187872159186, 'word_perplexity_stderr,none': 'N/A', 'byte_perplexity,none': 1.4902143610748824, 'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none': 0.575519871235033, 'bits_per_byte_stderr,none': 'N/A', 'alias': 'wikitext'}

python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode torchao-int4 python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4.pth --compile python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4.pth python eval.py --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4.pth --tasks wikitext --compile

For model checkpoints/meta-llama/Meta-Llama-3-8B/model_torchao-int4.pth
wikitext: {'word_perplexity,none': 8.59031159441983, 'word_perplexity_stderr,none': 'N/A', 'byte_perplexity,none': 1.4950796712267396, 'byte_perplexity_stderr,none': 'N/A', 'bits_per_byte,none': 0.5802223661766339, 'bits_per_byte_stderr,none': 'N/A', 'alias': 'wikitext'}

Reviewers:

Subscribers:

Tasks:

Tags:

[HDCharles]

@Chillee should i add info to the README or implement this differently?

[gchlebus]

Is the plan for this PR also to add fp8 support which is available in torchao?

[jerryzh168]

it seems that this one does not work with tp yet:

ENABLE_INTRA_NODE_COMM=1 torchrun --standalone --nproc_per_node=2 generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model_torchao-int4.pth

[rank1]: NotImplementedError: AffineQuantizedTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.split_with_sizes', overload='default')>, types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>,), arg_types=(<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>, <class 'list'>), kwarg_types={}

will need to implement this op in AQT to support this or change the tp implementation to DTensor I guess.