Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Feature Request] Create a BatchNorm layer #5

Open
ShreyasKhandekar opened this issue Oct 1, 2024 · 1 comment
Open

[Feature Request] Create a BatchNorm layer #5

ShreyasKhandekar opened this issue Oct 1, 2024 · 1 comment

Comments

@ShreyasKhandekar
Copy link
Contributor

Pytorch has a few normalization layers that ChAI does not currently implement. See https://pytorch.org/docs/stable/nn.html#normalization-layers for a list of some of the layers.

I think the most important one to implement first is BatchNorm1d which does batch normalization over a 2D or 3D input.
We can then move to higher dimensions like BatchNorm2d/3d when needed.
ShreyasKhandekar pushed a commit to ShreyasKhandekar/ChAI that referenced this issue Oct 18, 2024
@jeremiah-corrado
Merge pull request Iainmon#5 from jeremiah-corrado/locale-sharded-mod…
50701a5 
…ules

Adds support for distributing sub-modules across a given list of target locales. When a module is loaded with a `targetLocales` argument as follows, and its parent module has sub-modules, then those submodules are distributed evenly across locales.

```chapel
const model = loadModel(specFile = "spec.json", weightsFolder = "weights/", targetLocales = Locales);
```

For example, if there are 10 sub-modules at the first layer of the module tree, and the program is run with 2 locales, then the first 5 are loaded on Locale 0, and the next 5 are loaded on Locale 1. I'm referring to this as the "sharded" mode of multi-locale inference. It is intended for scenarios where a model cannot fit into memory on a single locale

This PR also adds tests for the Sharded and Cloned multi-locale inference modes.

**Future work:**
* consider the sub-module sizes (in memory) when distributing. The current design essentially assumes that all sub-modules are the same size
* support distributed tensors, allowing individual modules to be distributed across multiple locales
  * make the two modes of distribution orthogonal in the API s.t. they can be used independently to distribute very large models that are either "deep" (many layers), or "wide" (many parameters within layers), or both.
@ShreyasKhandekar
Copy link
Contributor Author

I have a rudimentary BatchNorm layer which can only do eval in my resnet implementation which should be factored out and made into a full fledged layer

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@ShreyasKhandekar