[ONNX] RMS Norm

[justinchuby]

• Implement rms norm using onnx RMSNormalization-23

• Use the correct eps for float32
  

  pytorch/aten/src/ATen/native/cuda/layer_norm_kernel.cu

  Lines 1844 to 1866 in eaadd12

  	auto acc_type = at::toAccumulateType(input.scalar_type(), /*is_cuda=*/true);
  	double eps_val;
  	if (acc_type == at::ScalarType::Float) {
  	eps_val = eps.value_or(std::numeric_limits<float>::epsilon());
  	} else {
  	eps_val = eps.value_or(std::numeric_limits<double>::epsilon());
  	}
  	Tensor Y = at::native::empty_like(
  	*X,
  	std::nullopt /* dtype */,
  	std::nullopt /* layout */,
  	std::nullopt /* device */,
  	std::nullopt /* pin_memory */,
  	LEGACY_CONTIGUOUS_MEMORY_FORMAT);
  	Tensor rstd = at::empty({M}, X->options().dtype(acc_type));
  	if (M > 0) {
  	RmsNormKernelImpl(*X, *gamma, M, N, eps_val, &Y, &rstd);
  	}
  	const auto input_shape = input.sizes();
  	const size_t axis = input.dim() - normalized_shape.size();

  
  

• Created facility to run tests with the reference runtime by extending ONNXProgram and assert_onnx_program.

Fix #159257

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159377

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 340dc41 with merge base f4bfac1 ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable) (gh) (#158876)
  sccache: error: couldn't connect to server

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[titaiwangms]

Where do we add the tests for op23 symbolic functions?

[justinchuby]

Where do we add the tests for op23 symbolic functions?

I added tests in small models e2e. I think we can extend the opinfo tests later.

[justinchuby]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 1 checks: pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[justinchuby]

In test: export with opset 23

[pytorchmergebot]

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

[justinchuby]

test_rms_norm needs 5e-5

[justinchuby]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 0 checks:

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

[justinchuby]

The default eps 1e-5 we use is too large. Reading the decomposed program we see the eps used in pytorch:

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[2, 5, 7, 3]"):
             # File: /Users/justinc/Documents/GitHub/pytorch/test/onnx/exporter/test_small_models_e2e.py:772 in forward, code: return torch.nn.functional.rms_norm(x, [7, 3])
            pow_1: "f32[2, 5, 7, 3]" = torch.ops.aten.pow.Tensor_Scalar(x, 2)
            mean: "f32[2, 5, 1, 1]" = torch.ops.aten.mean.dim(pow_1, [3, 2], True);  pow_1 = None
            add: "f32[2, 5, 1, 1]" = torch.ops.aten.add.Scalar(mean, 1.1920928955078125e-07);  mean = None
            rsqrt: "f32[2, 5, 1, 1]" = torch.ops.aten.rsqrt.default(add);  add = None
            mul: "f32[2, 5, 7, 3]" = torch.ops.aten.mul.Tensor(x, rsqrt);  rsqrt = None
            type_as: "f32[2, 5, 7, 3]" = torch.ops.aten.type_as.default(mul, x);  mul = x = None
            return (type_as,)
            
Graph signature: 
    # inputs
    x: USER_INPUT
    
    # outputs
    type_as: USER_OUTPUT
    
Range constraints: {}

[justinchuby]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 1 checks: pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here