Add decomposition for aten.native_layer_norm_backward op.

[gprateek93]

Signed-Off-By: Prateek Gupta prateek@nod-labs.com

[facebook-github-bot]

Hi @gprateek93!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@fb.com. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[Chillee]

Mostly looks good to me! A couple minor notes, and I'd like to wait for the CI to pass first.

[Chillee]

Tests look good, just fix the flake8 issues (and the other comments and we're good to go!)

[Chillee]

Oh actually, I realized our decomposition tests are currently skipped due to an issue with upstream, I'll test it out locally.

[Chillee]

Actually, it doesn't seem like it passes testing.

To test your decompositions, comment out https://github.com/pytorch/functorch/blob/main/test/test_ops.py#L973 and run

pytest test/test_ops.py::TestDecompositionOpInfoCPU -k layer_norm

[gprateek93]

Actually, it doesn't seem like it passes testing.

To test your decompositions, comment out https://github.com/pytorch/functorch/blob/main/test/test_ops.py#L973 and run

pytest test/test_ops.py::TestDecompositionOpInfoCPU -k layer_norm

I will check and make necessary changes. Thanks!

[gprateek93]

@Chillee Even after removing my commit, the test is failing. I am assuming the error is due to the decomposition of aten.native_layer_norm.

[Chillee]

@gprateek93 What's the error you see? There was a minor precision issue I just fixed.

The error I ran into with the decomposition though seemed more like a logical error

functorch/_src/decompositions.py:82: in native_layer_norm_backward
    b = aten.sum(grad_x_hat, inner_dims, True)
../pytorch/torch/_ops.py:142: in __call__
    return self._op(*args, **kwargs or {})
test/test_ops.py:1143: in __torch_dispatch__
    real_out = func(*tree_map(unwrap_tensor, args), **tree_map(unwrap_tensor, kwargs))
  IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

[gprateek93]

@gprateek93 What's the error you see? There was a minor precision issue I just fixed.

The error I ran into with the decomposition though seemed more like a logical error

functorch/_src/decompositions.py:82: in native_layer_norm_backward
    b = aten.sum(grad_x_hat, inner_dims, True)
../pytorch/torch/_ops.py:142: in __call__
    return self._op(*args, **kwargs or {})
test/test_ops.py:1143: in __torch_dispatch__
    real_out = func(*tree_map(unwrap_tensor, args), **tree_map(unwrap_tensor, kwargs))
  IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

Hi @Chillee, On commenting the line you mentioned, I am not getting any error in the code. But I may have got the logical bug. I am not able to verify with test. Can you check once?

[Chillee]

@gprateek93 Why are you unable to run the tests? Like, what happens when you try running them? If the decomposition passes you should see something like

========================================================================== 4 passed, 3106 deselected, 13 warnings in 6.30s ===========================================================================

When I run the decompositions with the newest commit in the PR there are still some test failures

FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_bfloat16 - RuntimeError: Difference from float64 is larger with decomposition native_layer_nor...
FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_float32 - RuntimeError: native_layer_norm_backward decomposition failed, max rel: 2.2540149688...
FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_float64 - RuntimeError: native_layer_norm_backward decomposition failed, max rel: 1.1596385197.

[gprateek93]

@gprateek93 Why are you unable to run the tests? Like, what happens when you try running them? If the decomposition passes you should see something like

========================================================================== 4 passed, 3106 deselected, 13 warnings in 6.30s ===========================================================================

When I run the decompositions with the newest commit in the PR there are still some test failures

FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_bfloat16 - RuntimeError: Difference from float64 is larger with decomposition native_layer_nor...
FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_float32 - RuntimeError: native_layer_norm_backward decomposition failed, max rel: 2.2540149688...
FAILED test/test_ops.py::TestDecompositionOpInfoCPU::test_decomposition_nn_functional_layer_norm_cpu_float64 - RuntimeError: native_layer_norm_backward decomposition failed, max rel: 1.1596385197.

So I am able to run the test, but I am getting all passed.
==================================== 3 passed, 3021 deselected, 7 warnings in 1.19s =====================================

The command I am running is : pytest test/test_ops.py::TestDecompositionOpInfoCPU -k layer_norm -v

[Chillee]

@gprateek93 Are you sure the decompositions are being run? Can you try changing it to something obviously wrong and see if it fails?

I'm not sure how you installed functorch, but perhaps if you installed it with pip you might need to reinstall it before your python changes will get picked up.

Otherwise, to install it in "development" mode, you can do python setup.py develop.

[gprateek93]

@Chillee I am able to get the failures now, But I am not sure why they are failing. I have tried to create one of the test cases below in google colab and ran it with both torch op and my decomposition, and the results are fairly equal.
sample_example

Could it be some kind of precision issue?

[Chillee]

@gprateek93 For the original PR, no, since the absolute value diff on that is > 2.

For the updated version in your gist, one of the errors might be an unavoidable precision issue. The other one seems to be an actual failure though.

I updated the main branch to output the failing arguments in an error.

[gprateek93]

@gprateek93 For the original PR, no, since the absolute value diff on that is > 2.

For the updated version in your gist, one of the errors might be an unavoidable precision issue. The other one seems to be an actual failure though.

I updated the main branch to output the failing arguments in an error.

Thank you for the updated test support. I can now view which argument is failing, I am working on this.

[gprateek93]

@Chillee apart from the precision issue, other test cases get passed in this current revision.