Skip to content

Fix 0.6B fine-tuning crash: project text embedding before adding to c…#336

Open
kru2710shna wants to merge 1 commit into
QwenLM:mainfrom
kru2710shna:fix/0.6b-text-projection-finetune
Open

Fix 0.6B fine-tuning crash: project text embedding before adding to c…#336
kru2710shna wants to merge 1 commit into
QwenLM:mainfrom
kru2710shna:fix/0.6b-text-projection-finetune

Conversation

@kru2710shna

Copy link
Copy Markdown

CAUSE
In finetuning/sft_12hz.py, the text and codec embeddings are added element-wise:

input_embeddings = input_text_embedding + input_codec_embedding

But they are defined at different widths in the talker model (modeling_qwen3_tts.py):

self.codec_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
self.text_embedding  = nn.Embedding(config.text_vocab_size, config.text_hidden_size)

So text_embedding has width text_hidden_size and codec_embedding has width hidden_size. These differ for 0.6B Base (text_hidden_size=2048, hidden_size=1024), so the add fails with:

RuntimeError: The size of tensor a (2048) must match the size of tensor b (1024) at non-singleton dimension 2

FIX
The talker already defines text_projection (Qwen3TTSTalkerResizeMLP, mapping text_hidden_size to hidden_size) and applies it on the inference path before combining text and codec embeddings. The fine-tuning script omitted it. This PR projects the text embedding through text_projection before the add, matching inference.
Because text_projection uses bias=True, the mask is re-applied after projection so padded positions remain zero:

input_text_embedding = model.talker.model.text_embedding(input_text_ids)
input_text_embedding = model.talker.text_projection(input_text_embedding)   # text_hidden_size -> hidden_size
input_text_embedding = input_text_embedding * text_embedding_mask           # re-zero padding after projection

Verification
Reproduced the crash on Qwen3-TTS-12Hz-0.6B-Base with training data prepared via prepare_data.py:

Before: RuntimeError: tensor a (2048) must match tensor b (1024) at the embedding add.
After: training proceeds and the step completes with finite loss (Epoch 0 | Step 0 | Loss: 11.82).

Fixes #335

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

An error occurred while fine-tuning the 0.6B model using the Fine-Tuning Script within the repository

1 participant