Repository navigation
perf(batchnorm): device-sized grid, one-pass statistics, finalize kernels - #18
Merged
Merged
Conversation
_get_shape_meta capped the reductions at 64 block rows, so at C=32 they ran 64 blocks of 256 threads: one per CU on the RX 9070 XT, which has room for 256. The grid now fills every resident block slot once, the same rule as choose_dweight_split_k: slots = mp_count * (max_threads_per_mp // 256) grid_y = max(1, min(ceil(N / 8), slots // grid_x)) 9070 XT (32 WGPs x 2048 threads): 256 blocks at every C. RTX 3050 (16 SMs x 1536 threads): 96 blocks, so C=128 drops from 256. At 256x32x32x32 on the 9070 XT the mean kernel goes from 0.159 to 0.063 ms and grad_sums from 0.153 to 0.078 ms. BatchNorm forward drops 0.19 ms per CIFAR-10 step. #pragma unroll 4 in the N loop made no difference, so it isn't used.
The forward read x three times (mean pass, variance pass, elementwise
pass) and launched about 15 kernels: two zero-inits, the divides, the
sqrt, and the running-stat EMA. It now launches three:
- stats: one pass accumulating d = x - x[0, c] and d*d. Shifting by a
sample of the channel keeps E[d^2] - E[d]^2 from cancelling when
|mean| >> std, without the per-element divide Welford needs. Each
block writes its pair to partials[2, grid_y, C] instead of using
atomicAdd, so there is no zero-init and the result is deterministic.
grad_sums shares the template and writes partials the same way (they
are summed with partials.sum here, until the next commit).
- bn_forward_finalize: reduces the partials and writes stats[4, C] =
(mean, var, inv_std, gamma * inv_std), then updates running_mean and
running_var in place. batch_mean, batch_var, inv_std and the new
scale attribute are row views of stats.
- _bn_apply_ew: out = fmaf(x - mean, scale, beta), shared by training
and inference. GPU inference moves from the folded x*scale + bias to
this centered form, which doesn't add rounding error growing with
|mean|/std. The CPU fallbacks are unchanged.
get_bn_mean_kernel, get_bn_var_kernel, _bn_forward_ew and
_bn_inference_ew are gone.
_forward_gpu now casts running_mean/running_var to fp32. Training on
NumPy under a float64 policy leaves them float64 and Model.to('cupy')
keeps that dtype; GPU inference then raised TypeError, and the in-place
update would have misread them.
BatchNorm forward on the 9070 XT: 1.018 -> 0.399 ms GPU and 0.388 ->
0.148 ms host per CIFAR-10 step.
Tests: the large-offset test now checks batch_var against the float64
variance of the same data (it fails on the naive E[x^2] - E[x]^2 and on
the old two-pass kernels); the forward/backward parity loops gain a
ragged (3, 37, 41, 40) shape with non-trivial gamma/beta; a new test
checks two GPU runs are bit-identical; and a regression test covers
float64 running statistics.
… kernel After grad_sums, the backward ran about ten tiny launches for dgamma, dvar, dmu and the scaled coefficients. bn_backward_finalize now reduces the grad_sums partials (sum of dy, sum of dy * (x - mean)) and writes grads[4, C] = (dgamma, dbeta, k1, k2) in one launch, with dx = dy * scale + (x - mean) * k1 + k2 k1 = -scale * inv_std^2 * sum(dy * (x - mean)) / N k2 = -scale * dbeta / N dweights and dbiases are row views of grads. The backward is three launches: grad_sums, finalize, elementwise. The two finalize kernels share the partials reduction (sum_partials) in the kernel prelude. BatchNorm backward on the 9070 XT: 1.067 -> 0.842 ms GPU and 0.358 -> 0.139 ms host per CIFAR-10 step. The elementwise pass is now at the DRAM floor (a plain cp.add with the same reads and writes takes as long). Model.manual_seed's note no longer lists BatchNorm among the kernels that accumulate with atomics; it now names MaxPool2d's backward with overlapping windows (xp.add.at) next to Conv2d.
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.
Context
BatchNorm was ~29 % of GPU time per CIFAR-10 step on the RX 9070 XT (forward
1.02 ms + backward 1.07 ms over 5 layers). The reductions launched too few
blocks, the forward read its input three times, and each direction launched
about 15 tiny kernels. This is Part 1 of the BatchNorm plan; fusing
ReLU/LeakyReLU into BatchNorm (Part 2) follows in a separate change. Three
commits, one per change below, each passing the suite on its own. None are
mechanical.
Reduction grid
The reductions were capped at 64 block rows, so at C=32 they ran 64 blocks: one
per CU on the 9070 XT, which can hold 256. The grid now fills every resident
block slot once (
mp_count * max_threads_per_mp / 256), the same rule aschoose_dweight_split_k. That is 256 blocks on the 9070 XT and 96 on theRTX 3050, where C=128 drops from 256 blocks to 96.
One-pass statistics
The forward is three launches instead of ~15:
d = x - x[0, c]andd*d. Shiftingby a sample keeps the one-pass variance from cancelling when |mean| >> std,
without the per-element divide Welford needs. Blocks write
partials[2, grid_y, C]instead of usingatomicAdd, so there is nozero-init and the result is deterministic.
stats[4, C](mean, var,inv_std, gamma * inv_std) and updates the running statistics in place.
fmaf(x - mean, scale, beta), shared by training andinference.
Backward finalize
bn_backward_finalizereplaces the ~10 tiny launches aftergrad_sums. Itwrites
(dgamma, dbeta, k1, k2), sodx = dy * scale + (x - mean) * k1 + k2.The backward is three launches. Its elementwise pass is at the DRAM floor: a
plain
cp.addwith the same two reads and one write takes as long.Results (RX 9070 XT,
cifar10_breakdown.py, median of 3)model.train()Kernels alone at 256x32x32x32: the mean + variance passes (0.159 + 0.156 ms)
became one stats pass (0.050 ms) plus a 0.009 ms finalize.
grad_sumswentfrom 0.153 to 0.048 ms. Single-kernel times at this shape vary between runs,
because x plus the output is 64 MB, exactly the Infinity Cache size, so the
breakdown above is the reliable number.
Behavior changes
gradients are bit-identical between runs. The
Model.manual_seednote nownames only Conv2d and MaxPool2d's backward with overlapping windows as
non-deterministic.
(x - mean) * scale + betainsteadof
x * scale + bias. Outputs differ at rounding level, and are moreaccurate when |mean| >> std. The CPU path is unchanged.
kernels were off by more than 1e-4 relative at mean 1e4, std 0.01.
set_precision('float64'), thento('cupy'), GPU inference raisedTypeError. The running statistics arenow cast to fp32 on GPU.
BatchNorm.scale(gamma * inv_std) is set by the GPUforward. The removed
get_bn_mean_kernel/get_bn_var_kernelwereinternal.
Testing
python3 -m unittest discover tests: 1148 tests pass on NumPy and CuPy (67skipped), at each commit.
test_gpu_numerical_stability_large_offsetnow checksbatch_varagainstthe float64 variance of the same data. It fails on a naive E[x^2] - E[x]^2
variant and on the old two-pass kernels.
test_gpu_training_is_deterministic: two GPU forward + backward runs givebit-identical results. Fails 3 out of 3 runs on
main.test_gpu_float64_running_stats: regression test for the fix above. Failson
main(TypeError), and also fails if the new cast is removed (thein-place update misreads fp64).
with non-trivial gamma/beta: N is not a multiple of 8, C is not a multiple
of 32, and the reduction spans many blocks.
the new kernels have only been compiled for ROCm, and the C=128 layer (96
blocks instead of 256) needs a before/after breakdown there.
Follow-ups (out of scope)
Model.finalize(). ReLU is 0.72 msper step, more than BatchNorm forward now takes.