From 489bacf518c50d4a5fe86b5002beb2411c85e618 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Tue, 23 Jan 2024 09:25:03 -0800 Subject: [PATCH 1/2] Pins jax==0.4.23 (xla matched), flax==0.7.5, te==1.2.1 --- .github/container/manifest.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml index 9d4240bc3..7303f219e 100644 --- a/.github/container/manifest.yaml +++ b/.github/container/manifest.yaml @@ -2,25 +2,25 @@ jax: url: https://github.com/google/jax.git tracking_ref: main - latest_verified_commit: 2356e57af47baa37286177676a5a4bcd0c7b8c40 + latest_verified_commit: jaxlib-v0.4.23 mode: git-clone xla: url: https://github.com/openxla/xla.git tracking_ref: main - latest_verified_commit: 17ec7b4031608d3ddaca83e0661ed3aa0393e9d9 + latest_verified_commit: a487d8ba5da8d0dec056972270e50de6748a1035 mode: git-clone flax: url: https://github.com/google/flax.git mirror_url: https://github.com/nvjax-svc-0/flax.git tracking_ref: main - latest_verified_commit: 7cbc3c13ecd8a1713732113dec7939d0b75342c6 + latest_verified_commit: v0.7.5 mode: git-clone patches: pull/3340/head: file://patches/flax/PR-3340.patch # Add Sharding Annotations to Flax Modules transformer-engine: url: https://github.com/NVIDIA/TransformerEngine.git tracking_ref: main - latest_verified_commit: bb759adc2a53c57cc387f531c580a16149e3d4ac + latest_verified_commit: v1.2.1 mode: git-clone t5x: url: https://github.com/google-research/t5x.git From d9b60e0f1e4ff3f4eeb5713cbacf90f80a37d751 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Tue, 23 Jan 2024 17:51:31 -0800 Subject: [PATCH 2/2] revert flax to previous manifest commit since rosetta t5x doesn't build due to patches being too new --- .github/container/manifest.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml index 7303f219e..7e00bdc87 100644 --- a/.github/container/manifest.yaml +++ b/.github/container/manifest.yaml @@ -13,7 +13,7 @@ flax: url: https://github.com/google/flax.git mirror_url: https://github.com/nvjax-svc-0/flax.git tracking_ref: main - latest_verified_commit: v0.7.5 + latest_verified_commit: 7cbc3c13ecd8a1713732113dec7939d0b75342c6 mode: git-clone patches: pull/3340/head: file://patches/flax/PR-3340.patch # Add Sharding Annotations to Flax Modules