-
Notifications
You must be signed in to change notification settings - Fork 4.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add a document for leveraging Advanced Matrix Extensions #2439
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/2439
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 52c96e2: NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
✅ Deploy Preview for pytorch-tutorials-preview ready!
To edit notification comments on pull requests, go to your Netlify site settings. |
recipes_source/amx.rst
Outdated
For more detailed information of oneDNN, see `oneDNN`_. | ||
|
||
The operation is fully handled by oneDNN according to the execution code path generated. I.e. when a supported operation gets executed into oneDNN implementation on a hardware platform with AMX support, AMX instructions will be invoked automatically inside oneDNN. | ||
No manual operations are required to enable this feature. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No manual operations are required to enable this feature. | |
Since oneDNN is the default acceleration library for CPU, no manual operations are required to enable the AMX support. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Apply the change.
recipes_source/amx.rst
Outdated
``conv_transpose1d``, | ||
``conv_transpose2d``, | ||
``conv_transpose3d``, | ||
``linear`` |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess we need a special note for quantized linear here that whether AMX kernel is chosen also depends on the policy of the quantization backend. Currently, the x86 quant backend uses fbgemm, not onednn while users can use onednn backend to turn on AMX for linear op. cc @Xia-Weiwen
In general, it is also true that whether to dispatch to AMX kernels is a backend/library choice. The backend/library would choose the most optimal kernels. It is worth noting in this tutorial.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add note.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes. However, I am not sure if it's OK to give such details in tutorial. 🤔
recipes_source/amx.rst
Outdated
|
||
Advanced Matrix Extensions (AMX), also known as Intel® Advanced Matrix Extensions (Intel® AMX), is an extension to the x86 instruction set architecture (ISA). | ||
Intel advances AI capabilities with 4th Gen Intel® Xeon® Scalable processors and Intel® AMX, delivering 3x to 10x higher inference and training performance versus the previous generation, see `Accelerate AI Workloads with Intel® AMX`_. | ||
AMX supports two data types, INT8 and BFloat16, compared to AVX512 FP32, it can achieve up to 32x and 16x acceleration, respectively, see figure 6 of `Accelerate AI Workloads with Intel® AMX`_. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are the speedups only on some particular newer hardware? Is the hardware consumer or enterprise centric?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
AMX is only available from the 4th gen of Xeon (codename sapphire rapids), it is enterprise centric.
recipes_source/amx.rst
Outdated
Confirm AMX is being utilized | ||
------------------------------ | ||
|
||
Set environment variable ``export ONEDNN_VERBOSE=1`` to get oneDNN verbose at runtime. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it would be nice to have some python function like is_x_available()
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add a python function torch.backends.mkldnn.verbose
.
Note: For quantized linear, whether to leverage AMX depends on which quantization backend to choose. | ||
At present, x86 quantization backend is used by default for quantized linear, using fbgemm, while users can specify onednn backend to turn on AMX for quantized linear. | ||
|
||
Guidelines of leveraging AMX with workloads |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would start with this section on how to use it and have the supported ops show up at the bottm
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for your suggestion. Do you mean to move this section above the supported ops ? Like this:
AMX in PyTorch
Guidelines of leveraging AMX with workloads
List supported ops
...
recipes_source/amx.rst
Outdated
Introduction | ||
============ | ||
|
||
Advanced Matrix Extensions (AMX), also known as Intel® Advanced Matrix Extensions (Intel® AMX), is an extension to the x86 instruction set architecture (ISA). |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I realize AMX is lower level than other Intel technologies but it's still worth rationalizing to an end user in a few lines why it's interesting for them to know about AMX vs Intel compiler technologies
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added more introduction to AMX and the benefits it can bring.
recipes_source/amx.rst
Outdated
|
||
Advanced Matrix Extensions (AMX), also known as Intel® Advanced Matrix Extensions (Intel® AMX), is an extension to the x86 instruction set architecture (ISA). | ||
Intel advances AI capabilities with 4th Gen Intel® Xeon® Scalable processors and Intel® AMX, delivering 3x to 10x higher inference and training performance versus the previous generation, see `Accelerate AI Workloads with Intel® AMX`_. | ||
AMX supports two data types, INT8 and BFloat16, compared to AVX512 FP32, it can achieve up to 32x and 16x acceleration, respectively, see figure 6 of `Accelerate AI Workloads with Intel® AMX`_. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we can directly copy the wording from
Compared to 3rd Gen Intel Xeon Scalable processors running Intel® Advanced Vector Extensions 512 Neural Network Instructions (Intel® AVX-512 VNNI), 4th Gen Intel Xeon Scalable processors running Intel AMX can perform 2,048 INT8 operations per cycle, rather than 256 INT8 operations per cycle. They can also perform 1,024 BF16 operations per cycle, as compared to 64 FP32 operations per cycle.
which is a quote from https://www.intel.com/content/www/us/en/products/docs/accelerator-engines/advanced-matrix-extensions/ai-solution-brief.html
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Quoted this.
recipes_source/amx.rst
Outdated
``conv1d``, | ||
``conv2d``, | ||
``conv3d``, | ||
``conv1d``, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why we have 2 sets of conv1d
, conv2d
, conv3d
here
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed typos.
recipes_source/amx.rst
Outdated
``addbmm``, | ||
``linear``, | ||
``matmul``, | ||
``_convolution`` |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_convolution
is not intended to be directly used, start with a _
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Removed _convolution
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please add a "summary" or "conclusion" section to summarize the document.
Added conclusion section. |
@msaroufim Could you please review this doc ? Thanks. |
@ngimel Could you please review this doc ? Thank you. |
@msaroufim Could you please review this doc ? Thank you. |
@kit1980 Could you please review this doc ? Thank you. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A couple editorial fixes for proper HTML rendering.
recipes_source/amx.rst
Outdated
to get higher performance out-of-box on x86 CPUs with AMX support. | ||
For more detailed information of oneDNN, see `oneDNN`_. | ||
|
||
The operation is fully handled by oneDNN according to the execution code path generated. I.e. when a supported operation gets executed into oneDNN implementation on a hardware platform with AMX support, AMX instructions will be invoked automatically inside oneDNN. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The operation is fully handled by oneDNN according to the execution code path generated. I.e. when a supported operation gets executed into oneDNN implementation on a hardware platform with AMX support, AMX instructions will be invoked automatically inside oneDNN. | |
The operation is fully handled by oneDNN according to the execution code path generated. For example, when a supported operation gets executed into oneDNN implementation on a hardware platform with AMX support, AMX instructions will be invoked automatically inside oneDNN. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for your comments.
Fixed.
recipes_source/amx.rst
Outdated
with torch.cpu.amp.autocast(): | ||
output = model(input) | ||
|
||
Note: Use channels last format to get better performance. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note: Use channels last format to get better performance. | |
.. note:: Use channels' last format to get better performance. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
|
||
When the generated graph model runs into oneDNN implementations with the supported operators, AMX accelerations will be activated. | ||
|
||
Note: When using PyTorch on CPUs that support AMX, the framework will automatically enable AMX usage by default. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note: When using PyTorch on CPUs that support AMX, the framework will automatically enable AMX usage by default. | |
.. note:: When using PyTorch on CPUs that support AMX, the framework will automatically enable AMX usage by default. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
- BF16 CPU ops that can leverage AMX: | ||
|
||
``conv1d``, | ||
``conv2d``, | ||
``conv3d``, | ||
``conv_transpose1d``, | ||
``conv_transpose2d``, | ||
``conv_transpose3d``, | ||
``bmm``, | ||
``mm``, | ||
``baddbmm``, | ||
``addmm``, | ||
``addbmm``, | ||
``linear``, | ||
``matmul``, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- BF16 CPU ops that can leverage AMX: | |
``conv1d``, | |
``conv2d``, | |
``conv3d``, | |
``conv_transpose1d``, | |
``conv_transpose2d``, | |
``conv_transpose3d``, | |
``bmm``, | |
``mm``, | |
``baddbmm``, | |
``addmm``, | |
``addbmm``, | |
``linear``, | |
``matmul``, | |
BF16 CPU ops that can leverage AMX: | |
- ``conv1d`` | |
- ``conv2d`` | |
- ``conv3d`` | |
- ``conv_transpose1d`` | |
- ``conv_transpose2d`` | |
- ``conv_transpose3d`` | |
- ``bmm`` | |
- ``mm`` | |
- ``baddbmm`` | |
- ``addmm`` | |
- ``addbmm`` | |
- ``linear`` | |
- ``matmul`` |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
``linear``, | ||
``matmul``, | ||
|
||
- Quantization CPU ops that can leverage AMX: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- Quantization CPU ops that can leverage AMX: | |
Quantization CPU ops that can leverage AMX: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
Confirm AMX is being utilized | ||
------------------------------ | ||
|
||
Set environment variable ``export ONEDNN_VERBOSE=1``, or use ``torch.backends.mkldnn.verbose`` to flexibly enable oneDNN to dump verbose messages. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Set environment variable ``export ONEDNN_VERBOSE=1``, or use ``torch.backends.mkldnn.verbose`` to flexibly enable oneDNN to dump verbose messages. | |
Set environment variable to ``export ONEDNN_VERBOSE=1`` or use ``torch.backends.mkldnn.verbose`` to flexibly enable oneDNN to dump verbose messages. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think to
should not be added here because the specific environment variable we want to use here is ONEDNN_VERBOSE
, whose value we set to 1
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for your comments. I will keep the original version for this sentence.
recipes_source/amx.rst
Outdated
- BFloat16 data type: | ||
|
||
Using ``torch.cpu.amp`` or ``torch.autocast("cpu")`` would utilize AMX acceleration for supported operators. | ||
|
||
:: | ||
|
||
model = model.to(memory_format=torch.channels_last) | ||
with torch.cpu.amp.autocast(): | ||
output = model(input) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- BFloat16 data type: | |
Using ``torch.cpu.amp`` or ``torch.autocast("cpu")`` would utilize AMX acceleration for supported operators. | |
:: | |
model = model.to(memory_format=torch.channels_last) | |
with torch.cpu.amp.autocast(): | |
output = model(input) | |
- BFloat16 data type: | |
- Using ``torch.cpu.amp`` or ``torch.autocast("cpu")`` would utilize AMX acceleration for supported operators. | |
:: | |
model = model.to(memory_format=torch.channels_last) | |
with torch.cpu.amp.autocast(): | |
output = model(input) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
|
||
- Quantization: | ||
|
||
Applying quantization would utilize AMX acceleration for supported operators. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Applying quantization would utilize AMX acceleration for supported operators. | |
- Applying quantization would utilize AMX acceleration for supported operators. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
recipes_source/amx.rst
Outdated
|
||
- torch.compile: | ||
|
||
When the generated graph model runs into oneDNN implementations with the supported operators, AMX accelerations will be activated. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When the generated graph model runs into oneDNN implementations with the supported operators, AMX accelerations will be activated. | |
- When the generated graph model runs into oneDNN implementations with the supported operators, AMX accelerations will be activated. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed
Fixes #2355
Description
Add a document about how to leverage AMX with PyTorch on the 4th Gen of Xeon.
Checklist
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @ZailiWang @ZhaoqiongZ @leslie-fang-intel @Xia-Weiwen @sekahler2 @zhuhaozhe @Valentine233