Kernels for GroupNorm #353
Merged
Conversation
This file contains 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
lancerts
reviewed
Nov 7, 2024
| c2 += tl.sum(wdy) | ||
|
|
||
| # Need to ensure additions to the same channel are atomic | ||
| tl.atomic_add(DW_ptr + channel_idx, dW.to(dtype)) |
There was a problem hiding this comment.
@ByronHsu is it possible for us to test on multiple GPU, specifically around
scope (str, optional) – Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are “gpu” (default), “cta” (cooperative thread array, thread block), or “sys” (stands for “SYSTEM”). The default value is “gpu”.
whether the default value works for multi-gpu.
There was a problem hiding this comment.
what kind of testing? run on a 4 GPUs env to ensure the kernel working fine on a single GPU? Not sure how this is related to multi gpu. My understanding is the kernel only happens on 1 gpu
|
Very solid PR! |
lancerts
previously approved these changes
Nov 7, 2024
lancerts
approved these changes
Nov 7, 2024
|
@pramodith can you update the readme to include groupnorm |
Will do tomorrow! |
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.
Summary
Implementation of group norm that achieves output parity with torch's
GroupNorm.This is feature is a part of #285
Details
The formulas/equations involved in
GroupNormare the same asLayerNorm/BatchNorm. The main differences lie in the axis along which the mean and std are computed + the dimensions of the Affine transformation parameters.In group norm W and B are of shape
(n_channels), however the mean and std are calculated over all the channels in a given group.Testing Done
Testing was done on a A100 PCIE and a A100 SXM-4.
We see an increase in speed, while the total memory used remains about the same. Note that benchmarking was done using a batch size of 128, Hidden Dim size of 512 and the number of channels per group fixed at 4.
These results look very similar to the layer norm benchmark too.
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence