Summary:
add a unit test to test the situation where a node is too large to fit into any device
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48938
Reviewed By: zhangguanheng66
Differential Revision: D25402967
Pulled By: scottxu0730
fbshipit-source-id: a2e2a3dc70d139fa678865ef03e67fa57eff4a1d
Summary:
add a unit test for the situation where devices have no enough memory
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48858
Reviewed By: malfet, gcatron
Differential Revision: D25341254
Pulled By: scottxu0730
fbshipit-source-id: c0524c22717b6c8afd67f5b0ad0f1851b973e4b7
Summary:
What this does is that given a `FxModule foo`, you can call `foo.to_folder('foo_folder', 'Foo')` and dump the current FX module into runnable Python code.
That is
```
foo = <fxModule>
foo = foo.to_folder('bar', 'Foo')
from bar import Foo
foo2 = Foo()
forall x, foo2(x) == Foo(x)
```
This has several use cases, largely lifted from jamesr66a's doc here: https://fb.quip.com/U6KHAFaP2cWa (FB-internal).
1. As we apply more heavy-weight function transformations with FX, figuring out what's going on can be quite a difficult experience. In particular, things that can typically be used for debugging (like `print` or `import pdb; pdb.set_trace()`) no longer work. This is particularly necessary if you're using a FX transform like `grad` or `vmap. With this, you simply open up the dumped file, and add `print`/`pdb` statements wherever you'd like.
2. This also provides an immense amount of user control. Some potential use-cases:
- Let's say an existing FX transform has some bug, or generates suboptimal code. Instead of needing to modify that FX transform, writing another FX pass that fixes the suboptimal code, or simply giving up on FX, they can workaround it by simply modifying the resulting code themselves.
- This allows users to check in their FX modules into source control.
- You could even imagine using this as part of some code-gen type workflow, where you write a function, `vmap` it to get the function you actually want, and then simply copy the output of the `vmap` function without needing FX at all in the final code.
An example:
```python
class Test(nn.Module):
def __init__(self):
super(Test, self).__init__()
self.W = torch.nn.Parameter(torch.randn(2))
self.linear = nn.Linear(2, 2)
self.attr = torch.randn(2)
self.attr2 = torch.randn(2)
def forward(self, x):
return self.linear(self.W + (self.attr + self.attr2) + x)
mod = fx.symbolic_trace(Test())
mod.to_folder('foo', 'Foo')
```
results in
```python
import torch
class Foo(torch.nn.Module):
def __init__(self):
super().__init__()
state_dict = torch.load('foo/state_dict.pt')
self.linear = torch.load('foo/linear.pt') # Linear(in_features=2, out_features=2, bias=True)
self.__tensor_constant0 = state_dict['__tensor_constant0']
self.W = torch.nn.Parameter(state_dict['W'])
def forward(self, x):
w = self.W
tensor_constant0 = self.__tensor_constant0
add_1 = w + tensor_constant0
add_2 = add_1 + x
linear_1 = self.linear(add_2)
return linear_1
```
Some current issues:
1. How do you actually ... save things like modules or parameters? I don't think FX is in the business of tracking initializations and such. Thus, the only way I see to do it is to dump the parameters/modules as blobs, and then load them in the generated initialization. This is a somewhat subpar user experience, and perhaps prevents it from being in some use cases (ie: you would need to check in the blobs into source control to save the model).
2. Currently, the only "atomic" modules we have are those in `torch.nn`. However, if we want to allow flexibility in this, and for example, allow "atomic" modules that are user-defined, then it's not clear how to allow those to be dumped in a way that we can then load elsewhere.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47544
Reviewed By: jamesr66a
Differential Revision: D25232917
Pulled By: Chillee
fbshipit-source-id: fd2b61a5f40e614fc94256a2957ed1d57fcf5492
Summary:
This PR add supports on AOT based partition. Given each node and its corresponding partition id, generate the partition, submodules and dag
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48336
Reviewed By: gcatron
Differential Revision: D25226899
Pulled By: scottxu0730
fbshipit-source-id: 8afab234afae67c6fd48e958a42b614f730a61d9
Summary:
Some interesting stuff going on. All benchmarks are tested with both my implementation as well as the current quantized fuser.
For these benchmarks, things like using MKLDNN/FBGEMM make a big differene.
## Manual compilation (everything turned off)
In the small case, things look good
```
non-fused: 1.174886703491211
fused: 0.7494957447052002
```
However, for `torchvision.resnet18`, we see
```
non-fused: 1.2272708415985107
fused: 3.7183213233947754
```
This is because Conv (no bias) -> Batch Norm is actually faster than Conv (bias) if you don't have any libraries...
## Nightly (CPU)
```
Toy
non-fused: 0.45807552337646484
fused: 0.34779977798461914
resnet18
non-fused: 0.14216232299804688
fused: 0.13438796997070312
resnet50
non-fused: 0.2999534606933594
fused: 0.29364800453186035
densenet161
non-fused: 0.6558926105499268
fused: 0.6190280914306641
inception_v3
non-fused: 1.2804391384124756
fused: 1.181272029876709
```
with MKLDNN.
We see a small performance gain across the board, with more significant performance gains for smaller models.
## Nightly (CUDA)
```
M
non-fused: 1.2220964431762695
fused: 1.0833759307861328
resnet18
non-fused: 0.09721899032592773
fused: 0.09089207649230957
resnet50
non-fused: 0.2053072452545166
fused: 0.19138741493225098
densenet161
non-fused: 0.6830024719238281
fused: 0.660109281539917
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47657
Reviewed By: eellison
Differential Revision: D25127546
Pulled By: Chillee
fbshipit-source-id: ecdf682038def046045fcc09faf9aeb6c459b5e3
Summary:
This is a partition search based on Kernighan-Lin algorithm. First, the graph is partitioned using size_based_partition, then nodes from different partitions are swapped until the cost reaches minimum.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48197
Reviewed By: gcatron
Differential Revision: D25097065
Pulled By: scottxu0730
fbshipit-source-id: 3a11286bf4e5a712ab2848b92d0b98cd3d6a89be
Summary:
This PR fixes the add_node and remove_node in partition class and also add a unit test for node manipulation in partition
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48016
Reviewed By: gcatron
Differential Revision: D24996368
Pulled By: scottxu0730
fbshipit-source-id: 0ddffd5ed3f95e5285fffcaee8c4b671929b4df3
Summary:
Change Partitioner.py file name to partitioner.py
Change GraphManipulation.py file name to graph_manipulation.py
Move test_replace_target_nodes_with() to test_fx_experimental.py
Remove the unnecessary argument in size_based_partition() in Partitioner class
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47914
Reviewed By: gcatron
Differential Revision: D24956653
Pulled By: scottxu0730
fbshipit-source-id: 25b65be7dc7d64e90ffdc59cf394446fee83c3e6
Summary:
In the cost_aware_partition, check the circular dependency in try_combining_partitions. Also fix the calculate of communication time between partitions.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47856
Reviewed By: gcatron
Differential Revision: D24926591
Pulled By: scottxu0730
fbshipit-source-id: c634608675ac14b13b2370a727e4fb05e1bb94f0
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47763
Changing the name due to the discussion in
https://github.com/pytorch/pytorch/pull/47399.
Test Plan:
```
python test/test_utils.py TestAssert.test_assert_true
python test/test_fx.py TestFX.test_symbolic_trace_assert
python test/test_fx_experimental.py
```
Imported from OSS
Reviewed By: ezyang
Differential Revision: D24891767
fbshipit-source-id: 01c7a5acd83bf9c962751552780930c242134dd2
Summary:
[WIP]This PR adds cost_aware_partition method in Partitioner class. The method partitions the fx graph module based on the latency of the whole graph.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47673
Reviewed By: gcatron
Differential Revision: D24896685
Pulled By: scottxu0730
fbshipit-source-id: 1b1651fe82ce56554f99d68da116e585c74099ed
Summary:
add get_device_to_partitions_mapping function in the Partitioner class to make size_based_partition more modular and organized. This function will also be used in the future cost_aware_partition
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47361
Reviewed By: gcatron
Differential Revision: D24760911
Pulled By: scottxu0730
fbshipit-source-id: 8cdda51b9a1145f9d13ebabbb98b4d9df5ebb6cd
Summary:
This PR adds the support to calculate the cost of a partitioned graph partition by partition based on the node cost. In a partitioned graph, top partitions (partitions without parents) are collected as the starting points, then use DFS to find the critical path among all partitions in the graph
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47280
Reviewed By: gcatron
Differential Revision: D24735932
Pulled By: scottxu0730
fbshipit-source-id: 96653a8208554d2c3624e6c8718628f7c13e320b
Summary:
This PR adds node-by-node cost function. Given a partition of nodes, get_latency_of_one_partition function will find the critical path in the partition and return its latency. A test unit is also provided. In the test unit, a graph module is partitioned into two partitions and the latency of each partition is tested.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47009
Reviewed By: gcatron
Differential Revision: D24692542
Pulled By: scottxu0730
fbshipit-source-id: 64c20954d842507be0d1afa2516d88f705e11224
Summary:
Placeholders and constants in the partition are counted twice when combining two partitions. This PR fixes it by adding add_node function into Partition class. A unit test is also updated to test if the partition size is correct
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47083
Reviewed By: gcatron
Differential Revision: D24634368
Pulled By: scottxu0730
fbshipit-source-id: ab408f29da4fbf729fd9741dcb3bdb3076dc30c4
Summary:
WIP: add support for different memory sizes on size_based_partition, so the size_based_partition could support different logical devices with different memory sizes. Compared to the original size_based_partition, the new one also supports partition to logical device mapping. Multiple partitions can be mapped into one device if the memory size is allowed. A test unit test_different_size_partition is also added.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46919
Reviewed By: gcatron, VitalyFedyunin
Differential Revision: D24603511
Pulled By: scottxu0730
fbshipit-source-id: 1ba37338ae054ad846b425fbb7e631d3b6c500b6
Summary:
WIP: This PR adds sparse_nn_partition into Partitioner class. It includes logical device assignment for all dag nodes. The basic idea is to do size_based_partition separately for embedding nodes and non-embedding nodes. A test unit is also added in test_fx_experimental.py.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46390
Reviewed By: gcatron
Differential Revision: D24555415
Pulled By: scottxu0730
fbshipit-source-id: 8772af946d5226883759a02a1c827cfdfce66097
Summary:
Reopen the PR: https://github.com/pytorch/pytorch/pull/45837
This PR add a new feature for Partitioner() class called size_based_partition. Given a list of devices with the same memory size, this function could distribute graph nodes into different devices. To implement this feature, several help functions are created in Partitioner.py and GraphManipulation.py.
An unit test is also added in test/test_fx_experimental.py
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46282
Reviewed By: gcatron
Differential Revision: D24288470
Pulled By: scottxu0730
fbshipit-source-id: e81b1e0c56e34f61e497d868882126216eba7538
