Conversation
|
PyTorch changes are fine - JAX changes, I don't really see it solving the stream problem and, since the allocation is still taking place at the Python level, doesn't seem to move the needle much, so going to request that is rolled back. An aside: you've got two existing PRs with pending edits, I've got to review changes for other contributor, and everyone has to do rebases every time main advances; do me a favor and try to keep the count of active PRs small. |
| jit.set_max_smem(0, forward_config_ref.smem); | ||
| jit.set_max_smem(4, forward_config_ref.smem); | ||
| jit.set_max_smem(FORWARD, forward_config_ref.smem); | ||
| jit.set_max_smem(DOUBLE_BACKWARD_A, forward_config_ref.smem); |
There was a problem hiding this comment.
Cool, thanks for making this an enum
|
|
||
| ### Unreleased | ||
| **Changed**: | ||
| - The deterministic convolution workspace is now allocated by the framework (PyTorch / JAX) each call |
There was a problem hiding this comment.
Mm - it was always allocated by the framework (albeit at the Python level) - more accurate to state that it's allocated inside the extension at the C++ level.
| ffi::AnyBuffer workspace, | ||
| ffi::AnyBuffer workspace_in, | ||
| ffi::AnyBuffer transpose_perm, | ||
| ffi::Result<ffi::AnyBuffer> workspace, |
There was a problem hiding this comment.
Why this change to the JAX function signatures? I don't see a reason why the workspace needs to be returned.
| return torch::zeros(sizes, ref.options()); | ||
| } | ||
|
|
||
| Tensor tensor_empty_bytes(const Tensor &ref, int64_t nbytes) { |
There was a problem hiding this comment.
Seems like this function is unused?
|
|
||
| @functools.cache | ||
| def conv_workspace_shape(kernel: str) -> jax.ShapeDtypeStruct: | ||
| size = json.loads(kernel)["kernel_prop"]["workspace_size"] |
There was a problem hiding this comment.
So, not too happy with this set of changes on the JAX side; json.loads is nontrivial. Yes, I know it will get optimized away with jax.jit, but no reason to excessively saddle the non-JIT version.
This is effectively allocating a new workspace for every call and not moving the allocation into the C++ layer (which was the ostensible target of your PyTorch changes, so I don't really see this change moving the needle. Please roll back the JAX workspace changes.
| @@ -0,0 +1,123 @@ | |||
| import pytest | |||
There was a problem hiding this comment.
Can you disable this test when pytest --jax is enabled?
| def step(): | ||
| return _fwd_bwd(compiled, X, Y, W, rows, cols, sender_perm, G) | ||
|
|
||
| for _ in range(4): |
There was a problem hiding this comment.
Why four times? And perhaps this test should live in compile_tests, not here?
| _assert_close(outputs, reference, deterministic) | ||
|
|
||
|
|
||
| def test_concurrent_streams_share_no_state(conv_and_inputs): |
There was a problem hiding this comment.
Would move to stream_tests.py
There was a problem hiding this comment.
The tests in this file could probably live in the existing test files - only one deals with cuda graphs, and that could probably go in the compile_tests.py file.
2 participants
This PR moves the responsibility of allocating the workspace to the frameworks (jax, torch) at each call site