Describe the bug
We are adding 2 variants of vanilla UNet to run on device via TTNN. Given these variants have trained weights, we can make a demo of our UNet efforts and track PCC robustly per op.
We are currently blocked for this given the issues bellow with the conv ops to run on device. In short, the last conv/classifier won't run on device and I can't find a single sharding strategy (either height or block sharding) for all the convs in the model that would run on device. Please read through for instructions on how to reproduce the issues and more details.
my branch at commit e110ed62fd2ef1cc8953fd1d8c62b600cf53e8be
run my unit test for the convs in the model as pytest tests/ttnn/unit_tests/operations/test_conv2d.py::test_unet_brain_conv
As of now, I have down-sized the unit tests to include math fidelity=LoFi and data-type=BFLOAT8_B only. There are two issues:
1. last convolution/the classifier in UNet fails with block sharding and fails with height sharding as well. I can not find a config override for act_block_h that would pass for the height sharding case2. I can't find configurations where all convs would pass as either block sharding or height sharding end to end. Meaning, we would need to change sharding strategy in between convs in the end to end model which can be expensive
Please see attached the model's graph end to end for your reference. I have commented every test case with the corresponding op from the graph for your convenience as in here
Expected behavior
Enable the last 1x1 conv to run on device.
Find the optimal configurations for block vs height sharding to run the model end to end.
Describe the bug We are adding 2 variants of vanilla UNet to run on device via TTNN. Given these variants have trained weights, we can make a demo of our UNet efforts and track PCC robustly per op. We are currently blocked for this given the issues bellow with the conv ops to run on device. In short, the last conv/classifier won't run on device and I can't find a single sharding strategy (either height or block sharding) for all the convs in the model that would run on device. Please read through for instructions on how to reproduce the issues and more details.
This issue is regarding the BRAIN MRI SEGMENTATION variant
To Reproduce Steps to reproduce the behavior:
pytest tests/ttnn/unit_tests/operations/test_conv2d.py::test_unet_brain_conv
1. last convolution/the classifier in UNet fails with block sharding and fails with height sharding as well. I can not find a config override for act_block_h that would pass for the height sharding case
2. I can't find configurations where all convs would pass as either block sharding or height sharding end to end. Meaning, we would need to change sharding strategy in between convs in the end to end model which can be expensive
Expected behavior Enable the last 1x1 conv to run on device. Find the optimal configurations for block vs height sharding to run the model end to end.
environment information: