Cannot load Qwen3 models

[toanhuynhnguyen]

Feature request

I cannot load Qwen3 models, such as when I run this command:

docker run --entrypoint optimum-cli \
    -v $(pwd)/data:/data \
    --privileged \
    -e HF_TOKEN=${HF_TOKEN} \
    -e HF_AUTO_CAST_TYPE="bf16" \
       ghcr.io/huggingface/neuronx-tgi:latest \
       export neuron \
       --model Qwen/Qwen3-8B \
       --auto_cast_type bf16 \
       --batch_size 20 \
       --sequence_length 13000 \
       /data1/qwen3_neuron_8B_bz20_13k

I get this error:

ubuntu@ip-172-31-50-88:~$ docker run --entrypoint optimum-cli \
    -v $(pwd)/data:/data \
    --privileged \
    -e HF_TOKEN=${HF_TOKEN} \
    -e HF_AUTO_CAST_TYPE="bf16" \
       ghcr.io/huggingface/neuronx-tgi:latest \
       export neuron \
       --model Qwen/Qwen3-8B \
       --auto_cast_type bf16 \
       --batch_size 20 \
       --sequence_length 13000 \
       /data1/qwen3_neuron_8B_bz20_13k
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/configuration_auto.py", line 1071, in from_pretrained
    config_class = CONFIG_MAPPING[config_dict["model_type"]]
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/configuration_auto.py", line 773, in __getitem__
    raise KeyError(key)
KeyError: 'qwen3'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/neuron/__main__.py", line 786, in <module>
    main()
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/neuron/__main__.py", line 732, in main
    input_shapes, neuron_config_class = get_input_shapes_and_config_class(task, args)
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/neuron/__main__.py", line 125, in get_input_shapes_and_config_class
    neuron_config_constructor = get_neuron_config_class(task, args.model)
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/neuron/__main__.py", line 132, in get_neuron_config_class
    config = AutoConfig.from_pretrained(model_id)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/configuration_auto.py", line 1073, in from_pretrained
    raise ValueError(
ValueError: The checkpoint you are trying to load has model type `qwen3` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.

You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is very new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git`
Traceback (most recent call last):
  File "/usr/local/bin/optimum-cli", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/optimum_cli.py", line 208, in main
    service.run()
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/export/neuronx.py", line 305, in run
    subprocess.run(full_command, shell=True, check=True)
  File "/usr/lib/python3.10/subprocess.py", line 526, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command 'python3 -m optimum.exporters.neuron --model Qwen/Qwen3-8B --auto_cast_type bf16 --batch_size 20 --sequence_length 13000 /data1/qwen3_neuron_8B_bz20_13k' returned non-zero exit status 1.

Motivation

Cannot load Qwen3 models

Your contribution

Cannot load Qwen3 models

[jimburtoft]

@toanhuynhnguyen

Support for Optimum Neuron in TGI has been upstreamed to the main TGI project. You want to use this docker image:

docker pull ghcr.io/huggingface/text-generation-inference:latest-neuron

(The image you are pulling is based on ON 0.0.28. The upstreamed version is 0.3.0)

You can see more examples here: https://github.com/huggingface/optimum-neuron/tree/main/benchmark/text-generation-inference/performance

[toanhuynhnguyen]

@jimburtoft Thank you so much.

[jimburtoft]

Hey @toanhuynhnguyen, it looks like the latest container only has Optimum Neuron 0.2.2.

[dacorvo]

The latest TGI image for neuron is now based on 0.3.0

[jimburtoft]

@dacorvo Maybe the latest-neuron tag isn't updated?

I did this test and it is showing an older version:

ubuntu@ip-172-31-51-185:~$ docker pull ghcr.io/huggingface/text-generation-inference:latest-neuron
latest-neuron: Pulling from huggingface/text-generation-inference
Digest: sha256:257123c725312dd67a29f050d6780a423648882c0f85a27790dd4a1152748f81
Status: Image is up to date for ghcr.io/huggingface/text-generation-inference:latest-neuron
ghcr.io/huggingface/text-generation-inference:latest-neuron

ubuntu@ip-172-31-51-185:~$ docker run -v $(pwd):/data --privileged --rm -it --entrypoint /bin/bash ghcr.io/huggingface/text-generation-inference:latest-neuron
root@fe548600047c:/# pip list | grep neuron
libneuronxla             2.2.1630.0
neuronx-cc               2.17.194.0+d312836f
neuronx-distributed      0.11.0
optimum-neuron           0.2.0
torch-neuronx            2.5.1.2.6.0
root@fe548600047c:/# 

[dacorvo]

You're probably right, it might be updated only when a new version is released. The alternative is to use the versions tagged by SHA.
Anyway, since @tengomucho released TGI 3.3.5, it should be good now.

[jimburtoft]

Yeah, the latest image is from 2 months ago. I'm assuming TGI 3.3.5 will have to trickle through the build/testing process before the updated images show up. Until then, I can build it myself or wait. Thanks!

[jimburtoft]

@toanhuynhnguyen I just checked and the TGI image was updated a while ago to include the latest Optimum Neuron code that can handle the Qwen models. As of today, TGI 3.3.6 was available as the latest:

https://github.com/huggingface/text-generation-inference/pkgs/container/text-generation-inference/517320827?tag=latest-neuron

[toanhuynhnguyen]

@jimburtoft Thank you so much.

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[github-actions]

This issue was closed because it has been stalled for 5 days with no activity.