Convert the stagger enter/leave into a proper contextlib function

Signed-off-by: Joshua Hursey <jhursey@us.ibm.com>

Author: jjhursey

aiu_fms_testing_utils/utils/__init__.py

@@ -10,30 +10,24 @@
 import json
 import random
 import math
+import contextlib
 
-def stagger_enter(limit: int):
+@contextlib.contextmanager
+def stagger_region(limit: int):
     """
-    Limit the number of concurrent processes into this section of code.
-    Processes return from this funciton when they are allowed to enter the section of code.
-    Must be paired with another call to stagger_leave() when exiting the section of code.
+    Limit the number of concurrent processes into this region of code.
+    Processes yield from this function when they are allowed to enter the region of code.
+    Processes return from this function when all of the processes have completed the region of code.
 
-    :param limit: Number of concurrent processes allowed in the code section if > 0.
+    :param limit: Number of concurrent processes allowed in the code region if > 0.
     """
     if limit > 0 and limit != world_size:
         for _set in range( math.ceil(world_size / float(limit)) ):
             if rank < (_set+1)*limit:
                 break
             torch.distributed.barrier()
         dprint(f"Stagger: Enter (Set: {_set+1} of {math.ceil(world_size / float(limit))})")
-
-def stagger_leave(limit: int):
-    """
-    Leave a section code where the number of concurrent processes is limited.
-    Processes return from this funciton when all of the processes have completed the section of code.
-    Must be paired with another call to stagger_enter() when entering the section of code.
-
-    :param limit: Number of concurrent processes allowed in the code section if > 0.
-    """
+    yield {}
     if limit > 0 and limit != world_size:
         for _set in range( math.ceil(world_size / float(limit)) ):
             if rank >= (_set+1)*limit:
@@ -49,15 +43,12 @@ def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int,
     if compile_dynamic_sendnn:
         max_new_tokens_warmup = 2
 
-    stagger_enter(stagger_update_lazyhandle)
-
-    pt_compile_model_time = time.time()
-    with torch_sendnn.warmup_mode():
-        generate(model, input_ids, max_new_tokens=max_new_tokens_warmup, max_seq_len=model.config.max_expected_seq_len, use_cache=True, do_sample=False, contiguous_cache=True, extra_kwargs=extra_kwargs)
-    pt_compile_model_time = time.time() - pt_compile_model_time
-    dprint(f"PT compile complete, took {pt_compile_model_time:.3f}s")
-
-    stagger_leave(stagger_update_lazyhandle)
+    with stagger_region(stagger_update_lazyhandle) as _s:
+        pt_compile_model_time = time.time()
+        with torch_sendnn.warmup_mode():
+            generate(model, input_ids, max_new_tokens=max_new_tokens_warmup, max_seq_len=model.config.max_expected_seq_len, use_cache=True, do_sample=False, contiguous_cache=True, extra_kwargs=extra_kwargs)
+        pt_compile_model_time = time.time() - pt_compile_model_time
+        dprint(f"PT compile complete, took {pt_compile_model_time:.3f}s")
 
 def ids_for_prompt(prompt, tokenizer):
     tokens = tokenizer.tokenize(prompt)

scripts/inference.py

@@ -12,7 +12,7 @@
 import math
 
 # Third Party
-from aiu_fms_testing_utils.utils import aiu_setup, warmup_model, stagger_enter, stagger_leave
+from aiu_fms_testing_utils.utils import aiu_setup, warmup_model, stagger_region
 from aiu_fms_testing_utils.utils.aiu_setup import dprint, rank, local_rank, world_size
 import numpy as np
 import torch
@@ -464,38 +464,35 @@ def select_int8_module(
 dprint(f"data_type={default_dtype}")
 dprint("="*60 + "\n")
 
-stagger_enter(args.stagger_load)
-
-model = get_model(
-    args.architecture,
-    args.variant,
-    model_path=args.model_path,
-    device_type="cpu" if is_aiu_backend else args.device_type,
-    data_type=default_dtype,
-    source=args.model_source,
-    distributed_strategy=distr_param,
-    group=dist.group.WORLD,
-    linear_config=linear_config,
-    fused_weights=fused_weights,
-)
-
-if args.quantization in ["gptq", "int8"]:
-    if rank == 0 and args.verbose > 0:
-        dprint("PARAMS:\n" + "\n".join(f"{k:60} {str(v.dtype):15} {str(v.device):10} {list(v.size())}" for k,v in model.named_parameters()))
-        dprint("BUFFERS:\n" + "\n".join(f"{k:60} {str(v.dtype):15} {str(v.device):10} {list(v.size())}" for k,v in model.named_buffers()))
+with stagger_region(args.stagger_load) as _s:
+    model = get_model(
+        args.architecture,
+        args.variant,
+        model_path=args.model_path,
+        device_type="cpu" if is_aiu_backend else args.device_type,
+        data_type=default_dtype,
+        source=args.model_source,
+        distributed_strategy=distr_param,
+        group=dist.group.WORLD,
+        linear_config=linear_config,
+        fused_weights=fused_weights,
+    )
+
+    if args.quantization in ["gptq", "int8"]:
+        if rank == 0 and args.verbose > 0:
+            dprint("PARAMS:\n" + "\n".join(f"{k:60} {str(v.dtype):15} {str(v.device):10} {list(v.size())}" for k,v in model.named_parameters()))
+            dprint("BUFFERS:\n" + "\n".join(f"{k:60} {str(v.dtype):15} {str(v.device):10} {list(v.size())}" for k,v in model.named_buffers()))
+            dprint("="*60 + "\n")
+        if args.architecture == "llama":
+            dprint("[NOTE] In Llama models, it's OK for bias and rotary embeddings to be marked as unused keys.")
+        dprint(model)
         dprint("="*60 + "\n")
-    if args.architecture == "llama":
-        dprint("[NOTE] In Llama models, it's OK for bias and rotary embeddings to be marked as unused keys.")
-    dprint(model)
-    dprint("="*60 + "\n")
-
-tokenizer = tokenizers.get_tokenizer(args.tokenizer)
-model.eval()
-torch.set_grad_enabled(False)
-loading_model_time = time.time() - loading_model_time
-dprint(f"loading complete, took {loading_model_time:.3f}s")
-
-stagger_leave(args.stagger_load)
+
+    tokenizer = tokenizers.get_tokenizer(args.tokenizer)
+    model.eval()
+    torch.set_grad_enabled(False)
+    loading_model_time = time.time() - loading_model_time
+    dprint(f"loading complete, took {loading_model_time:.3f}s")
 
 if args.compile:
     dprint("compiling model")