[PJRT:GPU] Treat GPU collective memory space as device memory space

[zhenying-liu]

This is a regression fix when using --xla_gpu_enable_nccl_user_buffers=true.
Return device memory space when collective memory space is used as an output on GPU.

[NaiyerRizz]

Hi @zhenying-liu
Seems like there are some conflicts in the branch, Can you please resolve those.
Thank you

[zhenying-liu]

Hi @zhenying-liu Seems like there are some conflicts in the branch, Can you please resolve those. Thank you

Thanks. Resolved conflicts.

[cheshire]

Any chance we could separate out XLA and PJRT changes?

I'm somewhat new to that code, but hardcoding memory spaces in layout.h seems strange, is that the usual approach?

[cheshire]

Is it possible to describe in general terms, not in specific fixes, what the change is doing? Is it fixing a crash? Is it performance improvement? If it's a bullet list of five things, maybe it should be 5 PRs?

[zhenying-liu]

Any chance we could separate out XLA and PJRT changes?

I'm somewhat new to that code, but hardcoding memory spaces in layout.h seems strange, is that the usual approach?

It's defined in XLA. It's harder to separate XLA and PJRT. But I can try my best to split.
I think defining memory spaces in layout.h is the best way that I can think of. Otherwise, defining it in different files could cause the conflicts. Currently both fast memory space and collect memory space has the value 1.

[zhenying-liu]

Is it possible to describe in general terms, not in specific fixes, what the change is doing? Is it fixing a crash? Is it performance improvement? If it's a bullet list of five things, maybe it should be 5 PRs?

There is a crash when using NCCL user buffer: --xla_gpu_enable_nccl_user_buffers=true. The error message is: "INVALID_ARGUMENT: Unexpected memory space 1 in output layout". It happened in pjrt_stream_executor_client.cc:MemoryKindFromSimpleShape().
Jax doesn't support any memory kind other than device and pinned_host. In this case, the output memory kind is neither device nor pinned_host. Our strategy is to treat CollectiveMemorySpace as device memory space to please Jax. But inside XLA, we keep using collective memory space to call nccl memory allocator to handle allocation/deallocation. The existing memory space color for collective memory space is 1, conflicting with Layout::kGenericFastMemorySpace. So we define a constant kCollectiveMemorySpace in layout.h with a unique value 6, and use this constant whenever it is needed.

[nouiz]

Is it possible to describe in general terms, not in specific fixes, what the change is doing? Is it fixing a crash? Is it performance improvement? If it's a bullet list of five things, maybe it should be 5 PRs?

There is a crash when using NCCL user buffer: --xla_gpu_enable_nccl_user_buffers=true. The error message is: "INVALID_ARGUMENT: Unexpected memory space 1 in output layout". It happened in pjrt_stream_executor_client.cc:MemoryKindFromSimpleShape(). Jax doesn't support any memory kind other than device and pinned_host. In this case, the output memory kind is neither device nor pinned_host. Our strategy is to treat CollectiveMemorySpace as device memory space to please Jax. But inside XLA, we keep using collective memory space to call nccl memory allocator to handle allocation/deallocation. The existing memory space color for collective memory space is 1, conflicting with Layout::kGenericFastMemorySpace. So we define a constant kCollectiveMemorySpace in layout.h with a unique value 6, and use this constant whenever it is needed.

This is a regression fix when using --xla_gpu_enable_nccl_user_buffers=true
@zhenying-liu , can you update the description to tell that?

[nouiz]

Is it possible to describe in general terms, not in specific fixes, what the change is doing? Is it fixing a crash? Is it performance improvement? If it's a bullet list of five things, maybe it should be 5 PRs?

This is a regression fix. Google JAX team want to make a releases, so merging this soon would be great.
@hawkinsp approved the changes, so all PJRT part should be good.
Splitting into multiple from outside would take extra times and this compete with doing a release of JAX very soon. So I let you @cheshire and @hawkinsp discuss that. Having regression in public version isn't great.

[cheshire]

Does this enum have to live in layout.h ? Maybe we could move it to a GPU-specific file? After all , GPU has to not collide only with GPU. Also having to recompiling XLA on any changes to this enum does not sound great.

[jaro-sevcik]

@cheshire At the same time, it would be good if the constants would not collide, so I would prefer to have all the constants in one place. How about creating a new header file just for the memory space constants?

[cheshire]

@jaro-sevcik Actually why not? If other device D has different memory spaces using same numbers, why does it matter? If they are allowed to collide, it means two backends can develop separately, which in general seems good? I guess here the "implementation details" need to be present at the API boundary, which makes it difficult.

Looking at the usage, conceptually, it seems like the design could be improved. The buffers allocated by XLA are internal, and should not be exposed to the API boundary (PjRT), but here they are. I guess ideally that logic should be moved out of PjRT into XLA.

[jaro-sevcik]

@cheshire Conceptually what you suggest sounds consistent. I would be concerned that someone takes a layout from an array, tries to use for some other device and gets surprising results without any warning, but perhaps that's already not possible because one cannot really sneak in memory spaces in layout.

In practice, this PR is holding up a JAX release, so we should decide what we want soon. Are you proposing to undo the move/rename of the kCollectiveMemorySpaceColor constant and change of its value to 6 (without worrying about the collective space id 1 colliding with kGenericFastMemorySpace) and just keep the changes to xla/pjrt/pjrt_stream_executor_client.cc and to xla/pjrt/gpu/se_gpu_pjrt_client_test.cc?

[cheshire]

In practice, this PR is holding up a JAX release

Is it? What I understand from @hawkinsp that some combination of flags became slower, which is not the property release has to uphold?

[zhenying-liu]

@cheshire I have the smallest changes to xla/pjrt/pjrt_stream_executor_client.cc and xla/pjrt/gpu/se_gpu_pjrt_client_test.cc ready. With that one, I have to use kGenericFastMemorySpace for collective space id 1. If so, we can fix the regression now and do the clean up the definition later.

[zhenying-liu]

@cheshire Please take another look at the code review as I changed the code as mentioned above.

[cheshire]

What happens if we execute it without the changes in this PR?

[zhenying-liu]

It will have the assertion failure in pjrt_stream_execution_client.cc:MemoryKindFromSimpleShape():
INVALID_ARGUMENT: Unexpected memory space 1 in output layout

[nouiz]

This approved PR is failing in Google internal import.
Who can look at it? @xla-rotation

[akuegel]

Merging this PR is blocked on a compile error. I wonder whether this actually compiles on your end?

[akuegel]

void function 'TestBody' should not return a value [-Wreturn-mismatch]
1239 | TF_ASSIGN_OR_RETURN(

[zhenying-liu]

Thanks for pointing this out. Sorry for the mistake. Fixed.