fix bf16 error

examples/offline_inference.py

@@ -12,7 +12,7 @@ if __name__ == '__main__':
     sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
     # Create an LLM.
-    llm = LLM(model="facebook/opt-125m",trust_remote_code=True, dtype="float16", enforce_eager=True)
+    llm = LLM(model="facebook/opt-125m",trust_remote_code=True, dtype="float16", enforce_eager=False)
     # Generate texts from the prompts. The output is a list of RequestOutput objects
     # that contain the prompt, generated text, and other information.
     outputs = llm.generate(prompts, sampling_params)

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -12,8 +12,10 @@ from vllm.distributed.device_communicators.custom_all_reduce_utils import (
 from vllm.distributed.parallel_state import (
     get_local_rank, get_tensor_model_parallel_cpu_group, is_in_the_same_node)
 from vllm.logger import init_logger
+from vllm.utils import is_hip
 
 try:
+    if (not is_hip()):
         import pynvml
 
         # Simulate ImportError if custom_ar ops are not supported.
@@ -29,6 +31,16 @@ try:
                 yield
             finally:
                 pynvml.nvmlShutdown()
+    else:
+        custom_ar = False
+        pynvml = None
+
+        @contextmanager
+        def _nvml():
+            try:
+                yield
+            finally:
+                pass
 
 except ImportError:
     # For AMD GPUs