Update jax array api #7

EiffL · 2024-03-01T19:20:36Z

This is a PR led by @ASKabalan to update the project to the JAX v0.4+ API

.gitmodules

README.md

benchmarks/autotune/autotune_4_256.txt

benchmarks/jaxpm.py

EiffL · 2024-03-02T16:47:33Z

I reformatted the code on main and in this branch using yapf to make it easier to review the bits that have changed

EiffL · 2024-03-02T17:01:53Z

I'm still working through it, but some top level comments:

We don't want to use a fork of cuDecomp, we should use the upstream version. If we need to add something to it, we can always submit a PR to the upstream project.
The repo for the library itself should be kept clean and tidy, so we don't want to commit here the results of experimentation runs. We can have a separate repo where we save the results of experiments. It is fine however to have the instrumentation scripts here.
As a library, jaxDecomp should not infringe on API boundaries. By that I mean, the scope of jaxDecomp should be strictly limited to providing fft ops as close as possible to drop in replacement for the official ones. It is out of scope to define utility functions like map_global_array_from_slices. Of course we can define them and use them internally if we need them, but the end users should not have to use them, nor be told about them, because it would impact the way they write their scripts beyond the strict minimum of having access to distributed ops.
The point here is to maintain clean interfaces between what different libraries provides, and have as minimal of a footprint as possible. The bigger the API we provide, the more ways it can break down when JAX changes and the costlier it is to maintain.

README.md

aboucaud · 2024-03-19T17:04:30Z

benchmarks/scripts/README.md

+# To use on Jean Zay with the TKC project
+


README file should be moved to the parent dir

I moved benchmarks to other repo.
Tell me if you still want the benchmark code to be here
I can add an 'example' folder instead

aboucaud · 2024-03-19T17:09:20Z

benchmarks/scripts/plotter.py

+        sorted_df = pd.DataFrame(sorted_dfs)
+        label = f"{method}-{backend}-{group['nodes'].values[0]}nodes" if nodes_in_label else f"{method}-{backend}"
+
+        ax.plot(sorted_df['x'].values, sorted_df['time'], marker='o', linestyle='-', label=label)
+
+    # add title nb of gpus
+    ax.set_title(f"Number of GPUs: {gpu_size}")
+    # Set x-axis ticks exactly at the provided data size values
+    ax.set_yscale('log')
+    data_sizes = list(dict.fromkeys(data_sizes))
+    data_sizes.sort()
+    ax.set_xticks(data_sizes)
+    ax.set_yticks([1e-3, 1e-2, 1e-1, 1e0])
+
+    # Set labels and title
+    ax.set_xlabel('Data Size')
+    ax.set_ylabel('Time')


FYI if you are not aware yet pandas has its own plotting API on top of matplotlib so you can do the same without all the burden once your dataframe is ready
https://pandas.pydata.org/pandas-docs/stable/user_guide/visualization.html

slurms/HOWTORUN.md

tests/test_fft.py

aboucaud · 2024-03-19T17:32:17Z

tests/test_fft.py

+# find a way to finalize pytest
+def test_end():
+  # Make sure that it is cleaned up
+  # This has to be this way because pytest runs the "global code" before running the tests
+  # There are other solutions https://stackoverflow.com/questions/41871278/pytest-run-a-function-at-the-end-of-the-tests


What about defining some fixtures with scope=session in a conftest.py ?

This could also be useful for the initialize at the beginning..

session is multiple pytest files, no?
I am not sure how pytest works.
But clean up needs to be done after every test

aboucaud · 2024-03-19T18:19:20Z

jaxdecomp/fft.py

@@ -1,9 +1,11 @@
 import jax.numpy as jnp
 from jax._src.numpy.fft import _fft_norm


_fft_norm is redefined below (code copied from jax source)

right, can we remove the code below and just use this function?

done in 3786c77

aboucaud · 2024-03-19T19:26:39Z

README.md

-### Step II: Building jaxDecomp
-
-This step is easier :-) From this directory, just run the following
+From this directory, just run the followg


https://www.todepond.com/wikiblogarden/better-computing/just/

EiffL

I added a few questions, but it looks pretty good!

If I understand correctly, right now there is a limitation that things will only work for a 3D cube, is this the same issue as the one I was too lazy to fix here:

Only transpose operations that preserve the size of local slices are supported
#1

.clang-format

CMakeLists.txt

benchmarks/mpi4jafft.py

jaxdecomp/_src/fft.py

tests/test_fft.py

EiffL · 2024-03-21T04:32:47Z

tests/test_fft.py


  # Check the forward FFT
-  assert_allclose(global_karray_slice.real, karray.real, atol=1e-10)
-  assert_allclose(global_karray_slice.imag, karray.imag, atol=1e-10)
+  # Non cube shape .. need to think about the slice coming from JAX


hummmm what does this mean? it doesn't work if the array has abritrary shape?

~~It does. (this is just a testing issue not a production code issue)~~
~~Issue is that pdims is [Z][X]~~
~~the code before used pdims as pdims[1] pdims[0] .. which obviously works because 2x4 and 4x2 will give the right FFT.~~
~~How ever when we construct the global array dims I use pdims for the partition specs and in this case they have to be accurate.~~

~~IE : (global is 512 512 512) a 2x4 pdims will give in the old code (256 128 512)~~
~~for cuDecomp pdims[0] is the Z axis, pdims[1] is the Y axis so If I do this :~~

The comment is not about pfft, it is about JAX fft and the tranpose needed
pfft works on abritrary shapes
Just testing against the array is not working the shapes of global_karray_slice are weird

EiffL · 2024-03-21T04:33:30Z

tests/test_fft.py

-                    global_shape[1] // pdims[1]].imag,
-      karray.imag,
-      atol=1e-10)
+  #assert_allclose(


why are we removing those?

shape of JAX arrays are wrong in the test.
This is not pfft error but just a test error
The FFT gives the right results

tests/test_fft.py

ASKabalan · 2024-04-27T22:26:10Z

File "/mnt/home/flanusse/repo/jaxDecomp/scripts/demo.py", line 58, in
recarray = slice_unpad(exchanged_reduced, padding_width, pdims)
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: custom_partitioner: TypeError: pad operand and padding_value must be same dtype: got complex64 and float32.

The issue is the fake padding_value that I give, it should be the same type as the array
Issue is here

Maybe doing arr[first:-last] is just better

I will fix once you finish reviewing

EiffL

Seems to be working :-)

EiffL · 2024-04-27T22:44:04Z

jaxdecomp/fft.py

-
+# Has to be jitted here because _fft_norm will act on non fully addressable global array
+# Which means this should be jit wrapped
+@partial(jit, static_argnums=(0,1,3))


@ASKabalan , same question here

ASKabalan · 2024-04-27T22:57:32Z

Seems to be working :-)

If you are talking about from mpi4py import MPI before jax.distributed() for me it doesn't
I use srun and not mpirun to run code, what do you use?

EiffL · 2024-04-27T23:00:54Z

I meant everything seems to run (minus details noted above).

I've removed all referenced to mpi4py because it's a pain to install, and is not necessary now that jax understands distribution itself.

I'm going to go ahead and merge, we can treat the rest of the comments as small pull requests. This branch is way better than the main branch already.

ASKabalan · 2024-04-27T23:03:06Z

Before merging, please see my comment about the distribted global_state
here

EiffL · 2024-04-27T23:04:10Z

hummm what's the comment?

ASKabalan · 2024-04-27T23:05:08Z

You must use jax.process_index() instead global_state.rank

EiffL · 2024-04-27T23:05:14Z

Maybe you haven't pushed your review? I can't see a comment at that line

README.md

benchmarks/autotune/autotune_4_256.txt

benchmarks/jaxpm.py

.clang-format

CMakeLists.txt

README.md

ASKabalan · 2024-04-27T22:16:17Z

README.md

 ## Usage

 The API is still under development, so it doesn't look very streamlined, but you
 can already do the following:
 ```python
-from mpi4py import MPI


WORLD communicator must be set in order for jax.distributed to work, without it it won't work

Ex :

import jax #from mpi4py import MPI <-- without this, it doesn't work jax.distributed.initialize() print(jax.devices())

@EiffL if you want to remove mpi4py as a requirement, we will have to init the WORLD communicator ourselfs .. I would have to recreate the init function that just instantiate the singleton like this

void init(){jd::GridDescriptorManager::getInstance();};

Then we do this

import jax import jaxdecomp jaxdecomp.init() jax.distributed.initialize() print(jax.devices())

Ok for you?

This may work for you, but not on JeanZay where we have to use srun

no, we shouldn't need to init things ourselves for jax.distributed to work.

ASKabalan · 2024-04-27T22:28:12Z

jaxdecomp/_src/padding.py

+        interiour_padding = 0
+
+        # unlike jnp.pad lax.pad can unpad if given negative values
+        return lax.pad(arr, padding_value=0.0, padding_config=((-first_x, -last_x , interiour_padding), (-first_y, -last_y , interiour_padding), (-first_z, -last_z , interiour_padding)))


Here, The padding_value has to be the same type as the array

Or maybe just using [first:-last] slicing
What do you think @EiffL

EiffL · 2024-04-27T23:08:53Z

Ah ok, I see your comments now.

WORLD communicator must be set in order for jax.distributed to work, without it it won't work

I have no problem running this with mpirun, jax.distributed works as expected, with no mpi4py. Where do you see an issue?

ASKabalan · 2024-04-27T23:09:50Z

when I use srun

I get this

CUDA backend failed to initialize: jaxlib/cuda/versions_helpers.cc:90: operation gpuDeviceGetAttribute( &major, GPU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device) failed: CUDA_ERROR_NOT_INITIALIZED (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Unless a communicator is set.
I can quickly recreate the init function and test

EiffL · 2024-04-27T23:10:19Z

Note that I'm instantiating the MPI context here:

jaxDecomp/src/grid_descriptor_mgr.cc

Line 26 in aededef

CHECK_MPI_EXIT(MPI_Init(nullptr, nullptr));

ASKabalan · 2024-04-27T23:11:48Z

This is lazily executed at the first PFFT execution.
Too late for that.

I will try to recreate the init and tell you

My comment explains it

EiffL · 2024-04-27T23:11:51Z

Yup, can you show me exactly the slurm config you use when you have this issue?

EiffL · 2024-04-27T23:13:02Z

Sure, but jax.distributed does not depend on mpi4py, we should not need it.

(everything (all tests and all scripts) work fine on my machine without mpi4py)

ASKabalan · 2024-04-27T23:19:18Z

I think I force pushed on your commit ..

EiffL · 2024-04-27T23:28:46Z

no no, we don't need to init MPI comm, send me your slurm script, I'll modify it, it's just a matter of gpu affinity with tasks.

ASKabalan · 2024-04-27T23:29:09Z

no no, we don't need to init MPI comm, send me your slurm script, I'll modify it, it's just a matter of gpu affinity with tasks.

Slack?

EiffL commented Mar 2, 2024

View reviewed changes

ASKabalan force-pushed the update-jax-array-api branch from 548e910 to 2333fd7 Compare March 4, 2024 11:11

romerojosh mentioned this pull request Mar 4, 2024

Issue with CMake when cuDecomp is used a subdirectory NVIDIA/cuDecomp#25

Closed

ASKabalan force-pushed the update-jax-array-api branch 3 times, most recently from d1b9683 to 0920183 Compare March 12, 2024 19:01

ASKabalan force-pushed the update-jax-array-api branch 5 times, most recently from e33d034 to 8989a8a Compare March 17, 2024 00:50