Open rohitdwivedula opened 1 month ago
cc @mwootton
I’m not too familiar with the AMD plugin code but dug in a bit. Some of the missing bits might be there in later versions, is it possible you can try out pytorch nightly ?
Here is what I found-
1) Device properties: This is already supported for HIP/AMD devices. https://github.com/pytorch/kineto/pull/927 (May 2024)
2) Memcpy bytes : Looks like there is a “size” field in memcpy events that should have this- (The metadataJson() function adds the extra fields you see in the trace) https://github.com/pytorch/kineto/blob/main/libkineto/src/RoctracerActivity_inl.h#L161-L170
3) I think correlation ID should be present as per latest code, on all runtime operations. We may need to investigate this in detail though (See metadataJson() again) https://github.com/pytorch/kineto/blob/main/libkineto/src/RoctracerActivity_inl.h#L187-L190
I heard that some collective events are missing correlation ID. But most likely these fixes will be in the AMD ROCTracer so @mwootton / @hongxiayang might be able to help out.
🐛 Describe the bug
Summary: Device information, correlation IDs, and the
bytes
field are missing intorch.profiler.profile
JSON dumps when this profiling class is used on AMD GPUs.test_kineto.py contains a script that (i) creates a tensor on each GPU, (ii) performs an allreduce.SUM operation across all ranks, while (iii) using the Kineto profiler to capture all CPU and CUDA events and (iv) dump the trace to JSON. When you run this script on AMD GPUs, the JSON is missing some attributes and information that are available when running the same profiler on Nvidia GPUs.
I ran the
test_kineto.py
script on (a) two Nvidia A100 80GB GPUs, and (b) two Instinct MI210 GPUs; using Python 3.9.19 and PyTorch 2. The output JSON files I got from these runs are available here (AMD Instinct MI 210) and here (for Nvidia A100).1. Missing
bytes
fieldFor example, this is what a
gpu_memcpy
operation in the Nvidia trace looks like:On AMD GPUs, a similar
gpu_memcpy
item in the JSON file looks like this:The
bytes
field is missing in the latter. The same problem exists for allgpu_memset
operations as well.2. Missing device information
The A100 Kineto trace contains a field
deviceProperties
- this is completely absent/missing from the AMD trace.3. Correlation ID
Each entry in the JSON generated on Nvidia GPUs contains a correlation ID:
This
correlation
field is not present in any of the events in the AMD GPU dump. For example, the MLCommons/chakra repository uses thetorch.profiler.profile
class, and assumes that thiscorrelation
field is present in the JSON in multiple places (e.g. here)Versions
Used the
rocm/pytorch:latest
docker image (image id:b80124b96134
) from DockerHub. Output ofcollect_env.py
:cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @robieta @chaekit @aaronenyeshi @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise