[fix] export failure with CUDA driver < 526 and pynvml>=11.5.0

- There is a bug that was fixed in the 526 driver release. For older driver versions the recommendation is to downgrade the pynvml version to 11.4.0 and use 11.5.0 only for drivers after 526.

Uses the legacy pynvml memory usage function even with pynvml 11.5.0 if the driver version is older than 526.

Mentioned in the issue as well: NVIDIA#808 (comment)

Author: hemant-co

tensorrt_llm/profiler.py

@@ -44,15 +44,19 @@
         "A required package 'pynvml' is not installed. Will not "
         "monitor the device memory usages. Please install the package "
         "first, e.g, 'pip install pynvml>=11.5.0'.")
-elif pynvml.__version__ < '11.5.0':
-    logger.warning(f'Found pynvml=={pynvml.__version__}. Please use '
-                   f'pynvml>=11.5.0 to get accurate memory usage')
-    # Support legacy pynvml. Note that an old API could return
-    # wrong GPU memory usage.
-    _device_get_memory_info_fn = pynvml.nvmlDeviceGetMemoryInfo
 else:
-    _device_get_memory_info_fn = partial(pynvml.nvmlDeviceGetMemoryInfo,
-                                         version=pynvml.nvmlMemory_v2)
+    pynvml.nvmlInit()
+    driver_version = pynvml.nvmlSystemGetDriverVersion()
+    if pynvml.__version__ < '11.5.0' or driver_version < '526':
+        logger.warning(f'Found pynvml=={pynvml.__version__} and cuda driver version {driver_version}.'
+                        ' Please use pynvml>=11.5.0 and cuda driver>=526 to get accurate memory usage')
+        # Support legacy pynvml. Note that an old API could return
+        # wrong GPU memory usage.
+        _device_get_memory_info_fn = pynvml.nvmlDeviceGetMemoryInfo
+    else:
+        _device_get_memory_info_fn = partial(pynvml.nvmlDeviceGetMemoryInfo,
+                                            version=pynvml.nvmlMemory_v2)
+    pynvml.nvmlShutdown()
 
 
 class Timer: