ecmwf / anemoi-training

Apache License 2.0
17 stars 16 forks source link

Better support for AMD GPUS #77

Open anaprietonem opened 1 month ago

anaprietonem commented 1 month ago

Is your feature request related to a problem? Please describe.

In terms of tracking system metrics from a profiler and mlflow perspective, the current code lacks some feature to better support other hardware different from Nvidia-GPUS. Since the only package we use to do this is pynvml and that is limited to NVIDIA GPUs.

Describe the solution you'd like

We could improve this, and at least provide better support to AMD GPUs since there is an open source package called pyrsmi (developed by ROCM), which does the same as pynvml, but for AMD ROCM hardware. We could define a custom SystemMetrics Monitor that can handle many hadwares. (see comments for potential implementation). This will work out of the box, with the same config settings as you would run your mlflow. Lastly, this will only monitor one node. I.e 8 amd gpus on the same node, and not across nodes, since we assume that the memory consumption, speed, etc.. would be the same for all gpus except "master gpu".

Describe alternatives you've considered

No response

Additional context

This solution was originally suggested by @einrone so many thanks for this!

Organisation

No response

anaprietonem commented 1 month ago

See reference implementation provided by Aram for AIFS monitor_system_metric.txt

mpvginde commented 1 month ago

Hi @anaprietonem, FYI: right now we are using this PR: https://github.com/mlflow/mlflow/pull/12694 from the people of MetNO, to do some basic monitoring of AMD GPUs.

HCookie commented 4 days ago

Was this addressed by #147 ?