[JAX] Custom Op Workspace Tensors from XLA Buffers #532

denera · 2023-11-22T21:01:51Z

Previously, custom ops allocated their workspace tensors via WorkspaceManager in jax/csrc/utils.h, which relied on direct cudaMalloc().

This PR removes the WorkspaceManager and exposes new Python interfaces to determine workspaces sizes in the custom-op primitives in order to request the same workspace memory allocation from XLA.

nouiz · 2023-11-22T21:51:38Z

@mingxu1067 to review.

transformer_engine/jax/csrc/modules.h

transformer_engine/jax/csrc/modules.cpp

mingxu1067 · 2023-11-23T02:48:28Z

transformer_engine/jax/cpp_extensions.py

        out_types = [
            ir.RankedTensorType.get(out_shape, output_type),
            ir.RankedTensorType.get(batch_shape, ir_mu_dtype),
            ir.RankedTensorType.get(batch_shape, ir_rsigma_dtype),
+            ir.RankedTensorType.get(work_meta[0], te_dtype_to_ir_dtype(work_meta[1])),


These changes make the number of outputs from lowering is different with abstract. I assume there would be a XLA assertion, could you test this to see if any potential issues?

Besides, these workspace tensors are never used, then it might be removed during XLA compiling. We might have to test this also.

The workspace tensors I added appear to be used in the C++ custom ops and NVTE kernels they call. They were being allocated with cudaMalloc in WorkspaceManager, which is what we want to replace with this PR. Did I misread/misunderstand something about this?

It is entirely possible I unintentionally skipped over additional changes needed in the abstract (or elsewhere) to make this work correctly. Do we need dummy outputs in the abstract to make sure we don't lose the workspaces to XLA optimization?

Yes the purpose of this PR is to remove cudaMalloc in WorkspaceManager.

I am just thinking there might be two issues

The mismatching number of outputs between lowering and abstract might trigger some assertion. If so, then we have to seperate abstract for inner_p and outer_p, then add dummy outouts to abstract of inner_p.

Ununsed workspace tensors might be removed via some XLA's memory optimization pass. But after carefully thinking, we might not to worry about this, since XLA has less knowledge to custom calls, therefore it usually skips those optimization to custom calls

If JAX doesn't raise an error due to different number of inputs/outputs, please raise this and give a repro (don't need to be a minimal one). I can look at adding the missing check in JAX.

XLA and all compilers can remove one full operation if all its outputs aren't used and is a side-effect free operation. But it can't remove only part of one operations outputs unless it is know that this instruction support this. So yes, XLA can't do that for custom_call, but in practice, it doesn't do that for native operations too . I don't know any such exception in XLA.

zlsh80826 · 2023-11-23T04:29:17Z

Would it be better to not expose the workspace data type to the cpp_extensions.py? We can do something like in JAX that only expose the required workspace size to python side and use somthing like ir.IntegerType.get_signless(8) as the data type. (jaxlib workspace reference)

Another thing is that we can calculate the workspace size in modules.cpp before returning back to the cpp_extensions.py, so that we can return only a single size_t instead of a vector.

transformer_engine/jax/csrc/modules.h

transformer_engine/jax/csrc/modules.cpp

transformer_engine/jax/csrc/modules.h

denera · 2023-12-15T21:57:45Z

/te-ci jax

denera · 2023-12-19T23:08:42Z

/te-ci jax

denera · 2023-12-20T19:58:18Z

@nouiz @mingxu1067 @zlsh80826 This is passing LayerNorm and FusedAttn tests on my end. I'm running the CI but this should be ready for final review now. Thanks!

denera · 2023-12-20T20:40:20Z

/te-ci jax

nouiz

LGTM at high level.

zlsh80826 · 2023-12-24T06:00:30Z

transformer_engine/jax/cpp_extensions.py

        out_types = [
            ir.RankedTensorType.get(out_shape, output_type),
            ir.RankedTensorType.get(batch_shape, ir_mu_dtype),
            ir.RankedTensorType.get(batch_shape, ir_rsigma_dtype),
+            ir.RankedTensorType.get(wkspace_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype)),
+            ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype))


Suggested change

ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype))

ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(barrier_aval.dtype))

zlsh80826 · 2023-12-26T16:12:37Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

zlsh80826 · 2023-12-26T16:18:31Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)
+        dgamma_part_size = reduce(operator.mul, dgamma_part_aval.shape)
+        dbeta_part_size = reduce(operator.mul, dbeta_part_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

dgamma_part_size = reduce(operator.mul, dgamma_part_aval.shape)

dbeta_part_size = reduce(operator.mul, dbeta_part_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

dgamma_part_size = dgamma_part_aval.size

dbeta_part_size = dbeta_part_aval.size

zlsh80826 · 2023-12-26T16:21:48Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

zlsh80826 · 2023-12-26T16:22:47Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)
+        dgamma_part_size = reduce(operator.mul, dgamma_part_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

dgamma_part_size = reduce(operator.mul, dgamma_part_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

dgamma_part_size = dgamma_part_aval.size

zlsh80826 · 2023-12-26T16:25:56Z

transformer_engine/jax/cpp_extensions.py

        out_types = [
            ir.RankedTensorType.get(x_shape, x_type.element_type),
            ir.RankedTensorType.get(g_shape, g_type.element_type),
            ir.RankedTensorType.get(b_shape, b_type.element_type),
+            ir.RankedTensorType.get(wkspace_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype)),
+            ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(barrier_aval.dtype)),
+            ir.RankedTensorType.get(dgamma_part_aval.shape,
+                                    jax_dtype_to_ir_dtype(dgamma_part_aval.dtype)),
+            ir.RankedTensorType.get(dbeta_part_aval.shape,
+                                    jax_dtype_to_ir_dtype(dbeta_part_aval.dtype))
        ]


This would work, but the current enumerated list is also fine to me.

Suggested change

out_types = [

ir.RankedTensorType.get(x_shape, x_type.element_type),

ir.RankedTensorType.get(g_shape, g_type.element_type),

ir.RankedTensorType.get(b_shape, b_type.element_type),

ir.RankedTensorType.get(wkspace_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype)),

ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(barrier_aval.dtype)),

ir.RankedTensorType.get(dgamma_part_aval.shape,

jax_dtype_to_ir_dtype(dgamma_part_aval.dtype)),

ir.RankedTensorType.get(dbeta_part_aval.shape,

jax_dtype_to_ir_dtype(dbeta_part_aval.dtype))

]

out_types = [

ir.RankedTensorType.get(output.shape, mlir.dtype_to_ir_type(output.dtype))

for output in ctx.avals_out

]

zlsh80826 · 2023-12-26T16:29:10Z

transformer_engine/jax/cpp_extensions.py

+        batch_size = reduce(operator.mul, batch_shape)
+
+        wkspace_aval = ctx.avals_out[-1]
+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

wkspace_size = wkspace_aval.size

zlsh80826 · 2023-12-26T16:30:54Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

zlsh80826 · 2023-12-26T16:31:45Z

transformer_engine/jax/cpp_extensions.py

+        wkspace_size = reduce(operator.mul, wkspace_aval.shape)
+        barrier_size = reduce(operator.mul, barrier_aval.shape)


Suggested change

wkspace_size = reduce(operator.mul, wkspace_aval.shape)

barrier_size = reduce(operator.mul, barrier_aval.shape)

wkspace_size = wkspace_aval.size

barrier_size = barrier_aval.size

denera · 2024-01-17T17:56:03Z

/te-ci jax

denera · 2024-01-22T15:44:39Z

/te-ci jax

denera · 2024-01-22T15:46:26Z

@cyanguwa There are some (minor) changes to the common fused attn kernels here and it would be great to get your feedback on them to make sure there won't be unintended consequences. Thanks!

…equest buffers from XLA for their workspace tensors. Signed-off-by: Alp Dener <adener@nvidia.com>

Signed-off-by: Alp Dener <adener@nvidia.com>

…eductions in primitives Signed-off-by: Alp Dener <adener@nvidia.com>

Signed-off-by: Alp Dener <adener@nvidia.com>

…ecuting dummy kernel call with nullptr stream Signed-off-by: Alp Dener <adener@nvidia.com>

Signed-off-by: Alp Dener <adener@nvidia.com>

* Removed cudaMalloc/WorkspaceManager in JAX csrc. JAX custom ops now request buffers from XLA for their workspace tensors. Signed-off-by: Alp Dener <adener@nvidia.com> * removed unused GEMM C++ API in TE-JAX Signed-off-by: Alp Dener <adener@nvidia.com> * fixed typo in layernorm_geglu_fp8_mlp and removed unnecessary shape reductions in primitives Signed-off-by: Alp Dener <adener@nvidia.com> * fixed import order for linting Signed-off-by: Alp Dener <adener@nvidia.com> * fixed custom op errors due to incorrect static arg nums in JAX jit Signed-off-by: Alp Dener <adener@nvidia.com> * shifted cudnnSetStream further down the kernel to avoid error when executing dummy kernel call with nullptr stream Signed-off-by: Alp Dener <adener@nvidia.com> * fixed linting errors for blank lines Signed-off-by: Alp Dener <adener@nvidia.com> --------- Signed-off-by: Alp Dener <adener@nvidia.com> Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

denera requested review from nouiz, mingxu1067 and zlsh80826 November 22, 2023 21:01

denera self-assigned this Nov 22, 2023

mingxu1067 reviewed Nov 23, 2023

View reviewed changes

transformer_engine/jax/csrc/modules.h Outdated Show resolved Hide resolved

mingxu1067 reviewed Nov 23, 2023

View reviewed changes

transformer_engine/jax/csrc/modules.cpp Outdated Show resolved Hide resolved

mingxu1067 reviewed Nov 23, 2023

View reviewed changes

transformer_engine/jax/csrc/modules.cpp Show resolved Hide resolved

mingxu1067 reviewed Nov 23, 2023

View reviewed changes

transformer_engine/jax/csrc/modules.cpp Outdated Show resolved Hide resolved

mingxu1067 reviewed Nov 23, 2023

View reviewed changes

zlsh80826 reviewed Nov 23, 2023

View reviewed changes

transformer_engine/jax/csrc/modules.h Outdated Show resolved Hide resolved

transformer_engine/jax/csrc/modules.cpp Outdated Show resolved Hide resolved

transformer_engine/jax/csrc/modules.h Outdated Show resolved Hide resolved

denera mentioned this pull request Nov 30, 2023

[JAX] Support layernorm/rmsnorm sm_margin control through environment variable #520

Merged

denera force-pushed the jax-gemm-workspace-xla-buffer branch 3 times, most recently from 9129b8e to 830b67f Compare December 15, 2023 21:39

denera force-pushed the jax-gemm-workspace-xla-buffer branch from 830b67f to 04b14b2 Compare December 15, 2023 23:57

denera marked this pull request as ready for review December 15, 2023 23:57

denera force-pushed the jax-gemm-workspace-xla-buffer branch 2 times, most recently from cfe0dda to a8c3fdc Compare December 19, 2023 22:02

denera force-pushed the jax-gemm-workspace-xla-buffer branch 2 times, most recently from 6bc5b3f to 1e6d94a Compare December 20, 2023 19:53

denera requested review from mingxu1067 and zlsh80826 December 20, 2023 20:48

nouiz approved these changes Dec 20, 2023

View reviewed changes

zlsh80826 reviewed Dec 26, 2023

View reviewed changes

denera force-pushed the jax-gemm-workspace-xla-buffer branch 6 times, most recently from 12bd429 to 088a79b Compare January 17, 2024 17:41

denera force-pushed the jax-gemm-workspace-xla-buffer branch from 088a79b to d5918b1 Compare January 22, 2024 15:44

denera requested a review from cyanguwa January 22, 2024 15:46

denera added 7 commits January 23, 2024 21:45

Removed cudaMalloc/WorkspaceManager in JAX csrc. JAX custom ops now r…

0dad593

…equest buffers from XLA for their workspace tensors. Signed-off-by: Alp Dener <adener@nvidia.com>

removed unused GEMM C++ API in TE-JAX

ed7048d

Signed-off-by: Alp Dener <adener@nvidia.com>

fixed typo in layernorm_geglu_fp8_mlp and removed unnecessary shape r…

67c0c4f

…eductions in primitives Signed-off-by: Alp Dener <adener@nvidia.com>

fixed import order for linting

ff0c1ab

Signed-off-by: Alp Dener <adener@nvidia.com>

fixed custom op errors due to incorrect static arg nums in JAX jit

558bc80

Signed-off-by: Alp Dener <adener@nvidia.com>

shifted cudnnSetStream further down the kernel to avoid error when ex…

dd3a91c

…ecuting dummy kernel call with nullptr stream Signed-off-by: Alp Dener <adener@nvidia.com>

fixed linting errors for blank lines

199d27f

Signed-off-by: Alp Dener <adener@nvidia.com>

denera force-pushed the jax-gemm-workspace-xla-buffer branch from 2c1703a to 199d27f Compare January 23, 2024 21:47

denera added the enhancement New feature or request label Jan 23, 2024

denera merged commit 4077ccc into NVIDIA:main Jan 29, 2024
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[JAX] Custom Op Workspace Tensors from XLA Buffers #532

[JAX] Custom Op Workspace Tensors from XLA Buffers #532

denera commented Nov 22, 2023

nouiz commented Nov 22, 2023

mingxu1067 Nov 23, 2023

denera Nov 23, 2023 •

edited

Loading

mingxu1067 Nov 23, 2023

nouiz Nov 23, 2023

zlsh80826 commented Nov 23, 2023 •

edited

Loading

denera commented Dec 15, 2023

denera commented Dec 19, 2023

denera commented Dec 20, 2023 •

edited

Loading

denera commented Dec 20, 2023

nouiz left a comment

zlsh80826 Dec 24, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

zlsh80826 Dec 26, 2023

denera commented Jan 17, 2024

denera commented Jan 22, 2024

denera commented Jan 22, 2024

	ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(wkspace_aval.dtype))
	ir.RankedTensorType.get(barrier_aval.shape, jax_dtype_to_ir_dtype(barrier_aval.dtype))

		wkspace_size = reduce(operator.mul, wkspace_aval.shape)
		barrier_size = reduce(operator.mul, barrier_aval.shape)

	wkspace_size = reduce(operator.mul, wkspace_aval.shape)
	wkspace_size = wkspace_aval.size

[JAX] Custom Op Workspace Tensors from XLA Buffers #532

[JAX] Custom Op Workspace Tensors from XLA Buffers #532

Conversation

denera commented Nov 22, 2023

nouiz commented Nov 22, 2023

Choose a reason for hiding this comment

denera Nov 23, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zlsh80826 commented Nov 23, 2023 • edited Loading

denera commented Dec 15, 2023

denera commented Dec 19, 2023

denera commented Dec 20, 2023 • edited Loading

denera commented Dec 20, 2023

nouiz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

denera commented Jan 17, 2024

denera commented Jan 22, 2024

denera commented Jan 22, 2024

denera Nov 23, 2023 •

edited

Loading

zlsh80826 commented Nov 23, 2023 •

edited

Loading

denera commented Dec 20, 2023 •

edited

Loading