clean up ruff check

Signed-off-by: kcirred <[email protected]>

Author: kcirred

aiu_fms_testing_utils/utils/model_setup.py

@@ -47,6 +47,8 @@ def get_device(args: argparse.Namespace) -> torch.device:
         device = torch.device(args.device_type, local_rank)
         torch.cuda.set_device(device)
     elif args.is_aiu_backend:
+        from torch_sendnn import torch_sendnn  # noqa: F401
+
         if args.distributed:
             aiu_setup.aiu_dist_setup(
                 distributed.get_rank(),

aiu_fms_testing_utils/utils/paged.py

@@ -3,6 +3,7 @@
 import time
 from typing import Any, Callable, List, MutableMapping, Optional, Tuple, Union
 import torch
+import fms.utils.spyre.paged  # noqa
 
 
 def adjust_inputs_to_batch(input_ids: torch.Tensor, **extra_kwargs):

scripts/generate_metrics.py

@@ -175,12 +175,11 @@ def find_eos_index(reference_tokens, eos_token_id):
     return result
 
 
-def filter_before_eos(level_metric, filter_indexes):
+def filter_before_eos(metrics, filter_indexes):
     from itertools import groupby
 
     filtered_results = [
-        list(g)[: filter_indexes[k]]
-        for k, g in groupby(level_metric, key=lambda x: x[0])
+        list(g)[: filter_indexes[k]] for k, g in groupby(metrics, key=lambda x: x[0])
     ]
     return [item for sublist in filtered_results for item in sublist]
 
@@ -197,10 +196,10 @@ def __prepare_inputs(batch_size, seq_length, tokenizer, seed=0):
     return input_ids, padding_kwargs
 
 
-def write_csv(metric, path, metric_name):
+def write_csv(metrics, path, metric_name):
     with open(path, "w") as f:
         f.write(f"{metric_name}\n")
-        for t in metric:
+        for t in metrics:
             f.write(f"{t[2].item()}\n")
         f.close()
 

scripts/inference.py

@@ -252,16 +252,20 @@
     from fms.utils.generation import generate
 
 if "fp8" in attn_name:
-    pass
+    import fms_mo.aiu_addons.fp8.fp8_attn  # noqa: F401
 
 if args.quantization == "gptq":
     if "aiu" in args.device_type:
         try:
+            from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear  # noqa
+
             print("Loaded `aiu_addons` functionalities")
         except ImportError:
             raise ImportError("Failed to import GPTQ addons from fms-mo.")
 elif args.quantization == "int8":
     try:
+        from fms_mo.aiu_addons.i8i8 import i8i8_aiu_adapter, i8i8_aiu_linear  # noqa
+
         print("Loaded `aiu_addons` functionalities")
     except ImportError:
         raise ImportError("Failed to import INT8 addons from fms-mo.")
@@ -301,6 +305,8 @@
     device = torch.device(args.device_type, local_rank)
     torch.cuda.set_device(device)
 elif is_aiu_backend:
+    from torch_sendnn import torch_sendnn  # noqa
+
     if not args.distributed:
         aiu_setup.aiu_setup(rank, world_size)
 

scripts/roberta.py

@@ -16,6 +16,7 @@
 from fms.utils.tp_wrapping import apply_tp
 
 # Import AIU Libraries
+from torch_sendnn import torch_sendnn  # noqa
 
 
 # ==============================================================

scripts/small-toy.py

@@ -297,6 +297,7 @@
     aiu_setup.aiu_dist_setup(dist.get_rank(), dist.get_world_size())
 
 # Always initialize AIU in this script
+from torch_sendnn import torch_sendnn  # noqa
 
 if not args.distributed:
     aiu_setup.aiu_setup(rank, world_size)

scripts/validation.py

@@ -315,12 +315,11 @@ def __find_eos_index(reference_tokens, eos_token_id, seq_length, max_new_tokens)
     return result
 
 
-def __filter_before_eos(level_metric, filter_indexes):
+def __filter_before_eos(metrics, filter_indexes):
     from itertools import groupby
 
     filtered_results = [
-        list(g)[: filter_indexes[k]]
-        for k, g in groupby(level_metric, key=lambda x: x[0])
+        list(g)[: filter_indexes[k]] for k, g in groupby(metrics, key=lambda x: x[0])
     ]
     return [item for sublist in filtered_results for item in sublist]