Refactor inference.py for LLM and RoBERTa support #34

andrea-fasoli · 2025-04-23T16:43:38Z

This PR implements a substantial refactoring of inference.py which becomes the single entry point for LLMs and RoBERTa models. Support covers non-quantized, GPTQ W4A16, and INT8 models.

inference.py code has been streamlined. It is now structured into the following sections:

args_parsing            define script arguments across all model configurations
aiu_setup               set up AIU environment variables
model_setup             define model dtype, device, and distributed strategy
quantization_setup      import FMS-MO addons and define linear_config for FMS
direct_quantization     quantize a non-quantized model to INT8 (WIP)
decoders                run token generation task with LLMs
encoders                run QA or MLM task with RoBERTa

Extensive code validation is needed prior merging.

Signed-off-by: Andrea Fasoli <andrea.fasoli@ibm.com>

andrea-fasoli · 2025-04-23T16:44:45Z

cc for review: @ani300 @JRosenkranz

ani300 · 2025-04-28T14:32:28Z

aiu_fms_testing_utils/utils/aiu_setup.py

@@ -57,6 +59,8 @@ def aiu_setup(rank=0, world_size=1, local_rank=0, local_size=1, verbose=False):
 def aiu_dist_setup(rank, world_size, local_rank=-0, local_size=-1, verbose=False):
    if local_rank < 0:
        local_rank = rank
+
+    # FIXME: local_size not in use ?


maybe it's no longer needed? if you can't find any reference to it feel free to delete

ani300 · 2025-04-28T14:33:09Z

aiu_fms_testing_utils/utils/aiu_setup.py

+    _target_cache_size = max(
+        int(args.max_new_tokens * 2),
+        int(args.min_pad_length * 2.5),
+        int(args.fixed_prompt_length * 2.5),
+    )
+    _prompt_size = max(int(args.min_pad_length), int(args.fixed_prompt_length))
+    if hasattr(torch._dynamo.config, "accumulated_cache_size_limit"):
+        if _target_cache_size > torch._dynamo.config.accumulated_cache_size_limit:
+            _prev = torch._dynamo.config.accumulated_cache_size_limit
+            torch._dynamo.config.accumulated_cache_size_limit = _target_cache_size
+            dprint(
+                "NOTICE: Adjusting torch._dynamo.config.accumulated_cache_size_limit "
+                f"from {_prev} to {torch._dynamo.config.accumulated_cache_size_limit} "
+                f"to accomodate prompt size of {_prompt_size} and decode tokens of "
+                f"{args.max_new_tokens}"
+            )
+
+    if _target_cache_size > torch._dynamo.config.cache_size_limit:
+        _prev = torch._dynamo.config.cache_size_limit
+        torch._dynamo.config.cache_size_limit = _target_cache_size
+        dprint(
+            f"NOTICE: Adjusting torch._dynamo.config.cache_size_limit from {_prev} to "
+            f"{torch._dynamo.config.cache_size_limit} to accomodate prompt size of "
+            f"{_prompt_size} and decode tokens of {args.max_new_tokens}"
+        )


this is only needed if compile_dynamic is disabled, can we gate it?

ani300 · 2025-04-28T14:33:49Z

aiu_fms_testing_utils/utils/aiu_setup.py

+    os.environ.setdefault("SENCORES", "32")
+    os.environ.setdefault("SENCORELETS", "2")
+    os.environ.setdefault("DATA_PREC", "fp16")
+    os.environ.setdefault("FLEX_OVERWRITE_NMB_FRAME", "1")


I think some of these are already set by default on the e2e_stable image, can we check and remove the ones we don't need anymore?

ani300 · 2025-04-28T14:34:11Z

aiu_fms_testing_utils/utils/aiu_setup.py

+    os.environ.setdefault("FLEX_OVERWRITE_NMB_FRAME", "1")
+    os.environ.setdefault("DTCOMPILER_KEEP_EXPORT", "true")
+
+    os.environ.setdefault("COMPILATION_MODE", "offline_decoder")


this one is only needed for decoder models, for roberta it will probably make it not work

ani300 · 2025-04-28T14:34:51Z

aiu_fms_testing_utils/utils/aiu_setup.py

+            print("must set AIU_WORLD_RANK_0")
+            exit()
+        os.environ.setdefault("FLEX_COMPUTE", "SENTIENT")
+        os.environ.setdefault("FLEX_DEVICE", "VFIO")


I think VFIO is now PF or VF, depending on the AIU setup, we can set PF by default as most cards are running in PF mode

ani300 · 2025-04-28T14:36:51Z