[Bug] Teacher template encoding fails on Agent tool-call samples during GKD training

[kerbeans]

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

Description

I encountered an HTTP 500 error from the external teacher during Megatron
GKD training on tool-call samples.

The dataset uses ms-swift's Agent dataset format:

• tools is a JSON-serialized string.
• messages contains system, user, assistant, and tool_call roles.
• Tool-call content is a JSON string.

The current configuration uses lmbda=0.0; student-side online vLLM
rollout is not enabled.

The issue still reproduces with official main at commit 426f89710,
installed using:

pip install --no-deps -e .

Dataset format

The pre-training dataset check reports:

{
    'tools': Value('string'),
    'messages': List({
        'role': Value('string'),
        'content': Value('string'),
    }),
    'session_id': Value('string'),
    'trace_id': Value('string'),
    'user_id': Value('string'),
}

The provided dataset contains 1069 samples, all ending with a tool_call
message. The original message content fields are strings, not raw-token
dictionaries.

Sanitized structural example (not yet verified as a standalone minimal
reproducer):

{
  "tools": "[{\"type\":\"function\",\"function\":{\"name\":\"lookup_status\",\"description\":\"Look up an item status\",\"parameters\":{\"type\":\"object\",\"properties\":{\"item_id\":{\"type\":\"string\"}},\"required\":[\"item_id\"]}}}]",
  "messages": [
    {"role": "system", "content": "Use the available tool when needed."},
    {"role": "user", "content": "Check the status of example-item."},
    {
      "role": "tool_call",
      "content": "{\"name\":\"lookup_status\",\"arguments\":{\"item_id\":\"example-item\"}}"
    }
  ]
}

Relevant configuration

Teacher:

• swift deploy --infer_backend vllm
• model_type=qwen3_5
• vllm_tensor_parallel_size=2
• max_logprobs=64

Student:

• Megatron GKD with an external teacher_model_server
• lmbda=1.0, beta=0.0,
• gkd_logits_topk=64
• enable_thinking=false
• add_non_thinking_prefix=true
• loss_scale=last_round+text_rule (custom loss-scale plugin)
• Custom dataset registration plugin
• max_length=122880, truncation_strategy=delete
• padding_free=true
• LoRA training

Actual behavior

The teacher returns HTTP 500 during template encoding.

Relevant traceback, with installation paths shortened:

POST /infer/ HTTP/1.1" 500 Internal Server Error

File "swift/pipelines/infer/deploy.py", line 331, in infer_handler
    return await asyncio.gather(*[self.infer_async(req, request_config) for req in infer_requests])

File "swift/infer_engine/vllm_engine.py", line 922, in infer_async
    inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)

File "swift/template/base.py", line 678, in encode
    encoded = self._encode_truncated(chosen)

File "swift/template/base.py", line 1512, in _encode_truncated
    encoded = Template._encode(self, inputs)

File "swift/template/base.py", line 1571, in _encode
    self._swift_encode(inputs) if template_backend == 'swift' else self._jinja_encode(inputs))

File "swift/template/base.py", line 1430, in _swift_encode
    response_content = self.tokenizer.decode(token_ids[-20:])

File "transformers/tokenization_utils_tokenizers.py", line 1032, in _decode
    text = self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)

TypeError: argument 'ids': 'str' object cannot be interpreted as an integer

Investigation

The earlier investigation identified an intermediate response of this form:

[
    '<think>...</think><tool_call>...</tool_call>',
    {
        'loss_scale': [0, 0, 0, 0, 1],
        'token_ids': [248068, 271, 248069, 271, 248046],
    },
]

This is an intermediate representation, not the source dataset format.
The token IDs shown above are tokenizer-specific.

The relevant code path supports constructing raw-token dictionaries for
teacher requests. Template preprocessing also converts tool calls to
assistant content and merges consecutive assistant messages.

However, _swift_encode() assumes that a list whose final element is
not a string can be passed directly to the tokenizer as integer IDs.
A mixed [str, dict] response violates that assumption.

The latest traceback confirms the same failure location and exception;
it does not itself print the intermediate response.

Expected behavior

The template should support mixed text and raw-token response content.

Suffix inspection should use the trailing raw-token component's
token_ids, while preserving the original response for actual encoding.
Decoding and re-tokenizing the entire response should be avoided because
it can change token alignment.

Related PR

Proposed fix: #10029.

How to Reproduce / 如何复现

1. Version and environment

• ms-swift: 4.6.0.dev0, official main commit 426f89710
• Installation: pip install --no-deps -e .
• OS: Ubuntu 22.04.5 LTS
• Python: 3.12.13
• PyTorch: 2.11.0+cu130
• Transformers: 5.12.0
• tokenizers: 0.22.2
• vLLM: 0.27.1
• Megatron-Core: 0.18.2
• CUDA build version: 13.0
• NVIDIA driver: 580.159.04
• GPUs: 6 × NVIDIA B300 SXM6 AC
• Model type: qwen3_5

2. Prepare an Agent-format dataset

The dataset uses ms-swift's Agent dataset format:
tools is a JSON string, and messages contains tool_call messages
whose content is also a JSON string.

The supplied dataset contains 1069 samples, all ending with a
tool_call message. The pre-training dataset check passes with:

{
    'tools': Value('string'),
    'messages': List({
        'role': Value('string'),
        'content': Value('string'),
    }),
    'session_id': Value('string'),
    'trace_id': Value('string'),
    'user_id': Value('string'),
}

A sanitized structural example is shown below. It illustrates the
format but has not been independently verified as a minimal reproducer:

{"tools":"[{\"type\":\"function\",\"function\":{\"name\":\"lookup_status\",\"description\":\"Look up an item status\",\"parameters\":{\"type\":\"object\",\"properties\":{\"item_id\":{\"type\":\"string\"}},\"required\":[\"item_id\"]}}}]","messages":[{"role":"system","content":"Use the available tool when needed."},{"role":"user","content":"Check the status of example-item."},{"role":"tool_call","content":"{\"name\":\"lookup_status\",\"arguments\":{\"item_id\":\"example-item\"}}"}]}

A custom dataset registration plugin and a custom
last_round+text_rule loss-scale plugin are used in the original run.

3. Start the external teacher

Start a fresh teacher process using the environment above:

CUDA_VISIBLE_DEVICES=4,5 swift deploy \
    --model /path/to/teacher-model \
    --model_type qwen3_5 \
    --infer_backend vllm \
    --host 127.0.0.1 \
    --port 8000 \
    --vllm_tensor_parallel_size 2 \
    --vllm_max_model_len 126976 \
    --max_length 126976 \
    --max_logprobs 64 \
    --load_args false \
    --vllm_gpu_memory_utilization 0.85

Do not reuse a teacher process started from the patched fork when
checking reproduction on official main.

4. Run Megatron GKD training

Relevant configuration from the training script is shown below.
Paths and the registered dataset name have been anonymized.
The custom plugins must be supplied; this excerpt is not standalone.

CUDA_VISIBLE_DEVICES=0,1,2,3 \
NPROC_PER_NODE=4 \
megatron rlhf \
    --rlhf_type gkd \
    --model /path/to/student-model \
    --mcore_model /path/to/mcore-model \
    --model_type qwen3_5 \
    --external_plugins /path/to/loss_scale.py /path/to/dataset.py \
    --dataset example_registered_dataset \
    --loss_scale last_round+text_rule \
    --enable_thinking false \
    --add_non_thinking_prefix true \
    --teacher_model_server http://127.0.0.1:8000 \
    --gkd_logits_topk 64 \
    --lmbda 1.0 \
    --beta 0.0 \
    --sleep-level 1 \
    --max_length 122880 \
    --truncation_strategy delete \
    --tuner_type lora \
    --lora_rank 8 \
    --lora_alpha 16 \
    --target_modules all-linear \
    --merge_lora false \
    --torch_dtype bfloat16 \
    --micro_batch_size 1 \
    --global_batch_size 16 \
    --tensor_model_parallel_size 4 \
    --pipeline_model_parallel_size 1 \
    --sequence_parallel true \
    --attention_backend fused \
    --padding_free true \
    --load_args false

5. Observe the teacher error on tool-call samples

The teacher receives a /infer/ request and fails during template
encoding:

POST /infer/ HTTP/1.1" 500 Internal Server Error

File "swift/template/base.py", line 1430, in _swift_encode
    response_content = self.tokenizer.decode(token_ids[-20:])

TypeError: argument 'ids': 'str' object cannot be interpreted as an integer

The earlier investigation identified an intermediate response containing
both text and a raw-token dictionary:

[
    '<think>...</think><tool_call>...</tool_call>',
    {
        'loss_scale': [0, 0, 0, 0, 1],
        'token_ids': [248068, 271, 248069, 271, 248046],
    },
]

This is an intermediate representation, not the original dataset format.
The old template code treats this mixed list as integer token IDs during
suffix inspection, causing the error above.

Expected behavior: support the mixed representation while preserving
the original raw token IDs and loss masks.

Related fix: #10029.

Additional Information / 补充信息

No response