Move Transformers

[ayerofieiev-tt]

tt_eager --> ttnn per op plan

We propose next order to breakdown this work into smaller pieces:

1. Replace usage of a given op in C++ with ttnn analog (example)
2. Replace usage in Python [Models, Demos, Examples, tests] (example)
3. Remove tt_lib bindings and entries from docs (bindings, docs)
4. Restructure op consolidating code in ttnn (see new structure)

Replacing usage in C++

Each of such usage should be replaced with a ttnn analog. For example repeat --> ttnn::repeat. This should be done for each operation. Missing operations should be added to ttnn.

• usually there are lots of usage in composite_ops.cpp, complex_ops.cpp, backward_ops.cpp
• lots of usage in tests, example 1, example 2

Replacing usage in Python

For every unary op, look for next entries in Tests/Sweeps, Demos, Models, Examples:

• ttl.tensor.repeat
• tt_lib.tensor.repeat
• ttnn.primary.tensor.repeat and replace them with ttnn.repeat. Example

⚠️ tt_lib operations might sometimes have a slightly different interface

• arguments may be in different order
• arguments may be positional when ttnn arguments may be kw only
• output_mem_config --> memory_config
• optional_output_tensor --> optional_tensor
• output_dtype --> dtype

Testing

For the best coverage, I recommend to run these workflows. If some of them fails, check if it is the same fail as on main:

• Post commit CI
• [post-commit] models tests
• Device perf regressions and output report
• Model perf regressions and output report
• Nightly fast dispatch model regression CI
• T3000 frequent tests
• T3000 demo tests

Scope of Work for Transformers

There are two sets of files to move which contain two sets of ops: any variation of split qkv heads and concat heads.

• [x] Split heads and concat heads from: tt_eager/tt_dnn/op_library/transformer_tms/transformer_tms.hpp
  • [x] operations.primary.transformers.split_query_key_value_and_split_heads used in:
    • models/experimental/bert_large_performant/unit_tests/test_bert_large_split_query_key_value_and_split_heads_sharded.py
    • models/experimental/bert_large_performant/unit_tests/test_bert_large_split_and_transform_qkv_heads.py
    • models/demos/metal_BERT_large_11/tt/mha.py
    • _tests/tt_eager/integration_tests/testbert.cpp (used as tt::operations::primary::transformers::split_query_key_value_and_split_heads)
  • [x] operations.primary.transformers.concatenate_heads used in:
    • models/experimental/bert_large_performant/unit_tests/test_bert_large_concatenate_heads.py
    • _tests/tt_eager/integration_tests/testbert.cpp (used as tt::operations::primary::transformers::concatenate_heads)
• Entirety of: tt_eager/tt_dnn/op_library/nlp_tms/nlp_tms.hpp
  • [ ] tensor.create_qkv_heads used in:
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_create_qkv_heads.py
  • ttnn/cpp/ttnn/operations/transformer.hpp (used as tt_metal::create_qkv_heads)
  • [ ] tensor.create_qkv_heads_from_separate_tensors used in:
  • models/demos/wormhole/stable_diffusion/tt2/ttnn_functional_cross_attention.py
  • ttnn/ttnn/experimental_operations/golden_functions.py
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_create_qkv_heads.py
  • [ ] tensor.nlp_create_qkv_heads_falcon7b used in:
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_create_qkv_heads.py
  • models/demos/falcon7b/tt/falcon_attention.py
  • ttnn/cpp/ttnn/operations/transformer.hpp (used as tt_metal::nlp_create_qkv_heads_falcon_7b)
  • [ ] tensor.nlp_create_qkv_heads_decode used in:
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_create_qkv_heads_decode.py
  • models/demos/t3000/mixtral8x7b/tt/mixtral_attention.py
  • models/experimental/llama2_70b/tests/unit_tests/test_nlp_create_head.py
  • models/demos/t3000/llama2_70b/tests/unit_tests/test_nlp_create_head.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • tests/ttnn/multichip_unit_tests/test_multidevice_TG.py
  • [ ] tensor.nlp_create_qkv_heads used in:
  • models/demos/t3000/falcon40b/tt/ops/falcon_nlp_create_qkv_heads.py
  • models/demos/t3000/falcon40b/tt/falcon_attention.py
  • models/demos/t3000/falcon40b/tests/test_falcon_model_single_chip.py
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_create_qkv_heads.py
  • models/demos/wormhole/mistral7b/tt/mistral_attention.py
  • models/experimental/llama2_70b/tests/unit_tests/test_nlp_create_head.py
  • models/demos/t3000/llama2_70b/tests/unit_tests/test_nlp_create_head.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • ttnn/cpp/ttnn/operations/transformer.hpp (used as tt_metal::nlp_create_qkv_heads)
  • [ ] tensor.nlp_concat_heads used in:
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_concat_heads.py
  • models/experimental/llama2_70b/tests/unit_tests/test_nlp_concat_head.py
  • models/demos/t3000/llama2_70b/tests/unit_tests/test_nlp_concat_head.py
  • models/experimental/llama2_70b/tests/unit_tests/test_attn_sdpa.py
  • models/demos/t3000/llama2_70b/tests/unit_tests/test_attn_sdpa.py
  • models/demos/t3000/falcon40b/tt/falcon_attention.py
  • models/demos/t3000/falcon40b/tt/falcon_attention.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/falcon7b/tt/falcon_attention.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_sharded.py
  • models/demos/falcon7b/tt/falcon_attention.py
  • models/demos/metal_BERT_large_11/tt/mha.py
  • [ ] tensor.nlp_concat_heads_decode used in:
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_concat_heads_decode.py
  • models/experimental/llama2_70b/tests/unit_tests/test_nlp_concat_head.py
  • models/demos/t3000/llama2_70b/tests/unit_tests/test_nlp_concat_head.py
  • models/demos/t3000/mixtral8x7b/tt/mixtral_attention.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • tests/ttnn/multichip_unit_tests/test_multidevice_TG.py
  • [ ] tensor.nlp_kv_cache_load_slice used in (this is a new op (?) but I will move it for completeness):
  • tests/tt_eager/python_api_testing/unit_testing/misc/test_nlp_kv_cache_load_slice.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • models/experimental/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/llama2_70b/tt/llama_attention_optimized.py
  • models/demos/t3000/mixtral8x7b/tt/mixtral_attention.py

[eyonland]

@TT-BrianLiu , is this work you could pick up here?

[TT-BrianLiu]

Yea, I can take a look

[TT-BrianLiu]

I have scoped out the work required for moving these ops (see the updated description). I will likely start with tt_eager/tt_dnn/op_library/transformer_tms/transformer_tms.hpp, then take it op by op for the ones in tt_eager/tt_dnn/op_library/nlp_tms/nlp_tms.hpp. At the end, tt_eager/tt_dnn/op_library/transformer_tms/transformer_tms.hpp will still exist since it contains other custom models ops that are completely unrelated but tt_eager/tt_dnn/op_library/transformer_tms/transformer_tms.hpp should be completely gone.

I plan to move everything into ttnn::experimental::transformer