Fix numerical instability in vector_norm when receiving large size tensor

[maybeLee]

I made this pull request because I encountered the numerical instability of the API torch.linalg.vector_norm when I used it to process a tensor with large shape such as ((1, 32, 224, 224, 160)).

Here is the code to reproduce the issue:

import numpy as np
import torch
x1 = np.ones((1, 32, 224, 224, 160))
ord = 2
print(np.size(x1))  # 256901120
res1 = torch.linalg.vector_norm(torch.tensor(x1, dtype=torch.float32), ord=ord)
res2 = torch.linalg.vector_norm(torch.tensor(x1, dtype=torch.float64), ord=ord)

print(res1, res2)  # tensor(11585.2373) tensor(16028.1353, dtype=torch.float64)
print(f"Expected result: {np.sqrt(np.size(x1))}")  # 16028.135262718493

When checking the code, I find that the issue may lie here

pytorch/aten/src/ATen/native/cpu/ReduceOpsKernel.cpp

Line 181 in 3d20cc1

acc_vec += data_vec * data_vec;

Where precision loss occurs during the accumulation of acc_vec when the acc_vec is too large while data_vec * data_vec is relatively small.

The fix I am applying is to use the Kahan summation algorithm (https://en.wikipedia.org/wiki/Kahan_summation_algorithm).

Please note that after this fix is applied, the average execution time is changed from 0.08817670345306397 to 0.15041379928588866.

Detailed Code

import numpy as np import time import torch

x1 = np.ones((1, 32, 224, 224, 160))

ord = 2
time_list = []
for i in range(10):
s_time = time.time()
res1 = torch.linalg.vector_norm(torch.tensor(x1, dtype=torch.float32), ord=ord)
time_list.append(time.time() - s_time)
print(np.mean(time_list), np.std(time_list))

Please let me know if the fix can be improved.

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 13 New Failures

As of commit e55d019 with merge base 16cb5d4 ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  inductor/test_torchinductor_dynamic_shapes.py::DynamicShapesCpuTests::test_dist_bf16_dynamic_shapes_cpu
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  inductor/test_torchinductor_dynamic_shapes.py::DynamicShapesCpuTests::test_dist_bf16_dynamic_shapes_cpu
• pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.11-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.11-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.12-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.8-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.8-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.8-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_linalg.py::TestLinalgCPU::test_norm_bfloat16_and_half_cpu_bfloat16
• pull / linux-jammy-py3.10-clang15-asan / test (default, 1, 6, linux.4xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16
• pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh)
  inductor/test_torchinductor.py::CpuTests::test_dist_bf16_cpu
• pull / linux-jammy-py3.8-gcc11 / test (default, 1, 3, linux.2xlarge) (gh)
  test_reductions.py::TestReductionsCPU::test_dim_reduction_lastdim_cpu_bfloat16

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

[linux-foundation-easycla]

• ❌ - login: @maybeLee . The commit (e55d019) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please submit a support request ticket.

[CaoE]

@maybeLee Thanks for your fix. If we perform group reduction on such large size cases, accuracy should be improved without introducing too much overhead.
For example (pseudocode):

acc_t acc_buffer[group_size]:
for (int64_t g = 0; g <  group_size; g++) {
  // Do reduction for each group
  acc_buffer[g] = group_reduce(...)

}

// Do reduction finally
double acc_value = reduce_finaliy(acc_buffer)
result = scalar_t(std::sqrt(acc_value);

[maybeLee]

Hi @CaoE , thanks for your reply. I think the original code has already done such group reduction?

pytorch/aten/src/ATen/native/cpu/ReduceOpsKernel.cpp

Lines 246 to 253 in 7b6e354

	for (; d < size - (size % Vec::size()); d += Vec::size()) {
	Vec data_vec = Vec::loadu(self_data + d);
	norm_two_reduce_step(acc_vec, data_vec);
	}
	acc_vec.store(buffer);
	for (int j = 1; j < fVec::size(); j++) {
	buffer[0] = buffer[0] + buffer[j];
	}

[CaoE]

@maybeLee This is also a group reduce, and it is divided by vec size. But even if the vec size is 32, the size of each group is also large: 256901120/32=8028160. We can set a parameter, e.g., group_size = 32768, and further group the reduction externally by group_size (32768 elements are reduced in the current way in each group).

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.