From 7f67a8e9db1fad3d69e6fa5ff565333e729034b1 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Tue, 5 May 2026 11:37:01 +0200 Subject: [PATCH 1/9] upgrade pytorch and flashattn wip --- pyproject.toml | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 00103cb8c..6dfd144e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,11 +76,11 @@ dev = [ [project.optional-dependencies] cpu = [ - 'torch==2.6.0', + 'torch==2.9.1', ] gpu = [ - "torch==2.6.0; sys_platform == 'linux'", + "torch==2.9.1; sys_platform == 'linux'", # flash-attn also has a torch dependency. "flash-attn", ] @@ -166,7 +166,7 @@ link-mode = "symlink" # by future releases of dependencies or sub-dependencies. # See https://docs.astral.sh/uv/reference/settings/#exclude-newer # TODO: pytorch does not publish valid release timestamps, so sadly it does not work. -exclude-newer = "2026-02-27T00:00:00Z" +exclude-newer = "2026-04-27T00:00:00Z" # The minimum version of uv required. # It is tightly controlled because the format of uv.lock has changed @@ -191,12 +191,6 @@ conflicts = [ ] -[[tool.uv.index]] -name = "pytorch-cu126" -url = "https://download.pytorch.org/whl/cu126" -explicit = true - - [tool.pyrefly] project-includes = ["src/"] project-excludes = [ @@ -231,25 +225,23 @@ weathergen-readers-extra = { workspace = true } flash-attn = [ -# The build of Cathal O'Brien is not compatible with the libc build on santis. -# Hardcode the reference to the swiss cluster for the time being. -# TODO: open issue -# { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.7.4+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, -# This version was rebuilt locally on santis and uploaded. - { url = "https://object-store.os-api.cci1.ecmwf.int/weathergenerator-dev/wheels/flash_attn-2.7.3-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, # There is no official support for flash-attn on macos ] torch = [ # Explicit pin for GPU - { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-linux_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + # { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-linux_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download.pytorch.org/whl/cu129/torch/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + # { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + { url = "https://download.pytorch.org/whl/cu129/torch/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, # Explicit pin for CPU. - { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-linux_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0-cp312-none-macosx_11_0_arm64.whl",marker = "sys_platform == 'darwin'", extra="cpu" }, + # { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, + { url = "https://download.pytorch.org/whl/cpu/torch/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, + # { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-linux_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, + { url = "https://download.pytorch.org/whl/cpu/torch/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, # No acceleration for now. # TODO: explore pytorch + metal backend { index = "pytorch-cpu", marker = "sys_platform == 'darwin'", extra="gpu"}, From dc6d05b75f56e524102dea552a34f006666df00c Mon Sep 17 00:00:00 2001 From: florianscheidl Date: Tue, 5 May 2026 11:58:30 +0200 Subject: [PATCH 2/9] update download paths --- pyproject.toml | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6dfd144e1..000fcbb48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -230,18 +230,13 @@ flash-attn = [ # There is no official support for flash-attn on macos ] - torch = [ # Explicit pin for GPU - # { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-linux_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download.pytorch.org/whl/cu129/torch/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - # { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, - { url = "https://download.pytorch.org/whl/cu129/torch/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, # Explicit pin for CPU. - # { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, - { url = "https://download.pytorch.org/whl/cpu/torch/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, - # { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-linux_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, - { url = "https://download.pytorch.org/whl/cpu/torch/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, # No acceleration for now. # TODO: explore pytorch + metal backend { index = "pytorch-cpu", marker = "sys_platform == 'darwin'", extra="gpu"}, From 5a6ea7b50cdba043004be269735c645dd7301a78 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Tue, 5 May 2026 13:36:05 +0200 Subject: [PATCH 3/9] Add Macos support --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 000fcbb48..f66a50421 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,6 +237,8 @@ torch = [ # Explicit pin for CPU. { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl",marker = "sys_platform == 'darwin'", extra="cpu" }, + # No acceleration for now. # TODO: explore pytorch + metal backend { index = "pytorch-cpu", marker = "sys_platform == 'darwin'", extra="gpu"}, From be79b44b51bafa3aebf2f42e327a231fd5ae51f1 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Tue, 5 May 2026 13:55:27 +0200 Subject: [PATCH 4/9] remove macos specific support --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f66a50421..5617de837 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,7 +237,6 @@ torch = [ # Explicit pin for CPU. { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, - { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl",marker = "sys_platform == 'darwin'", extra="cpu" }, # No acceleration for now. # TODO: explore pytorch + metal backend From 2334509903b1b2159c3189c19b7b1f27b2fa05e5 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Tue, 5 May 2026 20:54:06 +0200 Subject: [PATCH 5/9] downgrade to cu126 again --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5617de837..8b87608dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -232,8 +232,8 @@ flash-attn = [ torch = [ # Explicit pin for GPU - { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, # Explicit pin for CPU. { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, From 28c3ee90c8d83ae8c99526e1fe7d04e8033662eb Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Wed, 6 May 2026 15:30:50 +0200 Subject: [PATCH 6/9] Specify build dependency to avoid dependency issues on specific wheels --- pyproject.toml | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8b87608dd..3750c4c61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -223,13 +223,6 @@ weathergen-evaluate = { workspace = true } weathergen-metrics = { workspace = true } weathergen-readers-extra = { workspace = true } - -flash-attn = [ - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, - # There is no official support for flash-attn on macos -] - torch = [ # Explicit pin for GPU { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, @@ -245,6 +238,9 @@ torch = [ # branch = "feature/any-zarr-version-in-pyproject/" anemoi-datasets = { git = "https://github.com/tjhunter/anemoi-datasets", tag = "0.0.1"} +[tool.uv.extra-build-dependencies] +flash-attn = ["torch"] + [tool.pytest.ini_options] log_cli = true log_cli_level = "INFO" From 12838d4ca45a95a5993636ad6d9e437626766309 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Wed, 6 May 2026 15:58:12 +0200 Subject: [PATCH 7/9] Revert "downgrade to cu126 again" This reverts commit 2334509903b1b2159c3189c19b7b1f27b2fa05e5. --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8b87608dd..5617de837 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -232,8 +232,8 @@ flash-attn = [ torch = [ # Explicit pin for GPU - { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, # Explicit pin for CPU. { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, From e62d52b3fdeba68f7fdda8363be067074c60b076 Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Wed, 6 May 2026 16:02:22 +0200 Subject: [PATCH 8/9] Reset by hand --- pyproject.toml | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6714bbae5..8b87608dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -223,10 +223,17 @@ weathergen-evaluate = { workspace = true } weathergen-metrics = { workspace = true } weathergen-readers-extra = { workspace = true } + +flash-attn = [ + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, + # There is no official support for flash-attn on macos +] + torch = [ # Explicit pin for GPU - { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download-r2.pytorch.org/whl/cu129/torch-2.9.1+cu129-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, # Explicit pin for CPU. { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, @@ -238,9 +245,6 @@ torch = [ # branch = "feature/any-zarr-version-in-pyproject/" anemoi-datasets = { git = "https://github.com/tjhunter/anemoi-datasets", tag = "0.0.1"} -[tool.uv.extra-build-dependencies] -flash-attn = ["torch"] - [tool.pytest.ini_options] log_cli = true log_cli_level = "INFO" From e249a03e9e495fd86adee4f9507fe810a11a33df Mon Sep 17 00:00:00 2001 From: Florian Scheidl Date: Fri, 22 May 2026 15:52:17 +0200 Subject: [PATCH 9/9] Trying pinned ABIfalse for x_86 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8b87608dd..9fe5dc158 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,7 +226,7 @@ weathergen-readers-extra = { workspace = true } flash-attn = [ { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, + { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.8.3+cu12torch2.9cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, # There is no official support for flash-attn on macos ]