Optimize Rocm backend + remove backend option - #3
Merged
Conversation
On MI250X use_compile made training 2.5x slower (44 vs 18 s/epoch). Kernel profiling found two causes: rocBLAS handles the tiny 9x9 per-pair contraction matrices badly, and the backward of the c[t1, t2] gather is index_put_(accumulate=True) - millions of pairs atomically accumulating into a ~1300-element tensor (11.4 ms, 51% of the step on MI250X). The new mulsum backend gathers the coefficient tables with a one-hot matmul (forward adds exact zeros; backward becomes one clean GEMM, no atomics) and contracts with broadcast-multiply + sum, which Inductor fuses into a single Triton kernel with no BLAS calls. auto now resolves to mulsum under use_compile on all platforms, and to loop eager on ROCm. The make_fx compiled-autograd force path uses mulsum too. MI250X full training: 3.1 s/epoch (was 44 compiled / 18 eager), loss digit-for-digit identical to the eager baseline. A2000: 5.7 vs 8.4 ms/step against the old compiled-bmm path; still faster at 16 types. New tests: loop/bmm/mulsum numerical equivalence (float64, 1e-10) and auto-resolution incl. ROCm.
Policy from an 8-variant x {1,2,4,8,16,32,64,87}-type sweep (NEP89
dataset subsets) on MI250X, V100 and A2000:
- GPU + compile: mulsum at every type count (fastest everywhere).
- GPU eager: mulsum up to 32 types; bmm above (eager mulsum keeps the
one-hot matrix for the backward - ~8 GiB at 64+ types - while bmm
stays flat and is the fastest non-mulsum eager path). loop is never
optimal on any GPU tested and leaves the GPU auto policy.
- CPU/MPS: unchanged (loop below 20 types, bmm above).
- mulsum internal switch: one-hot matmul everywhere on ROCm (atomics
are the bottleneck at low type counts: gather 53-303 ms vs one-hot
8-41 ms/step) and up to 32 types on CUDA; plain gather above 32 on
CUDA (V100 87 types: 42.6 vs 58.2 ms/step and 0.1 vs 2.5 GiB).
Both forms are bit-identical.
The backend parameter is removed from train_nep, train_nep_sharded and
the predict entry points - resolution is device-aware and automatic.
The q_scaler pass follows the eager policy (its old loop path took
40+ s at 87 types; mulsum/bmm take ~1 s on V100). Banner now prints
only the force mode.
Also fixes CompiledAutogradForce on torch builds that do not
auto-import torch.backends.opt_einsum (ROCm 2.7.x) - make_fx autograd
now works on AMD and is the fastest path there at 64+ types.
New tests: mulsum one-hot/gather vs bmm exact match; auto-policy
matrix. Full suite: 121 passed.
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
1 participant
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.
Optimize Rocm backend + remove backend option