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. |
github-actions
bot
added
docathon-h1-2023
✅ Deploy Preview for pytorch-tutorials-preview ready!
To edit notification comments on pull requests, go to your Netlify site settings. |
CaoE
changed the title
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.
| 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. |
recipes_source/amx.rst
Outdated
| ``conv_transpose1d``, | ||
| ``conv_transpose2d``, | ||
| ``conv_transpose3d``, | ||
| ``linear`` |
There was a problem hiding this comment.
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.
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.
Are the speedups only on some particular newer hardware? Is the hardware consumer or enterprise centric?
There was a problem hiding this comment.
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.
it would be nice to have some python function like is_x_available()
There was a problem hiding this comment.
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.
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.
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.
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.
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.
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
recipes_source/amx.rst
Outdated
| ``conv1d``, | ||
| ``conv2d``, | ||
| ``conv3d``, | ||
| ``conv1d``, |
There was a problem hiding this comment.
why we have 2 sets of conv1d, conv2d, conv3d here
recipes_source/amx.rst
Outdated
| ``addbmm``, | ||
| ``linear``, | ||
| ``matmul``, | ||
| ``_convolution`` |
There was a problem hiding this comment.
_convolution is not intended to be directly used, start with a _
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. |
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.
| 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.
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.
| Note: Use channels last format to get better performance. | |
| .. note:: Use channels' last format to get better performance. |
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.
| 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. |
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.
| - 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`` |
recipes_source/amx.rst
Outdated
| ``linear``, | ||
| ``matmul``, | ||
|
|
||
| - Quantization CPU ops that can leverage AMX: |
There was a problem hiding this comment.
| - Quantization CPU ops that can leverage AMX: | |
| Quantization CPU ops that can leverage AMX: |
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.
| 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.
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.
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.
| - 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) |
recipes_source/amx.rst
Outdated
|
|
||
| - Quantization: | ||
|
|
||
| Applying quantization would utilize AMX acceleration for supported operators. |
There was a problem hiding this comment.
| Applying quantization would utilize AMX acceleration for supported operators. | |
| - Applying quantization would utilize AMX acceleration for supported operators. |
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.
| 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. |
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