Hugging Face's logo

raahemnabeel
/
qwen3-coder-next

blackhole
p150x2
p150x4
p300
p300x2
tt-model-cache
tt-model-catalog
tt-model-container
vllm-plugin
Model card Files
xet
 Community
qwen3-coder-next / code /models /tt_transformers /README.md
raahemnabeel's picture
raahemnabeel
Add files using upload-large-folder tool
d18366a verified
preview code
|
Raw History Blame Contribute Delete
36.1 kB

TT-Transformers

This code can run large language models such as the Llama3 family, Qwen2.5, Mistral, DeepSeek-R1-Distill variants and similar. Tensor-parallelism automatically distributes workloads across all available chips.

The current version is verified to work with the following models:

Model Hardware <org/model>
DeepSeek R1 Distill Llama 70B LoudBox / QuietBox / Galaxy deepseek-ai/DeepSeek-R1-Distill-Llama-70B
EXAONE 4.5 33B (text; host-vision hybrid demo) P150x8 (BH LoudBox) LGAI-EXAONE/EXAONE-4.5-33B
Llama 3.1 8B n150 / p100 / p150 meta-llama/Llama-3.1-8B
Llama 3.1 70B LoudBox / QuietBox / Galaxy meta-llama/Llama-3.1-70B
Llama 3.2 1B n150 meta-llama/Llama-3.2-1B
Llama 3.2 3B n150 meta-llama/Llama-3.2-3B
Llama 3.2 11B Vision n300 meta-llama/Llama-3.2-11B-Vision
Llama 3.2 90B Vision LoudBox / QuietBox meta-llama/Llama-3.2-90B-Vision
Mistral 7B Instruct v0.3 n150 mistralai/Mistral-7B-Instruct-v0.3
Mistral Small 3.1 24B Instruct T3K mistralai/Mistral-Small-3.1-24B-Instruct-2503
Mixtral 8x7B Instruct v0.1 LoudBox / QuietBox mistralai/Mixtral-8x7B-Instruct-v0.1
Qwen 2.5 7B n300 Qwen/Qwen2.5-7B
Qwen 2.5 Coder 32B LoudBox / QuietBox Qwen/Qwen2.5-Coder-32B
Qwen 2.5 72B LoudBox / QuietBox Qwen/Qwen2.5-72B
Qwen 3 32B LoudBox / QuietBox Qwen/Qwen3-32B
Phi-3-mini-128k-instruct n150 / n300 microsoft/Phi-3-mini-128k-instruct
Explore an extended list of compatible models (experimental)
Model Hardware <org/model>.
Bellatrix-Tiny-0.5B N150 prithivMLmods/Bellatrix-Tiny-0.5B
Bellatrix-Tiny-1B N150 prithivMLmods/Bellatrix-Tiny-1B
Bellatrix-Tiny-1B-R1 N150 prithivMLmods/Bellatrix-Tiny-1B-R1
deepseek-coder-6.7b-instruct N300 deepseek-ai/deepseek-coder-6.7b-instruct
deepseek-llm-7b-base N300 deepseek-ai/deepseek-llm-7b-base
deepseek-llm-7b-chat N300 deepseek-ai/deepseek-llm-7b-chat
DeepSeek-R1-Distill-Llama-8B N300 deepseek-ai/DeepSeek-R1-Distill-Llama-8B
DeepSeek-R1-Distill-Llama-8B N300 unsloth/DeepSeek-R1-Distill-Llama-8B
DeepSeek-R1-Distill-Llama-8B-abliterated N300 huihui-ai/DeepSeek-R1-Distill-Llama-8B-abliterated
DeepSeek-R1-Distill-Qwen-1.5B N150 deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
DeepSeek-R1-Distill-Qwen-32B T3K deepseek-ai/DeepSeek-R1-Distill-Qwen-32B
DeepSeek-R1-Distill-Qwen-32B-Japanese T3K cyberagent/DeepSeek-R1-Distill-Qwen-32B-Japanese
DeepSeek-R1-Distill-Qwen-32B-abliterated T3K huihui-ai/DeepSeek-R1-Distill-Qwen-32B-abliterated
DeepSeek-R1-Distill-Qwen-7B N300 deepseek-ai/DeepSeek-R1-Distill-Qwen-7B
DeepSeek-R1-Distill-Qwen-7B-abliterated-v2 N300 huihui-ai/DeepSeek-R1-Distill-Qwen-7B-abliterated-v2
DeepSeek-R1-Distill-Qwen-7B-Japanese N300 lightblue/DeepSeek-R1-Distill-Qwen-7B-Japanese
Dobby-Mini-Leashed-Llama-3.1-8B N300 SentientAGI/Dobby-Mini-Leashed-Llama-3.1-8B
Dobby-Mini-Unhinged-Llama-3.1-8B N300 SentientAGI/Dobby-Mini-Unhinged-Llama-3.1-8B
Dolphin3.0-Mistral-24B T3K cognitivecomputations/Dolphin3.0-Mistral-24B
Dolphin3.0-R1-Mistral-24B T3K cognitivecomputations/Dolphin3.0-R1-Mistral-24B
EuroLLM-9B-Instruct N300 utter-project/EuroLLM-9B-Instruct
GPT4chan-24B T3K v2ray/GPT4chan-24B
Llama-2-7b-chat-hf N300 meta-llama/Llama-2-7b-hf
Llama-2-7b-hf N300 meta-llama/Llama-2-7b-hf
Llama-3.1-8B-Instruct N300 meta-llama/Llama-3.1-8B-Instruct
Llama-3.2-1B-Instruct N150 meta-llama/Llama-3.2-1B-Instruct
Llama-3.2-3B-Instruct N150 meta-llama/Llama-3.2-3B-Instruct
Llama-3.2-6B-AlgoCode N150 prithivMLmods/Llama-3.2-6B-AlgoCode
Llama-Krikri-8B-Base N150 ilsp/Llama-Krikri-8B-Base
Llama-Krikri-8B-Instruct N300 ilsp/Llama-Krikri-8B-Instruct
LwQ-10B-Instruct N300 prithivMLmods/LwQ-10B-Instruct
Magnum-v4-Cydonia-vXXX-22B T3K Kaoeiri/Magnum-v4-Cydonia-vXXX-22B
Megatron-Opus-7B-Exp N300 prithivMLmods/Megatron-Opus-7B-Exp
Meta-Llama-3-8B N300 meta-llama/Meta-Llama-3-8B
Meta-Llama-3-8B-Instruct N300 meta-llama/Meta-Llama-3-8B-Instruct
Mistral-7B-Instruct-v0.1 N150 mistralai/Mistral-7B-Instruct-v0.1
Mistral-7B-Instruct-v0.2 N150 mistralai/Mistral-7B-Instruct-v0.2
Mistral-7B-v0.1 N150 mistralai/Mistral-7B-v0.1
Mistral-Small-24B-Base-2501 T3K mistralai/Mistral-Small-24B-Base-2501
Mistral-Small-24B-Instruct-2501 T3K mistralai/Mistral-Small-24B-Instruct-2501
Mistral-Small-3.1-24B-Instruct-2503 T3K mistralai/Mistral-Small-3.1-24B-Instruct-2503
NuminaMath-7B-TIR N300 AI-MO/NuminaMath-7B-TIR
OpenThinker-7B N300 open-thoughts/OpenThinker-7B
Primal-Mini-3B-Exp N150 prithivMLmods/Primal-Mini-3B-Exp
Qwen2.5-0.5B N150 Qwen/Qwen2.5-0.5B
Qwen2.5-0.5B-Instruct N150 Qwen/Qwen2.5-0.5B-Instruct
Qwen2.5-1.5B N150 Qwen/Qwen2.5-1.5B
Qwen2.5-1.5B-Instruct N150 Qwen/Qwen2.5-1.5B-Instruct
Qwen2.5-3B-Instruct N150 Qwen/Qwen2.5-3B-Instruct
Qwen2.5-32B T3K Qwen/Qwen2.5-32B
Qwen2.5-Coder-7B-Instruct N300 Qwen/Qwen2.5-Coder-7B-Instruct
Qwen2.5-Math-1.5B N150 Qwen/Qwen2.5-Math-1.5B
Qwen2.5-Math-7B N300 Qwen/Qwen2.5-Math-7B
Reasoning-Distilled-ta-7B N300 prithivMLmods/Reasoning-Distilled-ta-7B
s1-32B T3K simplescaling/s1-32B
Selene-1-Mini-Llama-3.1-8B N300 AtlaAI/Selene-1-Mini-Llama-3.1-8B
Sky-T1-32B-Preview T3K NovaSky-AI/Sky-T1-32B-Preview
SmallThinker-3B-Preview N150 PowerInfer/SmallThinker-3B-Preview
SmolLM2-1.7B-Instruct N150 HuggingFaceTB/SmolLM2-1.7B-Instruct
Sqweeks-7B-Instruct N300 prithivMLmods/Sqweeks-7B-Instruct
TinyLlama-1.1B-Chat-v1.0 N150 TinyLlama/TinyLlama-1.1B-Chat-v1.0
TinySwallow-1.5B N150 SakanaAI/TinySwallow-1.5B
TinySwallow-1.5B-Instruct N150 SakanaAI/TinySwallow-1.5B-Instruct
Triangulum-v2-10B N300 prithivMLmods/Triangulum-v2-10B
WebMind-7B-v0.1 N300 prithivMLmods/WebMind-7B-v0.1
WizardLM-7B-Uncensored N300 cognitivecomputations/WizardLM-7B-Uncensored
Xwen-7B-Chat N300 xwen-team/Xwen-7B-Chat
YuE-s1-7B-anneal-en-cot N300 m-a-p/YuE-s1-7B-anneal-en-cot
YuE-s2-1B-general N150 m-a-p/YuE-s2-1B-general

Prerequisites

Install TT-Metalium and TTNN.

How to Run

To configure the weights and run a demo, choose Automatic Download for a quick setup or Manual Download for more control over the files.

Automatic Setup

Login to HuggingFace using your token

Use huggingface-cli login or set the token with the command export HF_TOKEN=<token>

Choose your Model

Set the HF_MODEL environment variable to the HuggingFace org/name of the model to be run.

This will automatically download the weights into your HuggingFace cache directory and run the model directly.

  • Check the models chart on the top of the page and substitute the <org/model> on the following command:
export HF_MODEL=<org/name>

You are all set!

Manual Setup

If you wish, you can manually download the weights from HuggingFace as described by the following section:

Download Weights from HuggingFace

Navigate to HuggingFace and download the weights of the model you wish to run (check models chart for compatibility and direct links to HuggingFace models).

  • Ensure your model directory has the following structure:
/path_to/org/model/
    config.json
    generation_config.json
    model-00001-of-00062.safetensors
    ...
  • Set HF_MODEL to the directory of the downloaded weights:
export HF_MODEL=<path_to_downloaded_directory>

You are all set!

Run the Demo

  • Make sure HF_MODEL has been set:

export HF_MODEL=<hf_model_name or hf_downloaded_directory>

  • Run the Demo

These are some example commands on how to run the demo. For more information about the parameters, check parameters description.

# Batch-1
pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1"

# Batch-32
pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-32"

# Long context with custom parameters
pytest models/tt_transformers/demo/simple_text_demo.py -k "long-context" --max_seq_len=16384

# Long-context
pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and long"

The above examples are run in ModelOptimizations.performance mode. You can override this by setting the optimizations or the decoder_config_file argument in the demo. To use instead the accuracy mode you can call the above tests with -k "accuracy and ..." instead of performance.

NOTE: trace region sizes are declared in models/model_trace_region_sizes.yaml and resolved at device-open time via get_supported_trace_region_size (which delegates to resolve_trace_region_size). A (model, SKU) pair without a YAML entry is not an error: resolution logs an info message and falls back to TRACE_REGION_SIZE_DYNAMIC (0, dynamic allocation). Add an explicit entry when a model needs a fixed reserved trace region.

Details

Extra compatibility settings for non-Llama models

If you are bringing up a new model that is similar to these but is not listed above, you may also need to set additional environment variables:

  • MAX_PREFILL_CHUNK_SIZE - this determines how many thousands of tokens are prefilled in one go. For optimal performance pick 128. Depending on the model dimensions and hardware you're running on, there may not be enough L1 to prefill 128K tokens at once, in which case you can reduce this in powers of 2 down to 4.
  • PAD_MLP_CORES - models with a hidden_dim that is not a nice power of 2 may not have a valid layout or may run with lower performance. You can set this to a multiple of 8 between 8 and 64; 16 and 32 commonly work well if this is required.

You should also watch out for:

  • RoPE encoding style. llama3, yarn and of course none are supported. HuggingFace models encode the complex numbers in RoPE as r1, r2, ..., i1, i2, ... whereas Meta models encode the complex numbers as r1, i1, r2, i2, ... - TTT uses a Meta-style implementation of the RoPE op and when loading a HuggingFace it will reshuffle the weights of the pre-rope attention weights to interleave their outputs in this style (see reverse_permute in load_checkpoints.py). When using TTT this happens invisibly and should not affect you, but when bringing up a new model or modifying TTT being aware of this will make some otherwise confusing things clearer.
  • Our [accuracy test](demo/simple_text_demo.py -k ci-token-matching) will require you to generate some reference logits and perhaps update the test to use them.
  • We parallelise attention over the number of heads. If this number is e.g. 14 then you will not be able to run it on more than 2 chips (because 14/2=7, a prime number). We do not support head-padding or similar mitigations at this time but a PR would be cool.

Huggingface models specify their architecture in the config.json file. The following architectures are known to work:

  • LlamaForCausalLM
  • Qwen2ForCausalLM
  • Qwen3ForCausalLM
  • MistralForCausalLM
  • Mistral3ForConditionalGeneration
  • Phi3ForCausalLM
  • Exaone4_5_ForConditionalGeneration (text decoder only)

At the time of writing this covers the majority of popular HuggingFace text-generation models. If you find another architecture that works or extend TT-Transformers to support one we would love to accept a PR!

Environment Variables Description

  • HF_MODEL is the HuggingFace org/name of the model you want to run or the path to the downloaded Huggingface weights.
  • TT_CACHE_PATH is optional. It sets the path for ttnn's weight cache files. See below for more details.
  • MESH_DEVICE is optional. It allows you to use fewer devices than are available. See below for more details.

On the first execution of each model, TTNN will create weight cache files for that model, to speed up future runs. These cache files only need to be created once for each model and device. These files are stored in one of three places:

  1. TT_CACHE_PATH if you have set it.
  2. HF_MODEL/device_name if a path to downloaded weights was specified using HF_MODEL.
  3. model_cache/HF_MODEL/device_name if a HuggingFace model name was specified using HF_MODEL.

The device name used is:

  • N150 for N150
  • N300 for N300
  • T3K for LoudBox / QuietBox
  • TG for Galaxy
  • P150x8 for BH LoudBox

By default tensor parallelism is used to run the model over all available chips. You can instead run on a smaller mesh either for testing or for performance reasons (for very small models the communication overhead of tensor parallelism may be larger than the performance gained). To use a smaller mesh, set MESH_DEVICE to one of the supported devices: N150, N300, T3K, TG or P150x8.

Example: export MESH_DEVICE=N150, will enable running one a single chip of a multi-chip system.

Parameters Description

The simple_text_demo.py script includes the following main modes of operation and is parametrized to support other configurations.

  • batch-1: Runs a small prompt (128 tokens) for a single user
  • batch-32: Runs a small prompt (128 tokens) for a a batch of 32 users
  • long-context: Runs a large prompt (64k tokens) for a single user
  • reasoning-1: Runs a reasoning prompt for a single user (generates up to 15k tokens)

If you want to provide your own demo configuration, please take a look at the pytest parametrize calls in models/tt_transformers/demo/simple_text_demo.py. For convenience we list all the supported params below:

  • input_prompts (string): input json file with prompts to process. See models/tt_transformers/demo/*.json for a list of input files
  • instruct (bool): Whether to use Llama instruct weights or general weights
  • repeat_batches (int): Number of consecutive batches of users to run (default: 1)
  • max_seq_len (int): Maximum context length supported by the model (refer to the table above)
  • batch_size (int): Number of users in a batch (Supports 1/2/4/8/16/32 batches)
  • max_generated_tokens (int): Maximum number of tokens to generate for each user (Note that the users will stop generation before this limit if they reach a eos token)
  • paged_attention (bool): Whether to use paged attention or default attention (vLLM support (WIP) requires paged attention)
  • page_params (dict): Page parameters for paged attention - [block_size, max_num_blocks]. For smaller context lengths use block_size=32 and max_num_blocks=1024, for larger context use block_size=64 and max_num_blocks=2048
  • sampling_params (dict): Sampling parameters for decoding -[temperature, top_p]. If temperature is set to 0, argmax (greedy decode) is used.
  • stop_at_eos (bool): Flag to stop decoding when the model generates an EoS token
  • optimizations (ModelOptimizations): Optimization level to use for the model [accuracy, performance]. Applied uniformly across all decoders unless an override config exists in models/tt_transformers/model_params/<model-name>
  • decoder_config_file (DecodersPrecision): Fine-grained optimization control that allows specifying a configuration file to set different settings for each decoder.

Please note that using argmax with batch_size > 1 or using top-p sampling with any batch size, these ops will be run on host. This is because those ops are not yet fully supported on device. A decrease in performance is expected when these configurations are enabled.

Optimization overrides

Some models require a unique set of optimizations defined in models/tt_transformers/model_params/<model-name>. To override the default optimizations, you can define files named models/tt_transformers/tt/model_config/PERFORMANCE_DECODER_CONFIG_FILENAME and models/tt_transformers/tt/model_config/ACCURACY_DECODER_CONFIG_FILENAME in the appropriate models/tt_transformers/model_params/<model-name> directory to override the ModelOptimizations.performance and ModelOptimizations.accuracy optimizations respectively. For example, to override the default "performance" optimizations for Llama3.1-8B-Instruct, a file named performance_decoder_config.json has been created in the models/tt_transformers/model_params/Llama3.1-8B-Instruct directory. The content to write in override files is described in the custom optimizations section. Optimizations are applied with the following prioritization:

  1. from override config (if it exists)
  2. from the optimizations argument

Custom input arguments

To facilitate testing different configurations, simple_text_demo.py supports argument overrides. The full list of overrides is included in models/tt_transformers/demo/conftest.py.

An example usage where the batch-1 test is modified to run with 16 users and keep generating tokens until 1024 are generated:

pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --batch_size 16 --max_generated_tokens 1024 --stop_at_eos 0

Custom optimizations

To apply the same settings across all decoders, the optimizations argument can be used. optimizations offers a wide range of configurations for precision and math fidelity. The user can override the configurations of the data types of the weight tensors and activation tensors and the math fidelity of the kernels that works on those tensors, using the --optimizations argument on the command line. For example:

pytest models/tt_transformers/demo/simple_text_demo.py -k "accuracy and batch-1" --optimizations 'precision_cfg = {ff1_3: bfp4, ff2: bfp4, wqkv: bfp8, wo: bfp8, kv_cache: bfp8, activation: mixed}, fidelity_cfg = {li_ff1_3: hifi2, li_ff2: lofi, li_qkv_decode: hifi2, li_o_decode: hifi2, sdpa_decode: hifi2na, li_qkv_prefill: hifi2, li_o_prefill: hifi2fp16, sdpa_prefill: hifi4}'

Please refer to model_config.py for the full list of supported key-value pairs in the --optimizations argument. Centralized performance and accuracy targets are defined in models/model_targets.yaml. The lt utility can still export markdown snapshots for local Pareto analysis.

To apply non-uniform settings across the decoders, the user can provide a JSON file using the decoder_config_file argument to specify the configuration for each decoder. For example

pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --decoder_config_file 'models/tt_transformers/demo/config_16_decoders.json'

When a component is not specified (e.g., FF2 is missing for decoder 2 in models/tt_transformers/demo/config_16_decoders.json), the baseline configuration is used for that component.

Expected performance and accuracy

See models/model_targets.yaml for expected performance and accuracy targets across supported configurations. Accuracy of the network architectures is measured by exact token matching using teacher forcing method. During inference the previous token is replaced by the ground truth token while the network generates the next token. This allows to avoid accumulating errors when comparisons on a finer level (tokens) assessed in comparison to other known metrics that compare quality and context of the answer. Token accuracy can be reported by passing the argument shown below:

pytest models/tt_transformers/demo/simple_text_demo.py -k "performance and batch-1" --token_accuracy True

Implementation notes

Chunked prefill (text-only): All of the compatible model/device combinations support a max prefill context-length of 128k, with the exception of Llama3.1-8B and Llama3.2-11B on N150 which have a max of 32k (due to a lack of memory). To support these large max context-lengths, chunked prefill is performed with different max chunk sizes as shown in the table below.

Max Prefill Chunk Sizes (text-only):

N150 N300 T3K TG
Llama3.2-1B 128k tokens 128k tokens 128k tokens 128k tokens
Llama3.2-3B 8k tokens 128k tokens 128k tokens 128k tokens
Qwen2.5-7B 4k tokens 32k tokens 128k tokens 128k tokens
Llama3.1-8B 4k tokens 64k tokens 128k tokens 128k tokens
Llama3.2-11B 4k tokens 64k tokens 128k tokens 128k tokens
Mistral-Small-3.1-24B 8k tokens 128k tokens 128k tokens 128k tokens
Llama3.1-70B Not supported Not supported 32k tokens 128k tokens
Llama3.2-90B Not supported Not supported 32k tokens Not supported
DeepSeek-R1-Distill-Llama3.3-70B Not supported Not supported 32k tokens 128k tokens
  • These max chunk sizes are specific to max context length 128k and are configured via MAX_PREFILL_CHUNK_SIZES_DIV1024 in model_config.py. If the max context length is set to a smaller value using the max_seq_len flag (see Run the demo), these chunk sizes can possibly be increased due to using a smaller KV cache.
  • EXAONE-4.5-33B runs on P150x8 (BH LoudBox) only, with a single 128k-token prefill chunk. Its max context length is 128k: the checkpoint advertises 256k, but chunked prefill is not supported on its sliding-window layers.

Chunked prefill (Llama3.2-11B multimodal): Llama3.2-11B multimodal is currently only supported on N300 and T3000. On N300, a max prefill context length of 8k is supported, while T3000 supports a max context length of 128k.

Chunked prefill (Mistral-Small-3.1-24B multimodal): Mistral-Small-3.1-24B-Instruct-2503 (Pixtral vision) is currently supported on T3000. On T3000, a max prefill context length of 128k is supported.

Chunked prefill (EXAONE-4.5-33B multimodal): EXAONE-4.5-33B is supported on P150x8 as a text decoder. Image input is experimental: the hybrid demo exaone_45_vision_hybrid.py runs the vision tower on the host (or on device with --vision-device tt, using models/experimental/exaone45_vl) and the text decoder on P150x8, with trace disabled for the vision prefill. Vision input is not wired through vLLM.

Memory Optimization

HuggingFace Model Caching Control

To help manage memory usage, you can control whether the HuggingFace model is cached in memory using the cache_hf parameter via command line or code:

# Default: disables caching to conserve memory usage
model_args = ModelArgs(
    mesh_device,
    cache_hf=False,  # Default: Reduces memory usage by not keeping HF model in memory
    max_batch_size=1,
    max_seq_len=2048
)

# Optional: enables caching for faster repeated access
model_args = ModelArgs(
    mesh_device,
    cache_hf=True,  # Cache HF model for better performance running reference tests
    max_batch_size=4,
    max_seq_len=4096
)

When to disable caching (cache_hf=False):

  • Running on systems with limited memory (< 256GB)
  • Loading large models (70B+ parameters)
  • Using the model for single inference runs
  • When you don't need reference model comparisons

When to keep caching enabled (cache_hf=True, default):

  • Sufficient memory available
  • Comparisons with torch model is needed
  • Minimizing test duration is prioritized over memory usage
  • Running reference model tests

The cache_hf parameter affects:

  • load_state_dict() method: Controls whether HF model is cached after loading
  • reference_transformer() method: Controls whether to reuse cached model or load fresh

Memory Impact:

  • Disabling caching saves approximately the full model size in memory
  • For a 70B model, this can save ~140GB+ of memory usage
  • Increased test duration as model needs to be reloaded for reference operations