Update GPT Fast handler to add option to get GPU ID from context (#2872)

* Update GPT Fast handler to add option to get GPU ID from context * Fix lint checks * Fix lint checks * Fix lint checks * Update handler --------- Co-authored-by: Ubuntu <[email protected]>
pytorch · Jan 4, 2024 · 39ea211 · 39ea211
1 parent 2eae8e5
commit 39ea211
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/large_models/gpt_fast/handler.py b/examples/large_models/gpt_fast/handler.py
@@ -44,9 +44,13 @@ def __init__(self):
 
     def initialize(self, ctx):
         self.context = ctx
+        properties = self.context.system_properties
+        gpu_id = properties.get("gpu_id")
+        if gpu_id is not None and int(gpu_id) < 0:
+            raise ValueError("Invalid gpu_id")
         rank = maybe_init_dist()
 
-        self.local_rank = rank if rank is not None else 0
+        self.local_rank = rank if rank is not None else int(gpu_id)
 
         if torch.cuda.is_available():
             self.map_location = "cuda"