it's indeed the libomp.so from conda in the nightly wheel is causing the issue. I replaced it with debian libgomp, the performance is restored, 2x improved.
I'm checking how I can switch to libgomp for the aarch64 nightly and release wheel building.

cc @malfet

@snadampal so, should we package libomp from debian in our build scripts?

I would love to help get away from these conda packaged deps and instead use something more OS native (i.e. - build in a container using whats provided via apt, yum, dnf, etc.)

The issue is observed for the pytorch 2.3 release candidate wheels as well.
pip3 install torch==2.3.0 --index-url https://download.pytorch.org/whl/test/cpu

In the release wheel package, I see we are already packaging both the omp libraries.

libgomp-0f9e2209.so.1.0.0 --> from debian
libomp-b8e5bcfb.so. --> from conda

but looks like all the libraries are linked to libomp-b8e5bcfb.so. so, it might be more than packaging, we will have to switch the dependent libraries too.
for experimenting, I renamed libgomp-0f9e2209.so.1.0.0to libomp-b8e5bcfb.so and saw the performance restored.

@bryantbiggs I think this is a general direction build process is going towards: no Anaconda, just use what's in pypa docker

I will check if the new scripts can be updated to remove conda dependency.

Otherwise we anyway have the fallback option: The old scripts https://github.com/pytorch/builder/blob/main/aarch64_linux/aarch64_wheel_ci_build.py are native manylinx os builds. they are being maintained, so, we can switch to them for the CD.

one small correction to my previous statement. we are packaging only the coda libomp.so for the nightly and release wheels, not both conda and manylinux.

In the release wheel package, I see we are already packaging both the omp libraries.

coming to the solution, it's the same.
if we replace libomp.so with libgomp.so, the performance is back , but I don't think it's a good direction just to replace omp library with manylinux but keeping rest from conda, they might cause problems somewhere else.
so, I'm planning to switch everything to manylinux and pip.

looking at how the wheel building scripts are integrated into nightly wheel workflow, everything is happening inside the manylinux docker. and the docker is missing many packages including openblas.
btw, even the older scripts are using conda for those missing packages, but only difference is, they are embedding the libgomp explicitly from host, instead of the one from conda.
as @malfet hinted, packaging libgomp into the wheel might be the only way for now.

any other thoughts?

@snadampal In original script I build OpenBLAS from source, because one that comes with OS was lacking OpenMP integration. And the only reason to use conda were to install cmake and ninja that were missing in PyPI at the time. Now one can(and should) completely eliminate Conda dependency for wheel builds

Right now the scripts are using manylinux2014 os which is pretty old, comes with python 3.6, that's why we had to rely on conda for lot of packages. Given the EOL for it is June 30th, 2024. (link), I'm upgrading the docker to manylinux_2_28, the latest one. this might remove conda dependency. I'm trying it.

it's becoming more involved than I initially thought. manylinux 2_28 comes with gcc-12 with which pytorch compilation is failing on aarch64.
@bryantbiggs , would you be able to give it a try?

ya, I'll see if I can take a look - still working my through building from source without conda for CUDA based build

I have upgraded the docker to manylinux 2_28 and removed conda dependency completely, everything installed from manylinux or pypi. This solves the libomp performance issues.

here is the draft PR:
#1781

I had to disable pytorch tests building, via BUILD_TEST=0 to make it work, that's where GCC-12 breaks coming from, I will look into it next.

I have fixed the pytorch test build issue, in fact it seems to be a known issue pytorch/pytorch#99278, and there was a PR too pytorch/pytorch#99468.
Looks like it was observed on x86 as well, so, it would be good if the PR can be merged.

With this PR, torch build is working fine in manylinux 2_28 docker with gcc12 toolchain.

For now, I'm using gcc-11 toolchain on manylinux2_28, so, I'm not blocked on PyTorch test build PR mentioned above.

I'm observing that compared to the default llvm libomp used in official torch wheels, gnu libgomp is giving 2x better perf for torch compile mode, but 10% slower perf for eager mode.
Unless this 10% slowness is root-caused, it's not good to switch to libgomp. I'm checking why the 10% drop, more importantly why libomp is 2x slower for torch.compile, is there any mixup in torch dynamo?

Looks like there is no clear winner for libgomp vs libompon aarch64 linux. We ran several benchmarks from huggingface, and plotted libgomp perf as the reference blue line, and the relative libomp perf as orange bars. From the chart, it's clear that while libomp is better for few models, libgomp is better in general, so, we are going ahead with switching to libgomp for aarch64 linux, which has already shown upto 2x better perf for torch.compile.