Skip to content

Commit 5a73f9e

Browse files
authored
Publish dynamic shape doc. (#5631)
* add bounded dynamic shape doc * got more done * Publish dynamic shape doc. * remove unwanted file * fix typo * add rfc to the doc * fix a typo
1 parent 987873c commit 5a73f9e

File tree

2 files changed

+47
-0
lines changed

2 files changed

+47
-0
lines changed
34.2 KB
Loading

docs/dynamic_shape.md

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Dynamic shape
2+
3+
Dynamic shape refers to the variable nature of a tensor shape where its shape depends on the value of another upstream tensor. For example:
4+
```
5+
>>> import torch, torch_xla
6+
>>> in_tensor = torch.randint(low=0, high=2, size=(5,5), device='xla:0')
7+
>>> out_tensor = torch.nonzero(in_tensor)
8+
```
9+
the shape of `out_tensor` depends on the value of `in_tensor` and is bounded by the shape of `in_tensor`. In other words, if you do
10+
```
11+
>>> print(out_tensor.shape)
12+
torch.Size([<=25, 2])
13+
```
14+
you can see the first dimension depends on the value of `in_tensor` and its maximum value is 25. We call the first dimension the dynamic dimension. The second dimension does not depend on any upstream tensors so we call it the static dimension.
15+
16+
Dynamic shape can be further categorized into bounded dynamic shape and unbounded dynamic shape.
17+
- bounded dynamic shape: refers to a shape whose dynamic dimensions are bounded by static values. It works for accelerators that require static memory allocation (e.g. TPU).
18+
- unbounded dynamic shape: refers to a shape whose dynamic dimensions can be infinitely large. It works for accelerators that don’t require static memory allocation (e.g. GPU).
19+
20+
Today, only the bounded dynamic shape is supported and it is in the experimental phase.
21+
22+
## Bounded dynamic shape
23+
24+
Currently, we support multi-layer perceptron models (MLP) with dynamic size input on TPU.
25+
26+
This feature is controlled by a flag `XLA_EXPERIMENTAL="nonzero:masked_select"`. To run a model with the feature enabled, you can do:
27+
```
28+
XLA_EXPERIMENTAL="nonzero:masked_select:masked_scatter" python your_scripts.py
29+
```
30+
31+
Here are some numbers we get when we run the MLP model for 100 iterations:
32+
33+
| | No dynamic shape | With dynamic shape |
34+
| :--- | :----: | ---: |
35+
| End-to-end training time | 29.49 | 20.03 |
36+
| Number of compilations | 102 | 49 |
37+
| Compilation cache hit | 198 | 1953 |
38+
39+
![Performance comparison (a) without dynamic shape (b) with dynamic shape](assets/dynamic_shape_mlp_perf.png)
40+
41+
One of the motivations of the dynamic shape is to reduce the number of excessive recompilation when the shape keeps changing between iterations. From the figure above, you can see the number of compilations reduced by half which results in the drop of the training time.
42+
43+
To try it out, run
44+
```
45+
XLA_EXPERIMENTAL="nonzero:masked_select" PJRT_DEVICE=TPU python3 pytorch/xla/test/ds/test_dynamic_shape_models.py TestDynamicShapeModels.test_backward_pass_with_dynamic_input
46+
```
47+
For more details on how we plan to expand the dynamic shape support on PyTorch/XLA in the future, feel free to review our [RFC](https://github.com/pytorch/xla/issues/3884).

0 commit comments

Comments
 (0)