Repository navigation
Fine-tuning (experimental): differentiable attention kernel, head-only and whole-model training - #2
Merged
Merged
Conversation
…coder finetune.py trains the head (5.6M of WeatherNext 2's 184M parameters) with the official training loss (fair CRPS, 2 members); the encoder, grid->mesh GNN and transformer stay frozen, so no attention backward pass is needed. New patch option `remat` (off by default) rematerialises each blocked GNN / encoder / decoder step in the backward pass; inference equivalence re-run unchanged (highest 1.98e-5, default 4.87e-3 max rel RMS). - ifs.training_example: one-step examples from IFS open data (+6 h analysis, 0-6 h IFS precipitation as targets) - scripts/verify_head_grads.py: Mini (1°) head gradients vs the official path, judged against a second official implementation as the rounding floor - scripts/finetune_head.py: 0.25° training (20.0 GiB temp, 3.84 s/step on an RTX 5090); lr 3e-6 + warmup lowered held-out loss 1.5 %, lr 3e-5 raised it - [train] extra (optax); evidence in docs/validation.md §8 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The fused Pallas attention gets a FlashAttention-2 backward (jax.custom_vjp): the forward also stores the log-sum-exp, and two tiled kernels compute dQ and dK/dV over the same tile lists (the k-hop mask is symmetric). All kernels use one dot precision; under `highest` the backward runs in 32-row programs (bwd_split=2) to fit sm_120's shared memory. One layer of the real mask: 42 ms forward+backward at TF32 vs 229 ms for XLA autodiff. - remat now also covers each transformer layer and each XLA attention chunk (jax.checkpoint there: hk.remat inside a raw lax.map leaks Haiku state) - sequential_samples: FGN's 2 members in an hk.scan, one at a time in the backward (45.8 -> 25.5 GiB at 0.25°; independent remat calls did not help) - finetune.split(params, "all"), finetune_head.py --modules all; a 0.25° full-model step takes 8.5 s and 28.6 GiB (pool 0.95); 20 steps lowered held-out one-step loss 1.8-2.4 % - whole-model gradients vs the official path on Modal H200 (cloud/modal_grads.py; official backward 102 GiB even at 1°): at TF32 at the official-vs-official floor; at highest with strict_fp32 at the floor - verify_head_grads.py: --modules all, --loss linear (smooth reference), XLA-autodiff path; attn_bench.py --grad - inference re-validated: highest 1.97e-5, default 4.74e-3 (6.54e-3 in one run, TF32 kernel-selection variance); evidence in docs/validation.md §9 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…head memory - scripts/finetune_head.py -> scripts/finetune.py and verify_head_grads.py -> verify_grads.py (both cover more than the head now) - finetune.py runs the 2 ensemble members one at a time by default: same loss, head-only 0.25° step 20.0 -> 12.3 GiB temp (+1.8 GiB arguments), ~5 % slower; 10 steps ran in an RTX 5090 pool capped to a 24 GB card's (peak 14.06 GiB) - cloud/modal_canary.py installs the [weathernext] extra (needed since 0.1.1); canary on A100 / A10G / L4 / T4: all run and fit the 8 GB pool, fp32 vs a fresh RTX 5090 run 1.7e-5 / 1.9e-5 / 9.7e-5; the A100 TF32-level deviation (1.32e-2) reproduced as on v0.1.0, recorded in docs/validation.md §4 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Head-only (5.6M parameters, within a 24 GB card's pool) and whole-model (184M, 32 GB card) fine-tuning, with the memory and time per step from docs/validation.md §8-§9, the prototype's limits, and that the "same model" claims cover inference only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Makes WeatherNext 2 trainable on one GPU, as an opt-in addition: inference is unchanged and re-validated.
What's new
pallas_attn.py): FlashAttention-2 backward as two more tiled Pallas kernels (dQ per query block, dK/dV per key block over the transposed tile list, which is the forward's for the symmetric k-hop mask). One layer of the real 0.25° mask: 42 ms forward+backward at TF32 vs 229 ms for XLA autodiff.faster_weathernext.finetune: the official training loss (fair CRPS, 2 members), head-only (mesh_to_grid_gnn+grid_decoder, 5.6M params) or every parameter (split(params, "all")); official-format checkpoints.remat(blocked GNN/encoder/decoder blocks, transformer layers, XLA attention chunks) andsequential_samples(the 2 ensemble members in anhk.scan, one at a time in the backward).ifs.training_example,scripts/finetune.py,scripts/verify_grads.py,cloud/modal_grads.py,attn_bench.py --grad;[train]extra (optax).Validation (docs/validation.md §8–§9)
strict_fp32at the floor, 3xTF32 1.7× (smooth loss) / 5× (CRPS).Known, not addressed here
🤖 Generated with Claude Code