Adding compilation support for io_encrypt flag (#393)

As part of Model-IP I/O encryption feature qaic-exec has added a new
command line option: `-io-encrypt` to enable support for
decryption/encryption of network I/O.
Added support in Qeff to pass this param to qaic-exec when a model needs
to be compiled with I/O encryption support.

CLI command: `python -m QEfficient.cloud.infer --model_name gpt2
--batch_size 1 --prompt_len 32 --ctx_len 128 --mxfp6 --num_cores 16
--device_group [0] --prompt "My name is" --mos 1
--aic_enable_depth_first --io-encrypt "chacha20"`

Note: If the flag is passed via infer CLI, CLI will exit after QPC
generation.

---------

Signed-off-by: Asmita Goswami <[email protected]>

Author: asmigosw

QEfficient/cloud/compile.py

@@ -85,17 +85,29 @@
     parser.add_argument(
         "--enable_qnn",
         "--enable-qnn",
-        action="store_true",
+        nargs="?",
+        const=True,
+        type=str,
         default=False,
         help="Enables QNN. Optionally, a configuration file can be provided with [--enable_qnn CONFIG_FILE].\
              If not provided, the default configuration will be used.\
              Sample Config: QEfficient/compile/qnn_config.json",
     )
-    parser.add_argument(
-        "qnn_config",
-        nargs="?",
-        type=str,
-    )
-    # FIXME(ochougul): Allow extra compilation arguments
-    args = parser.parse_args()
-    QEfficient.compile(**vars(args))
+
+    args, compiler_options = parser.parse_known_args()
+
+    if isinstance(args.enable_qnn, str):
+        args.qnn_config = args.enable_qnn
+        args.enable_qnn = True
+
+    compiler_options_dict = {}
+    for i in range(0, len(compiler_options)):
+        if compiler_options[i].startswith("--"):
+            key = compiler_options[i].lstrip("-").replace("-", "_")
+            value = (
+                compiler_options[i + 1]
+                if i + 1 < len(compiler_options) and not compiler_options[i + 1].startswith("-")
+                else True
+            )
+            compiler_options_dict[key] = value
+    QEfficient.compile(**args.__dict__, **compiler_options_dict)

QEfficient/cloud/infer.py

@@ -197,6 +197,10 @@ def main(
         **kwargs,
     )
 
+    #  If the io-encrypt flag is passed we will exit after QPC generation.
+    if kwargs.get("io_encrypt", None):
+        exit()
+
     #########
     # Execute
     #########

QEfficient/compile/compile_helper.py

@@ -64,9 +64,6 @@ def compile_kv_model_on_cloud_ai_100(
         DeprecationWarning,
         stacklevel=2,
     )
-    if kwargs:
-        # FIXME
-        raise NotImplementedError("Can't handle extra compilation args now!")
     aic_binary_dir = os.path.join(base_path, "qpcs")
 
     if os.path.isdir(aic_binary_dir):
@@ -111,6 +108,13 @@ def compile_kv_model_on_cloud_ai_100(
         with open(mdp_ts_config_path, "w") as file:
             json.dump(mdp_ts_config, file, indent=4)
         command.append(f"-mdp-load-partition-config={mdp_ts_config_path}")
+    for key, value in kwargs.items():
+        option = "-" + key.replace("_", "-")
+        if isinstance(value, bool):
+            if value:
+                command.append(option)
+            continue
+        command.append(f"{option}={value}")
     print("Running AI 100 compiler:", " ".join(command))
     result = subprocess.run(command, capture_output=True, text=True)
     if result.returncode != 0:
@@ -221,6 +225,13 @@ def compile(
             allow_mxint8_mdp_io=allow_mxint8_mdp_io,
             mos=mos,
             device_group=device_group,
+            **kwargs,
         )
-        logger.info(f"Compiled QPC files can be found here: {qpc_path}")
+        if kwargs.get("io_encrypt", None):
+            logger.warning(
+                f"Compilation for IO-Encrypt has been successfully completed at path: {qpc_path}. However, Efficient-Transformers do not support IO-Encrypt execution. Please run the execution separately"
+            )
+        else:
+            logger.info(f"Compiled QPC files can be found here: {qpc_path}")
+
     return qpc_path

QEfficient/transformers/models/modeling_auto.py

@@ -1635,6 +1635,11 @@ def compile(
             **compiler_options,
         )
 
+        if compiler_options.get("io_encrypt", None):
+            logger.warning(
+                "Compilation for IO-Encrypt has been successfully completed. However, Efficient-Transformers do not support IO-Encrypt execution. Please run the execution separately with QPC compiled without io-encrypt."
+            )
+
         return qpc_path
 
     # FIXME: Update this method to match with transformers AutoModelForCausalLM.generate