Uh oh!
There was an error while loading. Please reload this page.
Adding documents to TE/JAX - #87
Merged
Merged
Conversation
zlsh80826
commented
Mar 9, 2023
Collaborator
/te-ci |
ksivaman
commented
Mar 9, 2023
Member
@mingxu1067@jeng1220 Looks like there are a lot of duplicate commits here, could you please rebase with main? |
ksivaman
self-requested a review
March 9, 2023 07:14
ksivaman
commented
Mar 9, 2023
Member
Since #54 is merged now, we can remove the WIP tag too |
jeng1220
commented
Mar 9, 2023
Contributor
@mingxu1067 , |
mingxu1067
commented
Mar 9, 2023
CollaboratorAuthor
mingxu1067force-pushed
the
mingh/te_docs
branch
2 times, most recently
from
March 9, 2023 08:44
e823240 to
d449c12Comparejeng1220
reviewed
Mar 9, 2023
Comment on lines
5
to
10
Contributor
There was a problem hiding this comment.
@mingxu1067 ,
Could you help to make import order to be ordered alphabetically?
zlsh80826
commented
Mar 9, 2023
Collaborator
/te-ci |
jeng1220
reviewed
Mar 9, 2023
Contributor
There was a problem hiding this comment.
The copyright statement needs to be updated
ksivaman
commented
Mar 9, 2023
Member
Could we combine #88 with this PR? I think the scope of this PR is documentation and that change falls within this. I want to avoid burdening the ci with a bunch of small PRs especially when they can be combined. This would speed up the dev/review/merge process. @mingxu1067 |
mingxu1067
commented
Mar 10, 2023
CollaboratorAuthor
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
zlsh80826
commented
Mar 13, 2023
Collaborator
/te-ci |
| ------- | ||
| .. autoapiclass:: transformer_engine.jax.LayerNorm(epsilon=1e-6, layernorm_type='layernorm', **kwargs) | ||
| :members: __call__ |
Member
There was a problem hiding this comment.
Suggested change
| :members: __call__ | |
| :members: __call__ |
| axis name in Mesh used to shard batch along. | ||
| if it is None, then disabling data parallelism. | ||
| The axis name in Mesh used to shard batches along. | ||
| If it is None, then disabling data parallelism. |
Member
There was a problem hiding this comment.
Suggested change
| IfitisNone, thendisablingdataparallelism. | |
| IfitisNone, thendataparallelismisdisabled. |
| axis name in Mesh used to split model tensor along. | ||
| if it is None, then disabling tensor parallelism. | ||
| The axis name in Mesh used to split the hidden dimensions along. | ||
| If it is None, then disabling tensor parallelism. |
Member
There was a problem hiding this comment.
Suggested change
| IfitisNone, thendisablingtensorparallelism. | |
| IfitisNone, thentensorparallelismisdisabled. |
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Comment on lines
+194
to
+195
| The key in given RNGs via flax.linen.Module.apply that | ||
| to generate Dropout masks in the core attention. |
Member
There was a problem hiding this comment.
This doesn't read too well. Maybe just some grammar fix
CollaboratorAuthor
There was a problem hiding this comment.
Fixed as the suggestion below
Comment on lines
+194
to
+195
| The key in given RNGs via flax.linen.Module.apply that | ||
| to generate Dropout masks in the core attention. |
Member
There was a problem hiding this comment.
Suggested change
| ThekeyingivenRNGsviaflax.linen.Module.applythat | |
| togenerateDropoutmasksinthecoreattention. | |
| ThekeyinthegivenRNGsviaflax.linen.Module.applythatis | |
| usedtogenerateDropoutmasksinthecoreattention. |
CollaboratorAuthor
There was a problem hiding this comment.
Yes, it is. Change to the suggestion. THX
| kernel_init: Initializer, default = | ||
| flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'normal') | ||
| used for initializing weights of QKV and Output projection weights. | ||
| Used for initializing weights of QKV and Output projection weights. |
Member
There was a problem hiding this comment.
Suggested change
| UsedforinitializingweightsofQKVandOutputprojectionweights. | |
| UsedforinitializingtheQKVandOutputprojectionweights. |
| If set to False, the layer will not learn additive biases. | ||
| bias_init: Initializer, default = flax.linen.initializers.zeros | ||
| used for initializing bias of QKVO projections, only works when :attr:`use_bias=True`. | ||
| Used for initializing bias of QKVO projections, it only works when :attr:`use_bias=True`. |
Member
There was a problem hiding this comment.
Suggested change
| UsedforinitializingbiasofQKVOprojections, itonlyworkswhen :attr:`use_bias=True`. | |
| UsedforinitializingbiasofQKVOprojections, onlyusedwhen :attr:`use_bias=True`. |
| use_bias: bool, default = False | ||
| indicate whether to enable bias shifting for QKVO projections. | ||
| if set to False, the layer will not learn additive biases. | ||
| Indicate whether to enable bias shifting for QKVO projections. |
Member
There was a problem hiding this comment.
Suggested change
| IndicatewhethertoenablebiasshiftingforQKVOprojections. | |
| IndicatewhetherornottoenablebiasshiftingforQKVOprojections. |
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
| A helper to update Flax's Collection. | ||
| Collection = [dict, FrozenDict] | ||
| Collection = [dict, Flax's FrozenDict] |
Member
There was a problem hiding this comment.
Suggested change
| Collection= [dict, Flax'sFrozenDict] | |
| Collection= [dict, flax.core.frozen_dict.FrozenDict] |
| updated_scale_inv = 1/updated_scale | ||
| Collection = [dict, FrozenDict] | ||
| Collection = [dict, Flax's FrozenDict] |
Member
There was a problem hiding this comment.
Suggested change
| Collection= [dict, Flax'sFrozenDict] | |
| Collection= [dict, flax.core.frozen_dict.FrozenDict] |
| ---------- | ||
| scale_factor : float, default = 1.0 | ||
| scale the inputs along the last dimension before running softmax. | ||
| Scale the inputs along the last dimension before running softmax. |
Member
There was a problem hiding this comment.
Is this scaling only across the last dimension? The whole softmax input is scaled, right?
Uh oh!
There was an error while loading. Please reload this page.
Member
There was a problem hiding this comment.
Suggested change
| FC1andFC2. Itonlyworkswhen :attr:`use_bias=True`. | |
| FC1andFC2. Itonlyusedwhen :attr:`use_bias=True`. |
Member
There was a problem hiding this comment.
@mingxu1067 This pattern actually exists throughout, could you please go through all the cases here change "only works" -> "only used" so that it gives the correct picture?
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
| scale the inputs along the last dimension before running softmax. | ||
| softmax_type : SoftmaxType, default = 'layernorm' | ||
| indicate the type of softmax. | ||
| Scale the whole (inputs + bias) before running softmax. |
Member
There was a problem hiding this comment.
Suggested change
| Scalethewhole (inputs+bias) beforerunningsoftmax. | |
| Scalarfortheinputtosoftmax. |
| Jax | ||
| ======= | ||
| Types |
Member
There was a problem hiding this comment.
I think this gives a better idea of what these are supposed to be
CollaboratorAuthor
There was a problem hiding this comment.
Remove categories as PyTorch to make the style consistent.
Comment on lines
+26
to
+42
| .. autoapiclass:: transformer_engine.jax.DenseGeneral(features, layernorm_type='layernorm', use_bias=False, **kwargs) | ||
| :members: __call__ | ||
| .. autoapiclass:: transformer_engine.jax.LayerNormDenseGeneral(features, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs) | ||
| :members: __call__ | ||
| .. autoapiclass:: transformer_engine.jax.LayerNormMLP(intermediate_dim=2048, layernorm_type='layernorm', epsilon=1e-6, use_bias=False, **kwargs) | ||
| :members: __call__ | ||
| .. autoapiclass:: transformer_engine.jax.RelativePositionBiases(num_buckets, max_distance, num_heads, **kwargs) | ||
| :members: __call__ | ||
| .. autoapiclass:: transformer_engine.jax.MultiHeadAttention(head_dim, num_heads, **kwargs) | ||
| :members: __call__ | ||
| .. autoapiclass:: transformer_engine.jax.TransformerLayer(hidden_size=512, mlp_hidden_size=2048, num_attention_heads=8, **kwargs) | ||
| :members: __call__ |
Member
There was a problem hiding this comment.
All of these modules are showing up in the Functions category below in the generated docs. @mingxu1067
CollaboratorAuthor
There was a problem hiding this comment.
Remove categories as PyTorch to make the style consistent.
nouiz
reviewed
Mar 13, 2023
| .. toctree:: | ||
| pytorch | ||
| jax |
Collaborator
There was a problem hiding this comment.
The main README should also mention JAX.
Grep for pytorch in that document and add JAX at those places and add a JAX example too as there is a PyTorch example.
CollaboratorAuthor
There was a problem hiding this comment.
I think README.md need more refactor to show all FWs supported by TE. Diretort adding JAX to all Pytorch appearance might mess up reading. We can submit a split PR to refactor README.md.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
ksivaman
approved these changes
Mar 14, 2023
ksivaman
commented
Mar 14, 2023
Member
/te-ci |
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
Mar 16, 2023
* Updated TE/JAX docs Signed-off-by: Ming Huang <mingh@nvidia.com> * Adding TE/JAX docs' rst files Signed-off-by: Ming Huang <mingh@nvidia.com> * Set DType as pybind11::module_local() to avoid generic_type errors. Signed-off-by: Ming Huang <mingh@nvidia.com> * Updating license and exporting more modules Signed-off-by: Ming Huang <mingh@nvidia.com> * Adopting autoapi and removing enum_tools. Signed-off-by: Ming Huang <mingh@nvidia.com> * Fix typo Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Make jax.rst be style consistent. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from code review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Update the description of Softmax Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Removed categories in catalog as PyTorch Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> --------- Signed-off-by: Ming Huang <mingh@nvidia.com> Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
cyanguwa pushed a commit
to cyanguwa/TransformerEngine
that referenced
this pull request
Mar 31, 2023
* Updated TE/JAX docs Signed-off-by: Ming Huang <mingh@nvidia.com> * Adding TE/JAX docs' rst files Signed-off-by: Ming Huang <mingh@nvidia.com> * Set DType as pybind11::module_local() to avoid generic_type errors. Signed-off-by: Ming Huang <mingh@nvidia.com> * Updating license and exporting more modules Signed-off-by: Ming Huang <mingh@nvidia.com> * Adopting autoapi and removing enum_tools. Signed-off-by: Ming Huang <mingh@nvidia.com> * Fix typo Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Make jax.rst be style consistent. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from code review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Update the description of Softmax Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Removed categories in catalog as PyTorch Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> --------- Signed-off-by: Ming Huang <mingh@nvidia.com> Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwa pushed a commit
to cyanguwa/TransformerEngine
that referenced
this pull request
Apr 1, 2023
* Updated TE/JAX docs Signed-off-by: Ming Huang <mingh@nvidia.com> * Adding TE/JAX docs' rst files Signed-off-by: Ming Huang <mingh@nvidia.com> * Set DType as pybind11::module_local() to avoid generic_type errors. Signed-off-by: Ming Huang <mingh@nvidia.com> * Updating license and exporting more modules Signed-off-by: Ming Huang <mingh@nvidia.com> * Adopting autoapi and removing enum_tools. Signed-off-by: Ming Huang <mingh@nvidia.com> * Fix typo Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Make jax.rst be style consistent. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Fixing doc statements as the suggestion from code review. Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Update the description of Softmax Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> * Removed categories in catalog as PyTorch Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> --------- Signed-off-by: Ming Huang <mingh@nvidia.com> Signed-off-by: Ming-Xu Huang <mingh@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Charlene Yang <charleney@nvidia.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 freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
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.
No description provided.