liger_kernel with Qwen3 Embedding fine-tuning - KeyError: 'last_hidden_state'

[MLKoz2]

Describe the bug
KeyError: 'last_hidden_state' during fine-tuning Qwen3 Embedding model with liger_kernel

Your hardware and system info
RTX 6000 Ada

Additional context
with liger_kernel false - everything is fine

[tastelikefeet]

Please attach your cmd and let me check

[JhaceLam]

@tastelikefeet I had the same issue.

nproc_per_node=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)
CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((nproc_per_node-1)))  \
NPROC_PER_NODE=$nproc_per_node \
MASTER_PORT=$(echo $ARNOLD_WORKER_0_PORT | cut -d ',' -f 1) \swift sft \
    --model ../models/Qwen3-Embedding-0.6B \
    --task_type embedding \
    --model_type qwen3_emb \
    --train_type full \
    --dataset \
        ../data/train_original_part1of4.jsonl \
        ../data/train_original_part2of4.jsonl \
    --split_dataset_ratio 0.05 \
    --eval_strategy steps \
    --output_dir output \
    --eval_steps 100 \
    --num_train_epochs 1 \
    --save_steps 100 \
    --save_total_limit 5 \
    --load_best_model_at_end \
    --per_device_train_batch_size 32 \
    --per_device_eval_batch_size 32 \
    --gradient_accumulation_steps 4 \
    --learning_rate 6e-6 \
    --loss_type infonce \
    --label_names labels \
    --dataloader_drop_last true \
    --deepspeed zero2 \
    --attn_impl flash_attn \
    --use_liger_kernel true

And followings is the detailed error of rank0 (error from other ranks are similiar):

Train:   0%|          | 0/1855 [00:00<?, ?it/s][rank1]: Traceback (most recent call last):
[rank1]:   File "/opt/tiger/ms-swift/swift/cli/sft.py", line 12, in <module>
[rank1]:     sft_main()
[rank1]:   File "/opt/tiger/ms-swift/swift/llm/train/sft.py", line 274, in sft_main
[rank1]:     return SwiftSft(args).main()
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/llm/base.py", line 49, in main
[rank1]:     result = self.run()
[rank1]:              ^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/llm/train/sft.py", line 122, in run
[rank1]:     return self.train(trainer)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/llm/train/sft.py", line 183, in train
[rank1]:     trainer.train(trainer.args.resume_from_checkpoint)
[rank1]:   File "/opt/tiger/ms-swift/swift/trainers/trainers.py", line 55, in train
[rank1]:     return super().train(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/trainers/mixin.py", line 419, in train
[rank1]:     res = super().train(*args, **kwargs)
[rank1]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/tiger/.local/lib/python3.11/site-packages/transformers/trainer.py", line 2245, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/tiger/.local/lib/python3.11/site-packages/transformers/trainer.py", line 2560, in _inner_training_loop
[rank1]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/tiger/.local/lib/python3.11/site-packages/transformers/trainer.py", line 3736, in training_step
[rank1]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/trainers/trainers.py", line 58, in compute_loss
[rank1]:     loss, outputs = super().compute_loss(model, inputs, return_outputs=True)
[rank1]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/tiger/.local/lib/python3.11/site-packages/transformers/trainer.py", line 3815, in compute_loss
[rank1]:     loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/opt/tiger/ms-swift/swift/plugin/loss.py", line 408, in infonce_loss
[rank1]:     sentences = outputs['last_hidden_state']
[rank1]:                 ~~~~~~~^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/tiger/.local/lib/python3.11/site-packages/transformers/utils/generic.py", line 445, in __getitem__
[rank1]:     return inner_dict[k]
[rank1]:            ~~~~~~~~~~^^^
[rank1]: KeyError: 'last_hidden_state'

[tastelikefeet]

let me check

[tastelikefeet]

Fixed by #5222