Segformer model optimization

[mbahnasTT]

• Support Conv large weights on L1
• Support sharded input to Conv, to avoid S2I/I2S
• Conv is running on 8 cores only #14977
• Transpose to support sharded and bfp8 input
• Softmax to support sharded input
• LayerNorm to support height sharded input
• Optimize the Upsample in the Decoder module #14978
• use bfloat8_b where possible
• switch to LoFi fidelity

[ntarafdar]

@mbahnasTT , commenting on behalf of TMs team here. can you comment on the shard spec (height, width or block sharded, shard shape, core grid, shard orientation, and tensor shape) required for Transpose?

Once you give us that we can get it to work in a unit test.

@sjameelTT can look at adding that.

[mbahnasTT]

hi @ntarafdar , @sjameelTT , do you like we have a brief working session? the input height sharded spec is shown in the perf sheet.

#ops_perf_results_SEG_ENCODER_2024_11_01_04_39_36.xls

[sjameelTT]

Hey @mbahnasTT , the bfloat8_b transpose on WH should work soon, it was accidentally disabled in a previous patch. I will push a fix for that. The sharded transpose will need some additional work I believe.

[sjameelTT]

@mbahnasTT let's set up a meeting to discuss the sharded transpose spec?

[mbahnasTT]

@sjameelTT Hey Saad, for the transpose OPs in Segformer Encoder, the sharding spec of the input to transpose is:

• input tensor = [256, 32], TILE, BFLOAT8_B , shard spec = [32, 32] on 8 cores, Height-sharded

[mbahnasTT]

Hey @ntarafdar and @nardoTT
Segformer decoder untilize sharding spec

• input = [16384, 256] --> Height sharding on 32 cores --> spec = [512, 256]
• input = [4096, 256] --> Height sharding on 32 cores --> spec = [128, 256]
• input = [1024, 256] --> Height sharding on 64 cores --> spec = [32, 256]
• input = [256, 256] --> Block sharding on 32 cores --> spec = [32, 32]

Segformer encoder untilize sharding spec

• input = [16384, 32] --> Height sharding on 64 cores --> spec = [256, 32]
• input = [16384, 128] --> Height sharding on 64 cores --> spec = [256, 128]
• input = [4096, 64] --> Height sharding on 64 cores --> spec = [64, 64]
• input = [4096, 256] --> Block sharding on 16 cores --> spec = [512, 128]
• input = [1024, 160] --> Block sharding on 40 cores --> spec = [128, 32]
• input = [1024, 640] --> Block sharding on 40 cores --> spec = [128, 128]
• input = [256, 1024] --> Block sharding on 64 cores --> spec = [32, 128]

[mbahnasTT]

Hey Jay, @jaykru-tt
@ntarafdar asked me to share with you the concat use case in Segformer:

• The concat requested to have a RM/interleaved input
• My input was height shared since it's the output from upsample
• To the concat, there are 4 inputs, coming from 4 different upsample OPs, and different sharding schemes
• So instead of S2I for the 4 inputs, can I keep them RM/height-sharded?
• input = 4x [1,128,128,256] and concat at dim=3
• output = [1,128,128,1024]

[mbahnasTT]

Hey @bbradelTT , we'll need to add some features to improve the perf of Segformer model , please:

• Softmax to support sharded input
• LayerNorm to support height sharded input

I can share more details.

[nkpatel-tt]

@mbahnasTT could you please share details on decoder also?? I am looking at optimising upsample for decoder.

[mbahnasTT]

@nkpatel-tt here are some details:

• pytest tests/ttnn/integration_tests/segformer/test_segformer_decode_head.py
• attached a recent perf sheetops_perf_results_SEG_DECODER_2024_11_03_05_01_37.xls
• upsample code at https://github.com/tenstorrent/tt-metal/blob/main/models/demos/segformer/tt/ttnn_segformer_decode_head.py#L79
• 

[nardoTT]

Hey @ntarafdar and @nardoTT Segformer decoder untilize sharding spec

• input = [16384, 256] --> Height sharding on 32 cores --> spec = [512, 256]
• input = [4096, 256] --> Height sharding on 32 cores --> spec = [128, 256]
• input = [1024, 256] --> Height sharding on 64 cores --> spec = [32, 256]
• input = [256, 256] --> Block sharding on 32 cores --> spec = [32, 32]

Segformer encoder untilize sharding spec

• input = [16384, 32] --> Height sharding on 64 cores --> spec = [256, 32]
• input = [16384, 128] --> Height sharding on 64 cores --> spec = [256, 128]
• input = [4096, 64] --> Height sharding on 64 cores --> spec = [64, 64]
• input = [4096, 256] --> Block sharding on 16 cores --> spec = [512, 128]
• input = [1024, 160] --> Block sharding on 40 cores --> spec = [128, 32]
• input = [1024, 640] --> Block sharding on 40 cores --> spec = [128, 128]
• input = [256, 1024] --> Block sharding on 64 cores --> spec = [32, 128]

The input layout above is supported by the current untilize operation, and utilize has been tested for these inputs. The MLP operation is explicitly converting the tensor layout from sharded to interleaved in line 114 in ttnn_segformer_mlp.py. If this line is commented, the to_layout operation in line 57 in ttnn_segformer_decode_head.py would execute the untilize operation in sharded mode. I don't think the issue you are facing is untilize related.