Fix Gemma 3 mixed-batch token type IDs

[GoGiants1]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Gemma3VisionTemplate._encode previously emitted token_type_ids only for rows containing images. BaseTemplate.data_collator gathers this field only from rows where it is present, so a mixed image/text batch could have different batch dimensions:

input_ids.shape      = [2, sequence_length]
token_type_ids.shape = [1, sequence_length]

Gemma 3's eager causal-mask implementation indexes token_type_ids using the full input batch index, which then accesses a row that does not exist and triggers a CUDA device-side assertion.

This PR initializes token_type_ids to all zeros for every row. Zero is the existing text-token type, so text-only semantics are unchanged. Image rows still overwrite the initialized values with 1 at image-token positions. A lightweight regression test covers the text-only encoding path.

Minimal reproduction

The following script creates one image row and one text-only row, keeps them in the same batch, and runs one eager-attention training step. It requires access to google/gemma-3-4b-pt and one CUDA GPU.

import json
import subprocess
import tempfile
from pathlib import Path

rows = [
    {
        'messages': [
            {'role': 'user', 'content': '<image>Describe this image.'},
            {'role': 'assistant', 'content': 'A demo image.'},
        ],
        'images': ['https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'],
    },
    {
        'messages': [
            {'role': 'user', 'content': 'Say hello.'},
            {'role': 'assistant', 'content': 'Hello.'},
        ],
    },
]

with tempfile.TemporaryDirectory() as tmpdir:
    dataset = Path(tmpdir) / 'mixed.jsonl'
    dataset.write_text(''.join(json.dumps(row) + '\n' for row in rows))
    subprocess.run([
        'swift', 'sft',
        '--model', 'google/gemma-3-4b-pt',
        '--dataset', str(dataset),
        '--output_dir', str(Path(tmpdir) / 'output'),
        '--attn_impl', 'eager',
        '--per_device_train_batch_size', '2',
        '--max_steps', '1',
        '--dataset_shuffle', 'false',
        '--train_dataloader_shuffle', 'false',
        '--split_dataset_ratio', '0',
        '--report_to', 'none',
    ], check=True)

Before this change, the model forward fails while constructing the eager causal mask. The relevant excerpt from the observed multi-GPU training failure is:

File "transformers/models/gemma3/modeling_gemma3.py", line 953, in forward
    "full_attention": create_causal_mask(**mask_kwargs),
File "transformers/masking_utils.py", line 54, in and_mask
    result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device)
torch.AcceleratorError: CUDA error: device-side assert triggered

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:113: operator():
Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.

Experiment results

After the change:

test_text_only_encode_has_token_type_ids ... ok

----------------------------------------------------------------------
Ran 1 test in 0.004s

OK

Validation commands:

PYTHONPATH=. python -m unittest tests.general.test_gemma3_template -v
python -m compileall -q swift/template/templates/gemma.py tests/general/test_gemma3_template.py
git diff --check origin/main...HEAD

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.