[AOTI] Add a --dynamic-shapes option to export

[desertfire]

Summary: The --dynamic-shapes option will default to False. When the actual inputs are with static shapes, calling export with static shapes will make sure more Inductor optimizations take effect down the line. This change by itself improves average tokens/sec from 29.60 to 33.43 on A100. Some following PRs will provide further perf gains.

python3 torchchat.py export llama3 --quantize '{"precision": {"dtype":"bfloat16"}, "executor":{"accelerator":"cuda"}}' --output-dso-path /tmp/model16.so && python3 torchchat.py generate llama3 --dso-path /tmp/model16.so --prompt "Once upon a time," --max-new-tokens 256 --device cuda --num-samples 3

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1011

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 1 Pending

As of commit 334654e with merge base 912917f ():
💚 Looks good so far! There are no failures yet. 💚

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

[Jack-Khuu]

Great to see the perf gains, is anything lost by switching over to static? Does model support change?

[desertfire]

Great to see the perf gains, is anything lost by switching over to static? Does model support change?

ET export is doing the same so it should be fine,

torchchat/export_util/export_et.py

Lines 57 to 65 in 912917f

	input = (
	torch.tensor([[1]], dtype=torch.long, device=device),
	torch.tensor([0], dtype=torch.long, device=device),
	)
	state_dict = model.state_dict()
	state_dict_dtype = state_dict[next(iter(state_dict))].dtype
	target_precision = get_precision()
	dynamic_shapes = None

[Jack-Khuu]

Amazing, thanks for fixing this

One last ask: Can you put the repro commands for the #'s in your description?

[desertfire]

Amazing, thanks for fixing this

One last ask: Can you put the repro commands for the #'s in your description?

Done. We should add a benchmarking script so that everyone can run the same experiment.