diff --git a/aiu_fms_testing_utils/scripts/inference.py b/aiu_fms_testing_utils/scripts/inference.py index 45d52cc3..502cc07f 100644 --- a/aiu_fms_testing_utils/scripts/inference.py +++ b/aiu_fms_testing_utils/scripts/inference.py @@ -138,6 +138,13 @@ default="inductor", choices=["inductor", "eager", "aot_eager"], ) +parser.add_argument( + "--compile_backend_aiu", + type=str, + help="Backend for AIU compilation", + default="sendnn", + choices=["sendnn", "sendnn_compile_only"], +) parser.add_argument( "--compile_dynamic", action="store_true", @@ -387,7 +394,10 @@ print("must set AIU_WORLD_RANK_0") exit() os.environ.setdefault("FLEX_COMPUTE", "SENTIENT") - os.environ.setdefault("FLEX_DEVICE", "PF") + if args.compile_backend_aiu == "sendnn_compile_only": + os.environ.setdefault("FLEX_DEVICE", "COMPILE") + else: + os.environ.setdefault("FLEX_DEVICE", "PF") device = torch.device("cpu") else: @@ -597,7 +607,8 @@ def select_int8_module( fx_config.backed_size_oblivious = "paged" in attn_name if is_aiu_backend: model.compile( - backend="sendnn", options={"sendnn.dynamic": args.compile_dynamic_sendnn} + backend=args.compile_backend_aiu, + options={"sendnn.dynamic": args.compile_dynamic_sendnn}, ) else: # compiling can make first inference pass slow