| 824 | # --------------------------------------------------------------------------- |
| 825 | |
| 826 | def parse_args() -> argparse.Namespace: |
| 827 | parser = argparse.ArgumentParser( |
| 828 | description=( |
| 829 | "AutoKernel Model Profiler -- identify GPU kernel bottlenecks " |
| 830 | "in any PyTorch model." |
| 831 | ), |
| 832 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 833 | epilog=( |
| 834 | "Examples:\n" |
| 835 | " uv run profile.py --model models/llama_7b.py " |
| 836 | "--class-name LlamaModel --input-shape 1,2048 --dtype float16\n" |
| 837 | " uv run profile.py --module transformers " |
| 838 | "--class-name AutoModelForCausalLM " |
| 839 | "--pretrained meta-llama/Llama-2-7b-hf --input-shape 1,2048\n" |
| 840 | " uv run profile.py --model my_net.py " |
| 841 | "--class-name MyNet --input-shape 8,3,224,224 --dtype float32\n" |
| 842 | ), |
| 843 | ) |
| 844 | |
| 845 | # Model source |
| 846 | parser.add_argument( |
| 847 | "--model", |
| 848 | type=str, |
| 849 | default=None, |
| 850 | help="Path to a Python file containing the model class.", |
| 851 | ) |
| 852 | parser.add_argument( |
| 853 | "--module", |
| 854 | type=str, |
| 855 | default=None, |
| 856 | help="Python module to import the model from (e.g. 'transformers').", |
| 857 | ) |
| 858 | parser.add_argument( |
| 859 | "--class-name", |
| 860 | type=str, |
| 861 | required=True, |
| 862 | help="Name of the model class to instantiate.", |
| 863 | ) |
| 864 | parser.add_argument( |
| 865 | "--pretrained", |
| 866 | type=str, |
| 867 | default=None, |
| 868 | help="Pretrained model name/path for HuggingFace from_pretrained().", |
| 869 | ) |
| 870 | |
| 871 | # Input configuration |
| 872 | parser.add_argument( |
| 873 | "--input-shape", |
| 874 | type=str, |
| 875 | required=True, |
| 876 | help="Comma-separated input shape, e.g. '1,2048' or '8,3,224,224'.", |
| 877 | ) |
| 878 | parser.add_argument( |
| 879 | "--dtype", |
| 880 | type=str, |
| 881 | default="float16", |
| 882 | help="Data type: float16, bfloat16, float32 (default: float16).", |
| 883 | ) |