[Pallas] Support programmatically extracting the payload #6696

alanwaketan · 2024-03-08T02:25:35Z

Summary:
This pull request introduces _extract_backend_config to extrac the Mosaic payload from the custom call programmatically such that we don't need to copy & paste.

In order to run the test, JAX dependencies are added to the CI. The JAX version is using the nightly on the same day as our libtpu for the best compatibility. However, we need to figure out a way to update that automatically when we are updating our open-xla pin.

Test Plan:
python test/test_operations.py -v -k test_tpu_custom_call_pallas_extract_add_payload

alanwaketan · 2024-03-08T02:25:46Z

I will add the CI change later.

will-cromar · 2024-03-08T02:44:02Z

.github/workflows/tpu_ci.yml

        run: |
          pip install fsspec
          pip install rich
+          pip install -U --pre jaxlib==0.4.25.dev20240213 -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html


Can you add these dependencies to our setup.py? Maybe a pallas extras_require since AFAIK jax will still be optional. We can update the JAX version together with libtpu-nightly when we do pin updates

Sure, I wasn't quite sure to make it default for development, but I guess why not.

@will-cromar How are you going to add nightly dependencies to setup.py? I tried:

ptxla@t1v-n-94bbce5b-w-0:/workspaces/work/pytorch/xla$ git diff setup.py diff --git a/setup.py b/setup.py index e4fdaebee..66a1eb321 100644 --- a/setup.py +++ b/setup.py @@ -288,6 +288,8 @@ setup( # importlib.metadata backport required for PJRT plugin discovery prior # to Python 3.10 'importlib_metadata>=4.6;python_version<"3.10"', + 'jaxlib==0.4.25.dev20240213 @ https://storage.googleapis.com/jax-releases/jaxl ib_nightly_releases.html', + 'jax==0.4.25.dev20240213 @ https://storage.googleapis.com/jax-releases/jax_nig htly_releases.html', ], package_data={ 'torch_xla': ['lib/*.so*',],

but that doesn't work...

error in torch_xla setup command: 'install_requires' must be a string or list of strings containing valid project/version requirement specifiers; Parse error at "'@ https:'": Expected string_end

Add it to extras_require so it's not required by default:

extras_require={ # On Cloud TPU VM install with: # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html 'tpu': [f'libtpu-nightly=={_libtpu_version}'], # On nightly, install libtpu with `pip install torch_xla[tpuvm]` # Remove from release branches since this is not allowed by PyPI. 'tpuvm': [f'libtpu-nightly @ {_libtpu_storage_path}'], 'pallas': ['jaxlib==0.4.25.dev20240213', 'jax==0.4.25.dev20240213'] },

To uploadable to PyPI, you can't bake in a link to another index. The user will have to do pip install torch_xla[pallas] -f path/to/index like they do for TPU.

.... unless you can just rely on a stable release of jax that's already on PyPI. If JAX is not loading libtpu (I suspect it does not, otherwise we would crash), the JAX build date does not have to match libtpu-nightly. Maybe it's better to investigate that option before we make the feature stable so we can get this merged soon.

No, JAX will need Mosaic which is in the libtpu to generate the payload. So we need the nightly. Thanks, Will. Will do it in this way.

will-cromar · 2024-03-08T02:47:36Z

test/test_operations.py

    self.assertTrue(torch.allclose(output.cpu(), expected_output.cpu()))

+  @unittest.skipIf(xr.device_type() != 'TPU', "This test only works on TPU.")
+  def test_tpu_custom_call_pallas_extract_add_payload(self):


Could the custom kernel tests should be in a separate test file? IMO test_operations shouldn't depend on anything other than our basic requirements

torchvision is imported but we don't actually use it

yea I agree that pallas related stuff is better to be putted in a new test file.

I'm just too lazzzy...

lol the problem of the pallas is it introcued the JAX dependency. We might have to disable the test on TPUCI sometimes, it is better not to disable the whole test_operations.py

Moved to a new file.

JackCaoG · 2024-03-08T20:16:45Z

Do we need this to be merged into 2.3 release? This intorduce a JAX dependency which might cause issues. I would prefer to leave this in nightly until things stabilized a bit.

alanwaketan · 2024-03-08T20:29:51Z

Do we need this to be merged into 2.3 release? This intorduce a JAX dependency which might cause issues. I would prefer to leave this in nightly until things stabilized a bit.

Yea, it needs to be in 2.3. I guess we can keep jax as an extra dependencies for those who needs pallas.

alanwaketan · 2024-03-08T21:55:44Z

@will-cromar I fixed the jax thing. Hopefully, this time it works.

alanwaketan · 2024-03-08T22:36:52Z

Thanks Jack for approving!

will-cromar · 2024-03-08T22:43:13Z

test/test_pallas.py

+    # def add_vectors_kernel(x_ref, y_ref, o_ref):
+    #   x, y = x_ref[...], y_ref[...]
+    #   o_ref[...] = x + y
+    payload = "{\"custom_call_config\": {\"body\": \"TUzvUgFNTElSMTguMC4wZ2l0AAErCwEDBQcJAQMLAwUDDQcFDxEJBRMVA2lNDQFLBw8LEw8PDwsPMwsLCwtlCwsLCwsPCw8PEwsTDwsTDwsPDxMLDwUDYQENGwcTDxsPAsICHx0rLQUXAwMnKRURNx1HSRELAQUZHTM1AwsVFxkbHw0hDSMlBRsBAQUdDQlhZmZpbmVfbWFwPChkMCkgLT4gKGQwKT4ABR8FIQUjBSUFJxEDAQUpFS8JHQ8xFwUTAQUrFwUdAR05OwUtFwUlAR0/QQUvFUMJHQ9FFwUVAQUxFREJI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4AF0sDIQcdAycDIQcBAgIFBwEBAQEBAgQEpwUBEAEHAwEFAxEBEwcDFScHAQEBAQEBBwMDBwMDCwYDAwUFAQcHAwMHAwMLBgMDBQUDCwkGPQMFBQkNBwMLBwMDCwYLAwUFBRENBAsHDwURBQABBgMBBQEAdgcz2wsTGdkNCxMjIR0pJ0MNCwsTDw8PDQkLEWJ1aWx0aW4AZnVuYwB0cHUAYXJpdGgAdmVjdG9yAG1vZHVsZQByZXR1cm4AY29uc3RhbnQAYWRkaQBsb2FkAHN0b3JlAC9ob21lL2p3dGFuL3BhbGxhcy9wYWxsYXNfYWRkLnB5AGFkZF92ZWN0b3JzX2tlcm5lbABkaW1lbnNpb25fc2VtYW50aWNzAGZ1bmN0aW9uX3R5cGUAc2NhbGFyX3ByZWZldGNoAHNjcmF0Y2hfb3BlcmFuZHMAc3ltX25hbWUAbWFpbgB2YWx1ZQAvZ2V0W3RyZWU9UHlUcmVlRGVmKChDdXN0b21Ob2RlKE5ESW5kZXhlclsoUHlUcmVlRGVmKChDdXN0b21Ob2RlKFNsaWNlWygwLCA4KV0sIFtdKSwpKSwgKDgsKSwgKCkpXSwgW10pLCkpXQBhZGRfdmVjdG9ycwA8bW9kdWxlPgAvYWRkAC9zd2FwW3RyZWU9UHlUcmVlRGVmKChDdXN0b21Ob2RlKE5ESW5kZXhlclsoUHlUcmVlRGVmKChDdXN0b21Ob2RlKFNsaWNlWygwLCA4KV0sIFtdKSwpKSwgKDgsKSwgKCkpXSwgW10pLCkpXQA=\", \"needs_layout_passes\": true}}"


I was going to ask if this base64 (I assume) blob means anything human-readable you could encode at test time. Then I decoded it:

$ echo TUzvUgFNTElSMTguMC4wZ2l0AAErCwEDBQcJAQMLAwUDDQcFDxEJBRMVA2lNDQFLBw8LEw8PDwsPMwsLCwtlCwsLCwsPCw8PEwsTDwsTDwsPDxMLDwUDYQENGwcTDxsPAsICHx0rLQUXAwMnKRURNx1HSRELAQUZHTM1AwsVFxkbHw0hDSMlBRsBAQUdDQlhZmZpbmVfbWFwPChkMCkgLT4gKGQwKT4ABR8FIQUjBSUFJxEDAQUpFS8JHQ8xFwUTAQUrFwUdAR05OwUtFwUlAR0/QQUvFUMJHQ9FFwUVAQUxFREJI3RwdS5tZW1vcnlfc3BhY2U8dm1lbT4AF0sDIQcdAycDIQcBAgIFBwEBAQEBAgQEpwUBEAEHAwEFAxEBEwcDFScHAQEBAQEBBwMDBwMDCwYDAwUFAQcHAwMHAwMLBgMDBQUDCwkGPQMFBQkNBwMLBwMDCwYLAwUFBRENBAsHDwURBQABBgMBBQEAdgcz2wsTGdkNCxMjIR0pJ0MNCwsTDw8PDQkLEWJ1aWx0aW4AZnVuYwB0cHUAYXJpdGgAdmVjdG9yAG1vZHVsZQByZXR1cm4AY29uc3RhbnQAYWRkaQBsb2FkAHN0b3JlAC9ob21lL2p3dGFuL3BhbGxhcy9wYWxsYXNfYWRkLnB5AGFkZF92ZWN0b3JzX2tlcm5lbABkaW1lbnNpb25fc2VtYW50aWNzAGZ1bmN0aW9uX3R5cGUAc2NhbGFyX3ByZWZldGNoAHNjcmF0Y2hfb3BlcmFuZHMAc3ltX25hbWUAbWFpbgB2YWx1ZQAvZ2V0W3RyZWU9UHlUcmVlRGVmKChDdXN0b21Ob2RlKE5ESW5kZXhlclsoUHlUcmVlRGVmKChDdXN0b21Ob2RlKFNsaWNlWygwLCA4KV0sIFtdKSwpKSwgKDgsKSwgKCkpXSwgW10pLCkpXQBhZGRfdmVjdG9ycwA8bW9kdWxlPgAvYWRkAC9zd2FwW3RyZWU9UHlUcmVlRGVmKChDdXN0b21Ob2RlKE5ESW5kZXhlclsoUHlUcmVlRGVmKChDdXN0b21Ob2RlKFNsaWNlWygwLCA4KV0sIFtdKSwpKSwgKDgsKSwgKCkpXSwgW10pLCkpXQA= | base64 --decode ML�RMLIR18.0.0git+ K iM 3 e +-')7GI 35 ffine_map<(d0) -> (d0)>!#%')/ 1+9;-%?A/C E1 #tpu.memory_space<vmem>K!'!�' = v3� � #!)'C builtinfunctpuarithvectormodulereturnconstantaddiloadstore/home/jwtan/pallas/pallas_add.pyadd_vectors_kerneldimension_semanticsfunction_typescalar_prefetchscratch_operandssym_namemainvalue/get[tree=PyTreeDef((CustomNode(NDIndexer[(PyTreeDef((CustomNode(Slice[(0, 8)], []),)), (8,), ())], []),))]add_vectors<module>/add/swap[tree=PyTreeDef((CustomNode(NDIndexer[(PyTreeDef((CustomNode(Slice[(0, 8)], []),)), (8,), ())], []),))]

You can secretly decode it in Google 3 haha...

alanwaketan · 2024-03-08T23:08:00Z

The dynamo issue in CI is unrelated...

Summary: This pull request introduces _extract_backend_config to extrac the Mosaic payload from the custom call programmatically such that we don't need to copy & paste. In order to run the test, JAX dependencies are added to the CI. Test Plan: python test/test_operations.py -v -k test_tpu_custom_call_pallas_extract_add_payload

alanwaketan requested review from JackCaoG and will-cromar March 8, 2024 02:25

alanwaketan self-assigned this Mar 8, 2024

alanwaketan requested a review from qihqi March 8, 2024 02:28

will-cromar reviewed Mar 8, 2024

View reviewed changes

alanwaketan force-pushed the alanwaketan/pallas_jax branch from 7ec4903 to 769fa66 Compare March 8, 2024 20:30

JackCaoG approved these changes Mar 8, 2024

View reviewed changes

will-cromar approved these changes Mar 8, 2024

View reviewed changes

alanwaketan added 8 commits March 8, 2024 23:23

Fix linters

acfc791

install jax in tpu ci

cd5fe09

Skip the test for non tpu ci

2efae7c

Move tests

a7e7c2a

Fix linters

f63e3c9

Improve jax dependency

3f05b90

Fix comment

18841a4

alanwaketan force-pushed the alanwaketan/pallas_jax branch from 0540fb8 to 18841a4 Compare March 8, 2024 23:23

alanwaketan added 3 commits March 9, 2024 02:11

Fix a type

07ad739

Add tpu ci

95075ed

Try fixing tpu ci

ae6f8dc

alanwaketan merged commit 3706791 into master Mar 9, 2024

alanwaketan deleted the alanwaketan/pallas_jax branch March 9, 2024 18:48

lsy323 mentioned this pull request Mar 11, 2024

Disable pallas test in tpu ci #6715

Merged

vanbasten23 restored the alanwaketan/pallas_jax branch March 11, 2024 23:08

alanwaketan mentioned this pull request Mar 12, 2024

[Pallas] Introduce make_kernel_from_pallas #6713

Merged

[Pallas] Support programmatically extracting the payload #6696

[Pallas] Support programmatically extracting the payload #6696

Uh oh!

Conversation

alanwaketan commented Mar 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alanwaketan commented Mar 8, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

will-cromar Mar 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JackCaoG commented Mar 8, 2024

Uh oh!

alanwaketan commented Mar 8, 2024

Uh oh!

alanwaketan commented Mar 8, 2024

Uh oh!

alanwaketan commented Mar 8, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alanwaketan commented Mar 8, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

alanwaketan commented Mar 8, 2024 •

edited

Loading

will-cromar Mar 8, 2024 •

edited

Loading