diff --git a/README.md b/README.md index 6f4b05a0..2bc24ab5 100644 --- a/README.md +++ b/README.md @@ -405,6 +405,7 @@ foundation for building customizable, enterprise-grade training services. | Component Type | Component Link | Component Function | Author | | -------------- | -------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- | ------------------- | +| Cookbook | [Qwen3-TTS LoRA](cookbook/qwen3_tts/README.md) | Single Ascend NPU training, explicit data splits and update-boundary resume | [Wallbreazzz](https://github.com/Wallbreazzz) | | Patch | [qwen3_moe_transformers4_patch](https://www.modelscope.cn/models/twinkle-kit/qwen3_moe_transformers4_patch) | Fixes Qwen3 MoE model hang issue during FSDP2 training, effective for transformers==4.x | ModelScope Official | ## Contributions diff --git a/cookbook/qwen3_tts/.gitignore b/cookbook/qwen3_tts/.gitignore new file mode 100644 index 00000000..76782217 --- /dev/null +++ b/cookbook/qwen3_tts/.gitignore @@ -0,0 +1,3 @@ +outputs/ +__pycache__/ +*.py[cod] diff --git a/cookbook/qwen3_tts/README.md b/cookbook/qwen3_tts/README.md new file mode 100644 index 00000000..353c6bff --- /dev/null +++ b/cookbook/qwen3_tts/README.md @@ -0,0 +1,162 @@ +# Qwen3-TTS LoRA training on one Ascend NPU + +This cookbook connects the Qwen3-TTS 12Hz Base model to Twinkle's local +Transformers training API. It provides explicit train/dev preparation, +single-sample micro-batches, LoRA on the Talker and Code Predictor, accumulation +with a final short window, and restart at completed optimizer updates. + +It is a single-card reference implementation, not a quality-tuned training +recipe. There are no changes to Twinkle core or its default dependencies. +Generation, speaker-similarity evaluation and experimental EOS handling are +outside this contribution. + +## Installation and supported scope + +Run from the repository root with an existing Ascend PyTorch environment: + +```bash +pip install -e . +pip install -r cookbook/qwen3_tts/requirements.txt +cd cookbook/qwen3_tts +export MODEL_PATH=/models/Qwen3-TTS-12Hz-1.7B-Base +export ASCEND_RT_VISIBLE_DEVICES=0 +``` + +The model directory must already contain the complete Base weights, config, +text tokenizer, processor and `speech_tokenizer/`. Preparation and training are +offline: missing assets cause an error, not a download. No virtual environment +or model download is created by these commands. Select one physical card with +`ASCEND_RT_VISIBLE_DEVICES`; `npu:0` is relative to that selection. + +The development baseline used Twinkle `5337cc3`, Qwen3-TTS `022e286`, Python +3.12, torch 2.10.0, torch_npu 2.10.0.post2, CANN 9.1.0 and an Ascend 910B. +The optional requirements pin the Qwen implementation and its Transformers +4.57.3 / accelerate 1.12.0 dependencies. They are not installed for ordinary +Twinkle users. Match torch/torch_npu to your Ascend image before installing +these extras. Codec and reference mel extraction use CPU fp32; the TTS model +uses NPU bf16 with eager attention. Multi-card, CUDA, padding-free execution +and other Qwen3-TTS model variants are not covered by the current evidence. + +## Prepare an existing split + +Provide two UTF-8 JSONL files. Each row has `id`, `text`, `audio`, `group_id` +and `speaker_id`. Paths are relative to the input JSONL file. For example: + +```json +{"id":"train-001","text":"A training utterance.","audio":"audio/train-001.wav","group_id":"session-a","speaker_id":"speaker-a"} +``` + +Use a different held-out group and text in the dev file. Audio must be finite +mono 16 or 24 kHz; 16 kHz is resampled to 24 kHz. The preparer requires one +speaker, unique IDs, nonempty train/dev sets and a reference chosen from train. +It rejects cross-split group, normalized text, file, decoded PCM and resampled +PCM duplicates. Group identifiers still need to reflect the real recording +source; these checks do not detect all near-duplicates. + +```bash +python scripts/prepare_split_manifest.py \ + --model_path "$MODEL_PATH" \ + --train_jsonl /data/splits/train.jsonl \ + --dev_jsonl /data/splits/dev.jsonl \ + --reference_id train-001 \ + --output_dir /data/derived/qwen3-tts +``` + +The output directory must be new. The script preserves the assigned split, +text and order; it never silently resplits or edits source audio. Output rows +contain `id`, `text`, integer `audio_codes` of shape `[T,16]` in `[0,2048)`, +and `ref_mel_path`. One train-reference mel `[F,128]` is shared by both splits. +`provenance.json` records input/output hashes and the source mapping. +The dev manifest is prepared for downstream evaluation and is **not** consumed +by this training loop. No dev loss, early stopping or quality score is implied. + +## Train and resume + +```bash +python scripts/train_qwen3_tts_lora.py \ + --model_path "$MODEL_PATH" \ + --train_jsonl /data/derived/qwen3-tts/train.jsonl \ + --output_dir outputs/run-1 \ + --epochs 3 --grad_accum 4 --shuffle --seed 42 \ + --lr 1e-4 --lr_decay 1.0 --lora_r 8 --lora_dropout 0.1 \ + --save_interval 10 +``` + +These are example hyperparameters, not a convergence recommendation. The +default learning rate stays constant. `--lr_decay` optionally supplies StepLR +gamma per optimizer update; the short resume regression explicitly uses 0.8 +to exercise scheduler restoration. LoRA targets `q_proj`/`v_proj` in both +branches, with alpha equal to twice the rank. All base parameters are frozen. + +Micro-batch size is one. The loop owns accumulation windows and keeps +Twinkle's group accumulation at one; `clip_grad_and_step` normalizes by the +actual number of micro-batches, including an epoch's short final window. +The speaker encoder runs without gradients. The Talker loss uses its internal +temporal shift; residual-codebook logits are supervised against corresponding +codebooks 1–15 without another shift along the codebook axis. Tests check +last-head gradients and the official residual assembly interface. This loss +choice is not a claim about how official pretrained weights were trained. + +For a controlled stop, add `--stop_after_updates 1`. Resume to a **new** output +directory with the same training options: + +```bash +python scripts/train_qwen3_tts_lora.py \ + --model_path "$MODEL_PATH" \ + --train_jsonl /data/derived/qwen3-tts/train.jsonl \ + --output_dir outputs/run-2 --resume_from outputs/run-1/checkpoint-update-1 \ + --epochs 3 --grad_accum 4 --shuffle --seed 42 \ + --lr 1e-4 --lr_decay 1.0 --lora_r 8 --lora_dropout 0.1 \ + --save_interval 10 +``` + +Checkpoints include adapter, optimizer, scheduler, RNG and loop progress. +Resume checks content fingerprints of model/manifest/reference mel, sample +plan and settings; it fails on mismatch instead of silently starting over. +Model hashing reads the local weight files and may take time. Only complete, +trusted checkpoints produced by this workflow should be resumed. Saving is +at update boundaries; recovery from arbitrary interruption during a checkpoint +write is not guaranteed. Inspect `trace.jsonl`, `run_config.json` and +`result.json` for the actual update count and checkpoint path. + +## Tests and evidence boundaries + +With the optional dependencies installed: + +```bash +python tests/run.py +MODEL_PATH="$MODEL_PATH" python tests/run.py --npu +``` + +The first command runs data, forward/loss, projection, manifest, preparation +and accumulation contracts. It includes official mel/residual-interface tests +but uses small CPU networks, not the pretrained TTS model. The second also +loads the real local model for a LoRA update/save/reload, public split +preparation-to-training, and continuous versus independent-process resume. +The two tiny LibriSpeech fixtures are attributed under +`tests/fixtures/librispeech/`; they are regression fixtures, not an independent +quality benchmark. Tests use temporary outputs and never require user data. + +For local CPU development without Twinkle or Qwen installed, install torch, +Transformers, NumPy, scipy and soundfile, then run `python tests/run.py +--isolated`. This mode explicitly substitutes Twinkle infrastructure and skips +two official-package checks. It is not NPU or native Twinkle validation. + +Before upstream packaging, real 1.7B single-NPU runs covered nonzero LoRA +gradients/updates, unchanged frozen parameters, exact saved/reloaded adapter +weights, accumulation tails and restart in separate processes. Final adapter +weights matched exactly in the resume comparison; optimizer states were +compared with rtol=1e-3/atol=1e-5. Those results do not establish long-training +stability, speaker similarity or fine-tuning quality. The reorganized PR tree +requires its own final NPU regression before the draft is marked ready. + +## Integration boundaries + +The local `tts/` package is added to the script path; it is not registered as +a new global Twinkle task. The subclass reuses forward/backward, optimizer, +clipping, saving and resume APIs. It currently overrides +`_construct_default_optimizer_group` for consistent tokenizer/processor/loss +defaults and calls `_lazy_wrap_model` when inspecting/loading wrapped weights. +These internal dependencies are isolated here and should be reviewed when +Twinkle changes. A public generic TTS abstraction should only be added if +maintainers identify a reusable need; this cookbook does not add one prematurely. diff --git a/cookbook/qwen3_tts/requirements.txt b/cookbook/qwen3_tts/requirements.txt new file mode 100644 index 00000000..50874f35 --- /dev/null +++ b/cookbook/qwen3_tts/requirements.txt @@ -0,0 +1,8 @@ +accelerate==1.12.0 +numpy>=2.0,<2.3 +peft==0.17.1 +qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@022e286b98fbec7e1e916cb940cdf532cd9f488e +safetensors +scipy +soundfile +transformers==4.57.3 diff --git a/cookbook/qwen3_tts/scripts/audio_features.py b/cookbook/qwen3_tts/scripts/audio_features.py new file mode 100644 index 00000000..cf955bd9 --- /dev/null +++ b/cookbook/qwen3_tts/scripts/audio_features.py @@ -0,0 +1,42 @@ +"""Official 24 kHz mel extraction and 16-codebook audio encoding.""" + +import numpy as np + +MEL_DIM = 128 +MEL_HOP = 256 +MEL_WIN = 1024 +MEL_FMAX = 12000.0 +TARGET_SR = 24000 + +def extract_mel(wav_24k: np.ndarray) -> np.ndarray: + """使用官方 mel 算法，返回单条 [F,128]，不使用替代 fbank。""" + import torch + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + + wav = np.asarray(wav_24k, dtype=np.float32) + if wav.ndim != 1: + raise ValueError('Expected mono waveform [samples]') + with torch.no_grad(): + mel = mel_spectrogram( + torch.from_numpy(wav).unsqueeze(0), n_fft=MEL_WIN, + num_mels=MEL_DIM, sampling_rate=TARGET_SR, hop_size=MEL_HOP, + win_size=MEL_WIN, fmin=0, fmax=MEL_FMAX, + ).transpose(1, 2) + return mel[0].cpu().numpy() + +def encode_audio_codes(wav_24k: np.ndarray, tokenizer) -> np.ndarray: + """单条输入必须返回一个 [T,16]；兼容列表、二维张量和单批张量。""" + result = tokenizer.encode(wav_24k, sr=TARGET_SR) # wav_24k 已升采样到 24kHz + codes = result.audio_codes + if isinstance(codes, (list, tuple)): + if len(codes) != 1: + raise ValueError('Single waveform must return exactly one audio_codes entry') + codes = codes[0] + if hasattr(codes, 'cpu'): + codes = codes.cpu().numpy() + codes = np.asarray(codes) + if codes.ndim == 3 and codes.shape[0] == 1: + codes = codes[0] + if codes.ndim != 2 or codes.shape[0] == 0 or codes.shape[1] != 16: + raise ValueError(f'Expected nonempty audio_codes [T,16], got {codes.shape}') + return codes.astype(np.int64) diff --git a/cookbook/qwen3_tts/scripts/prepare_split_manifest.py b/cookbook/qwen3_tts/scripts/prepare_split_manifest.py new file mode 100644 index 00000000..50b039c2 --- /dev/null +++ b/cookbook/qwen3_tts/scripts/prepare_split_manifest.py @@ -0,0 +1,128 @@ +"""Preserve an explicitly assigned single-speaker train/dev split; never resplit.""" +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import unicodedata + +import numpy as np + + +def sha(data): + return hashlib.sha256(data).hexdigest() + + +def read_audio(path): + import soundfile as sf + from scipy.signal import resample_poly + wave, sr = sf.read(path, dtype='float32') + if sr not in (16000, 24000) or wave.ndim != 1 or not len(wave) or not np.isfinite(wave).all(): + raise ValueError('Expected finite mono 16/24 kHz audio') + original = sha(np.asarray(wave, dtype='= 2048).any()): + raise ValueError('Expected integer 12Hz codes [T,16] in [0,2048)') + encoded.append(dict(id=row['id'], text=row['text'], audio_codes=codes.tolist(), + ref_mel_path='reference.npy')) + mapping.append(dict(row, split=split, reference_id=reference_id, + output_manifest=f'{split}.jsonl')) + (out / f'{split}.jsonl').write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in encoded), encoding='utf-8') + report = dict(format_version=1, counts={k: len(v) for k, v in splits.items()}, + reference_id=reference_id, reference_split='train', mapping=mapping, + source_manifest_sha256={k: sha(Path(p).read_bytes()) for k, p in + [('train', train_path), ('dev', dev_path)]}, + output_sha256={n: sha((out / n).read_bytes()) for n in + ('train.jsonl', 'dev.jsonl', 'reference.npy')}) + (out / 'provenance.json').write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + return report + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + for name in ('train_jsonl', 'dev_jsonl', 'output_dir', 'reference_id', 'model_path'): + parser.add_argument('--' + name, required=True) + args = parser.parse_args() + # Audit before loading weights or creating output. No automatic split/filter fallback. + audit_split(args.train_jsonl, args.dev_jsonl, args.reference_id) + os.environ['HF_HUB_OFFLINE'] = '1' + os.environ['TRANSFORMERS_OFFLINE'] = '1' + import torch + from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSTokenizer + from audio_features import encode_audio_codes, extract_mel + codec = Qwen3TTSTokenizer.from_pretrained(str(Path(args.model_path) / 'speech_tokenizer'), + local_files_only=True, device_map='cpu', torch_dtype=torch.float32) + codec.model.float().eval() + with torch.no_grad(): + report = prepare(args.train_jsonl, args.dev_jsonl, args.output_dir, args.reference_id, + lambda wave: encode_audio_codes(wave, codec), extract_mel) + from train_qwen3_tts_lora import model_fingerprint + report['codec_content_sha256'] = model_fingerprint(Path(args.model_path) / 'speech_tokenizer') + (Path(args.output_dir) / 'provenance.json').write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print('准备完成 / Preparation completed: ' + json.dumps(report['counts']), flush=True) + + +if __name__ == '__main__': + main() diff --git a/cookbook/qwen3_tts/scripts/train_qwen3_tts_lora.py b/cookbook/qwen3_tts/scripts/train_qwen3_tts_lora.py new file mode 100644 index 00000000..6b955167 --- /dev/null +++ b/cookbook/qwen3_tts/scripts/train_qwen3_tts_lora.py @@ -0,0 +1,152 @@ +"""Manifest-based single-NPU bf16 LoRA training with boundary-only resume. + +Each JSONL row: id, text, audio_codes [T,16], ref_mel_path [F,128] .npy. +Output directory must be new; resume writes to a new run directory. +""" +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import sys + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + + +def model_fingerprint(path): + digest = hashlib.sha256() + files = sorted(p for p in path.rglob('*') if p.is_file() and p.suffix in ('.safetensors', '.json', '.txt')) + if not any(p.suffix == '.safetensors' for p in files): + raise ValueError('Local model must contain safetensors weights') + for file in files: + digest.update(file.relative_to(path).as_posix().encode()) + with file.open('rb') as source: + for block in iter(lambda: source.read(1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +def parse_args(argv=None): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument('--model_path', required=True) + ap.add_argument('--train_jsonl', required=True) + ap.add_argument('--output_dir', required=True) + ap.add_argument('--resume_from') + ap.add_argument('--epochs', type=int, default=1) + ap.add_argument('--grad_accum', type=int, default=2) + ap.add_argument('--seed', type=int, default=42) + ap.add_argument('--shuffle', action='store_true') + ap.add_argument('--lr', type=float, default=1e-4) + ap.add_argument('--lr_decay', type=float, default=1.0, help='StepLR gamma per optimizer update; default keeps LR constant') + ap.add_argument('--lora_r', type=int, default=8) + ap.add_argument('--lora_dropout', type=float, default=0.1) + ap.add_argument('--save_interval', type=int, default=1, help='Optimizer updates between checkpoints') + ap.add_argument('--stop_after_updates', type=int, default=0, help='Absolute update limit; 0 runs all epochs') + args = ap.parse_args(argv) + if min(args.epochs, args.grad_accum, args.lora_r, args.save_interval) < 1 or args.stop_after_updates < 0: + ap.error('Counts must be positive (stop_after_updates may be 0)') + if not math.isfinite(args.lr_decay) or not 0 < args.lr_decay <= 1: + ap.error('lr_decay must be in (0, 1]') + if not math.isfinite(args.lr) or not args.lr > 0 or not 0 <= args.lora_dropout < 1: + ap.error('lr must be positive and dropout in [0,1)') + return args + + +def train(): + args = parse_args() + os.environ['HF_HUB_OFFLINE'] = '1' + os.environ['TRANSFORMERS_OFFLINE'] = '1' + import random + import numpy as np + import torch + import torch_npu # noqa: F401 + import twinkle + from twinkle import DeviceMesh + from peft import LoraConfig + from tts.qwen3.manifest import Manifest, sample_plan, planned_windows, validate_cursor + from tts.qwen3.runtime import load_twinkle_model, data_config, ADAPTER + from tts.qwen3.tokenization import load_text_tokenizer + from tts.qwen3.data import build_assistant_text, build_qwen3_sample, collate_qwen3_batch + from tts.qwen3.input_processor import Qwen3TTSInputProcessor + from tts.qwen3.loss import Qwen3TTSLoss + from tts.qwen3.training import Progress, train_window, save_boundary, resume_boundary + + manifest = Manifest(args.train_jsonl) + plan = sample_plan(len(manifest), args.epochs, args.seed, args.shuffle) + schedule = planned_windows(plan, args.grad_accum) + path = Path(args.model_path).resolve(strict=True) + print('计算模型内容指纹 / Hashing local model files', flush=True) + config = dict(format_version=1, manifest=manifest.fingerprint, model=model_fingerprint(path), + plan=plan, epochs=args.epochs, accumulation_steps=args.grad_accum, seed=args.seed, + shuffle=args.shuffle, lr=args.lr, lora_r=args.lora_r, lora_dropout=args.lora_dropout, + targets=['q_proj', 'v_proj'], scheduler=f'StepLR(1,{args.lr_decay})', tokenizer_fix=True) + fingerprint = hashlib.sha256(json.dumps(config, sort_keys=True).encode()).hexdigest() + out = Path(args.output_dir).resolve() + out.mkdir(parents=True, exist_ok=False) + (out / 'run_config.json').write_text(json.dumps(config, indent=2), encoding='utf-8') + torch.npu.set_device(0) + device_mesh = DeviceMesh.from_sizes(fsdp_size=1, dp_size=1) + twinkle.initialize(mode='local', global_device_mesh=device_mesh) + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.npu.manual_seed(args.seed) + model = load_twinkle_model(path, device_mesh) + model.model.speech_tokenizer = None # manifest is already codec-encoded + model.add_adapter_to_model(ADAPTER, LoraConfig(r=args.lora_r, lora_alpha=2 * args.lora_r, + target_modules=['q_proj', 'v_proj'], lora_dropout=args.lora_dropout, bias='none'), + gradient_accumulation_steps=1) + model.set_processor(Qwen3TTSInputProcessor, adapter_name=ADAPTER) + model.set_loss(Qwen3TTSLoss(), adapter_name=ADAPTER) + model.set_optimizer('AdamW', lr=args.lr, weight_decay=0.0, adapter_name=ADAPTER) + model.set_lr_scheduler('StepLR', step_size=1, gamma=args.lr_decay, adapter_name=ADAPTER) + model._lazy_wrap_model() + cfg = data_config(model.hf_config) + tokenizer = load_text_tokenizer(path) + suffix = tokenizer.encode('<|im_end|>\n<|im_start|>assistant\n', add_special_tokens=False) + + def batch(index): + row = manifest.rows[index] + ids = tokenizer.encode(build_assistant_text(row['text'])) + if not suffix or ids[-len(suffix):] != suffix: + raise ValueError('Unexpected assistant suffix') + sample = build_qwen3_sample(np.asarray(ids[:-len(suffix)]), np.asarray(row['audio_codes']), cfg) + mel = np.load(row['ref_mel_path'], allow_pickle=False) + raw = collate_qwen3_batch([sample], cfg, [mel]) + return model.optimizer_group[ADAPTER].processor(raw, model=model.strategy.unwrap_model(model.model)) + + progress = Progress(args.grad_accum) + if args.resume_from: + progress = resume_boundary(model, args.resume_from, ADAPTER, fingerprint, args.grad_accum) + validate_cursor(progress, schedule) + final_checkpoint = None + trace = out / 'trace.jsonl' + with trace.open('w', encoding='utf-8') as log: + for indices in schedule[progress.optimizer_steps:]: + if args.stop_after_updates and progress.optimizer_steps >= args.stop_after_updates: + break + losses = train_window(model, [batch(i) for i in indices], progress, ADAPTER) + record = dict(micro_steps=progress.micro_steps, optimizer_steps=progress.optimizer_steps, + consumed_samples=progress.consumed_samples, sample_ids=[manifest.rows[i]['id'] for i in indices], + losses=losses, learning_rates=model.optimizer_group[ADAPTER].lr_scheduler.get_last_lr()) + log.write(json.dumps(record, ensure_ascii=False) + '\n') + log.flush() + print(json.dumps(record, ensure_ascii=False), flush=True) + if progress.optimizer_steps % args.save_interval == 0: + final_checkpoint = save_boundary(model, progress, out, + f'checkpoint-update-{progress.optimizer_steps}', ADAPTER, fingerprint) + if progress.optimizer_steps and (final_checkpoint is None or + final_checkpoint.name != f'checkpoint-update-{progress.optimizer_steps}'): + final_checkpoint = save_boundary(model, progress, out, + f'checkpoint-update-{progress.optimizer_steps}', ADAPTER, fingerprint) + (out / 'result.json').write_text(json.dumps(dict(pid=os.getpid(), checkpoint=str(final_checkpoint) if final_checkpoint else None, + micro_steps=progress.micro_steps, optimizer_steps=progress.optimizer_steps, + consumed_samples=progress.consumed_samples, complete=progress.optimizer_steps == len(schedule)), + indent=2), encoding='utf-8') + print('本次运行完成 / Run completed', flush=True) + + +if __name__ == '__main__': + train() diff --git a/cookbook/qwen3_tts/tests/batch1/isolated_support.py b/cookbook/qwen3_tts/tests/batch1/isolated_support.py new file mode 100644 index 00000000..adf1a8c6 --- /dev/null +++ b/cookbook/qwen3_tts/tests/batch1/isolated_support.py @@ -0,0 +1,46 @@ +"""Explicit CPU test doubles for Twinkle infrastructure only. + +Production tensor operations and transformers loss remain real. This is NOT +Twinkle/NPU integration evidence. Only loaded with --isolated-twinkle. +""" +import sys +import types +from pathlib import Path +import torch + + +def install(root): + def module(name, **attrs): + obj = types.ModuleType(name) + obj.__dict__.update(attrs) + sys.modules[name] = obj + return obj + + for name, path in [('tts', 'tts'), ('tts.common', 'tts/common'), ('tts.qwen3', 'tts/qwen3')]: + module(name, __path__=[str(root / path)]) + + class Processor: + padding_map = {} + VLM_CONCAT_FIELDS = set() + def __init__(self, **kwargs): + self.padding_free = kwargs.get('padding_free', False) + + def pad_stack(values, pad_value, concat): + assert concat is False + return torch.nn.utils.rnn.pad_sequence(values, batch_first=True, padding_value=pad_value) + + module('twinkle', remote_function=lambda **kw: lambda fn: fn, + Platform=types.SimpleNamespace(get_local_device=lambda: 'cpu')) + module('twinkle.data_format', InputFeature=dict, Trajectory=dict, LossOutput=dict) + module('twinkle.data_format.output', LossOutput=dict) + module('twinkle.processor') + module('twinkle.processor.base', InputProcessor=Processor) + module('twinkle.utils') + module('twinkle.utils.torch_utils', pad_and_stack_tensors=pad_stack) + module('twinkle.loss') + module('twinkle.loss.base', Loss=object) + module('twinkle.infra', collect_tensor_dict=lambda x: x) + module('twinkle.model') + module('twinkle.model.transformers', TransformersModel=object) + from tts.common.multi_codebook_loss import MultiCodebookLoss + sys.modules['tts.common'].MultiCodebookLoss = MultiCodebookLoss diff --git a/cookbook/qwen3_tts/tests/batch1/test_contracts.py b/cookbook/qwen3_tts/tests/batch1/test_contracts.py new file mode 100644 index 00000000..8f907623 --- /dev/null +++ b/cookbook/qwen3_tts/tests/batch1/test_contracts.py @@ -0,0 +1,273 @@ +"""Run: python tests/batch1/test_contracts.py [--isolated-twinkle]. + +Default requires installed Twinkle + Qwen3-TTS + transformers. Isolated mode +substitutes only Twinkle infrastructure, and skips official mel parity. +All tensors, autograd and ForCausalLMLoss use real torch/transformers. +""" +import ast +import importlib.util +import inspect +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import torch +from torch import nn +from transformers.loss.loss_utils import ForCausalLMLoss + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +ISOLATED = '--isolated-twinkle' in sys.argv +if ISOLATED: + sys.argv.remove('--isolated-twinkle') + from isolated_support import install + install(ROOT) + +from tts.qwen3.data import Qwen3TTSDataConfig, build_qwen3_sample, collate_qwen3_batch +from tts.qwen3.input_processor import Qwen3TTSInputProcessor +from tts.qwen3.forward import qwen3_tts_forward_step +from tts.qwen3.loss import Qwen3TTSLoss +from tts.qwen3.model import Qwen3TTSModel + +spec = importlib.util.spec_from_file_location('prepare', ROOT / 'scripts/audio_features.py') +prepare = importlib.util.module_from_spec(spec) +spec.loader.exec_module(prepare) + + +def samples(n=2): + cfg = Qwen3TTSDataConfig( + tts_pad_token_id=0, tts_bos_token_id=1, tts_eos_token_id=2, + codec_pad_id=3, codec_bos_id=4, codec_eos_token_id=5, + codec_nothink_id=6, codec_think_bos_id=7, codec_think_eos_id=8) + rows = [build_qwen3_sample(np.arange(10, 15-i), + np.arange((3-i)*16).reshape(3-i, 16) % 20 + 9, cfg) for i in range(n)] + mels = [np.full((10, 128), i+1, dtype=np.float32) for i in range(n)] + return cfg, rows, mels + + +class Talker(nn.Module): + """Small differentiable double of Qwen's two-stage API, with real HF loss.""" + def __init__(self): + super().__init__() + self.config = types.SimpleNamespace(num_code_groups=16) + self.model = nn.Module() + self.model.text_embedding = nn.Embedding(64, 8) + self.text_projection = nn.Identity() + self.model.codec_embedding = nn.Embedding(64, 8) + self.code_predictor = nn.Module() + self.code_predictor.embeddings = nn.ModuleList([nn.Embedding(64, 8) for _ in range(15)]) + self.code_predictor.get_input_embeddings = lambda: self.code_predictor.embeddings + self.head = nn.Linear(8, 64) + self.heads = nn.ModuleList([nn.Linear(8, 64) for _ in range(15)]) + + def forward(self, inputs_embeds, attention_mask, labels, **kwargs): + self.labels = labels + self.hidden = inputs_embeds + self.logits = self.head(inputs_embeds) + self.logits.retain_grad() + return types.SimpleNamespace( + loss=ForCausalLMLoss(self.logits, labels, 64), + hidden_states=((inputs_embeds,), None)) + + def forward_sub_talker_finetune(self, codes, hidden): + self.selected_hidden = hidden + self.selected_codes = codes + logits = torch.stack([head(hidden) for head in self.heads], dim=1) + logits.retain_grad() + self.sub_logits = logits + return logits, ForCausalLMLoss(logits, codes[:, 1:], 64) + + +class Model(nn.Module): + dtype = torch.float32 + device = torch.device('cpu') + def __init__(self): + super().__init__() + self.talker = Talker() + self.speaker_encoder = lambda mels: mels.mean(1)[:, :8] + + +class Contracts(unittest.TestCase): + def test_initialize_accepts_local_mode(self): + tree = ast.parse((ROOT / 'scripts/train_qwen3_tts_lora.py').read_text(encoding='utf-8')) + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) and node.func.attr == 'initialize'] + self.assertEqual(len(calls), 1) + seen = [] + def initialize(mode, global_device_mesh): + self.assertIn(mode, ('local', 'ray')) # pinned Twinkle's accepted contract + self.assertEqual(mode, 'local') + seen.append(global_device_mesh) + eval(compile(ast.Expression(calls[0]), '', 'eval'), + {'twinkle': types.SimpleNamespace(initialize=initialize), 'device_mesh': 'single'}) + self.assertEqual(seen, ['single']) + + def test_codec_single_list_tensor_and_batch_tensor(self): + expected = np.arange(48).reshape(3, 16) + for value in ([torch.tensor(expected)], [expected], torch.tensor(expected), + torch.tensor(expected)[None], expected): + tokenizer = types.SimpleNamespace(encode=lambda *a, **kw: types.SimpleNamespace(audio_codes=value)) + np.testing.assert_array_equal(prepare.encode_audio_codes(np.ones(4096), tokenizer), expected) + + def test_codec_rejects_empty_multiple_and_wrong_shapes(self): + for value in ([], [np.zeros((2,16))]*2, np.zeros((2,3,16)), np.zeros((2,15)), np.zeros((0,16))): + with self.subTest(shape=str(type(value))), self.assertRaises(ValueError): + tokenizer = types.SimpleNamespace(encode=lambda *a, **kw: types.SimpleNamespace(audio_codes=value)) + prepare.encode_audio_codes(np.ones(4096), tokenizer) + + def test_numpy_mel_batch_identity_and_singleton(self): + for n in (1,2): + cfg, rows, mels = samples(n) + for refs in (mels, [m[None] for m in mels]): + batch = collate_qwen3_batch(rows, cfg, refs) + self.assertEqual(batch['ref_mels'].shape, (n,10,128)) + for i in range(n): + np.testing.assert_array_equal(batch['ref_mels'][i], mels[i]) + + def test_mel_length_and_count_fail_explicitly(self): + cfg, rows, mels = samples() + for refs in ([mels[0], mels[1][:7]], mels[:1], [], [mels[0], np.zeros((10,80))]): + with self.assertRaises(ValueError): + collate_qwen3_batch(rows, cfg, refs) + feats = [dict(row, ref_mels=mel) for row, mel in zip(rows, refs)] + if len(feats) == 2: + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor().collate_fn(feats) + + def test_processor_entire_call_samples_and_collated(self): + for n in (1,2): + cfg, rows, mels = samples(n) + expected = collate_qwen3_batch(rows, cfg, mels) + inputs = [[dict(row, ref_mels=mel) for row,mel in zip(rows,mels)], expected, + {k:torch.as_tensor(v) for k,v in expected.items()}] + for value in inputs: + actual = Qwen3TTSInputProcessor()(value, model=Model()) + self.assertEqual(set(actual), set(expected)) + for key in expected: + np.testing.assert_array_equal(actual[key].numpy(), expected[key]) + self.assertEqual(actual['codec_mask'].dtype, torch.bool) + + def test_processor_rejects_missing_mel_bad_rank_and_packing(self): + cfg, rows, mels = samples() + batch = collate_qwen3_batch(rows, cfg, mels) + missing = dict(batch); del missing['ref_mels'] + for value in (missing, dict(batch, ref_mels=np.zeros((20,128)))): + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor()(value, model=Model()) + with self.assertRaises(ValueError): + Qwen3TTSInputProcessor(padding_free=True)(batch, model=Model()) + + def test_wrapper_forward_alignment_and_every_head_gradient(self): + for n in (1,2): + cfg, rows, mels = samples(n) + batch = collate_qwen3_batch(rows, cfg, mels) + model = Model() + status = types.SimpleNamespace() + config = types.SimpleNamespace(processor=Qwen3TTSInputProcessor(), train_status=status) + wrapper = types.SimpleNamespace( + model=model, optimizer_group={'default':config}, + _get_default_group=lambda:'default', _lazy_wrap_model=lambda:None, + _not_encoded=lambda x:False, sp_strategy=None, hf_config=None) + outputs = inspect.unwrap(Qwen3TTSModel.forward)(wrapper, inputs=batch) + t = model.talker + labels = torch.as_tensor(batch['codec_0_labels']) + self.assertTrue(torch.equal(t.labels, labels)) + expected = nn.functional.cross_entropy(t.logits[:,:-1].reshape(-1,64), + labels[:,1:].reshape(-1)) + torch.testing.assert_close(outputs['talker_loss'], expected) + mask = torch.as_tensor(batch['codec_mask']) + torch.testing.assert_close(t.selected_hidden, t.hidden[mask]) + codes = torch.as_tensor(batch['codec_ids'])[mask] + torch.testing.assert_close(outputs['sub_talker_labels'], codes[:,1:]) + loss = Qwen3TTSLoss()({}, outputs)['loss'] + self.assertTrue(torch.isfinite(loss)) + loss.backward() + for b,row in enumerate(rows): + targets = np.where(row['codec_0_labels'] != -100)[0] + self.assertGreater(t.logits.grad[b,targets[0]-1].abs().sum().item(), 0) + self.assertGreater(t.logits.grad[b,targets[-1]-1].abs().sum().item(), 0) # EOS + shifted_valid = nn.functional.pad(labels[:,1:] != -100, (0,1)) + self.assertEqual(t.logits.grad[~shifted_valid].abs().sum().item(), 0) + # Every head receives its own label, including cb1 and cb15. + want = t.sub_logits.detach().softmax(-1) + want.scatter_add_(-1, codes[:,1:,None], -torch.ones_like(codes[:,1:,None], dtype=want.dtype)) + want *= 0.3 / (codes.shape[0]*15) + torch.testing.assert_close(t.sub_logits.grad, want) + for head in t.heads: + self.assertTrue(torch.isfinite(head.weight.grad).all()) + self.assertGreater(head.weight.grad.abs().sum().item(), 0) + self.assertIs(status.outputs['talker_loss'], outputs['talker_loss']) + self.assertEqual(status.loss_value, 0) + + def test_talker_time_shift_independent_of_processor(self): + cfg, rows, mels = samples() + batch = {key:torch.as_tensor(value) for key,value in collate_qwen3_batch(rows,cfg).items()} + batch['ref_mels'] = torch.as_tensor(np.stack(mels)) + model = Model() + outputs = qwen3_tts_forward_step(model,batch) + # Independent oracle: position t predicts original label[t+1], including + # the first frame and EOS. This fails the old external+internal shift. + expected = nn.functional.cross_entropy(model.talker.logits[:,:-1].reshape(-1,64), + batch['codec_0_labels'][:,1:].reshape(-1)) + torch.testing.assert_close(outputs['talker_loss'], expected) + self.assertEqual(model.talker.logits.shape[:2],batch['codec_0_labels'].shape) + + def test_official_mel_call_contract(self): + fake = types.ModuleType('qwen_tts.core.models.modeling_qwen3_tts') + def mel(wave, **kw): + self.assertEqual(wave.shape, (1,4096)) + self.assertEqual(kw, dict(n_fft=1024,num_mels=128,sampling_rate=24000, + hop_size=256,win_size=1024,fmin=0,fmax=12000.0)) + return torch.ones(1,128,16) + fake.mel_spectrogram = mel + with patch.dict(sys.modules, {'qwen_tts.core.models.modeling_qwen3_tts':fake}): + self.assertEqual(prepare.extract_mel(np.ones(4096)).shape, (16,128)) + + @unittest.skipIf(ISOLATED, 'Official qwen_tts not installed in isolated CPU run') + def test_official_sub_talker_methods_and_last_head_gradient(self): + from qwen_tts.core.models.modeling_qwen3_tts import ( + Qwen3TTSTalkerForConditionalGeneration, + Qwen3TTSTalkerCodePredictorModelForConditionalGeneration, + ) + # Run the installed official assembly and head-selection methods. Only + # the expensive transformer decoder is replaced by a tiny causal sum. + embeddings = nn.ModuleList([nn.Embedding(64,8) for _ in range(15)]) + heads = nn.ModuleList([nn.Linear(8,64) for _ in range(15)]) + cp = types.SimpleNamespace( + config=types.SimpleNamespace(num_code_groups=16, vocab_size=64, + output_attentions=False, output_hidden_states=False), + small_to_mtp_projection=nn.Identity(), lm_head=heads, + get_input_embeddings=lambda:embeddings, loss_function=ForCausalLMLoss, + model=lambda inputs_embeds, **kw: types.SimpleNamespace(last_hidden_state=inputs_embeds.cumsum(1))) + cp.forward_finetune = types.MethodType( + Qwen3TTSTalkerCodePredictorModelForConditionalGeneration.forward_finetune, cp) + talker = types.SimpleNamespace( + config=types.SimpleNamespace(num_code_groups=16,hidden_size=8), + code_predictor=cp, get_input_embeddings=lambda:embeddings[0]) + codes = torch.arange(32).reshape(2,16) + hidden = torch.randn(2,8,requires_grad=True) + logits, _ = Qwen3TTSTalkerForConditionalGeneration.forward_sub_talker_finetune(talker,codes,hidden) + logits.retain_grad() + outputs = dict(talker_loss=hidden.sum()*0,sub_talker_logits=logits,sub_talker_labels=codes[:,1:]) + Qwen3TTSLoss()({},outputs)['loss'].backward() + expected = logits.detach().softmax(-1) + expected.scatter_add_(-1,codes[:,1:,None],-torch.ones(2,15,1)) + torch.testing.assert_close(logits.grad, expected*(0.3/30)) + for head in heads: + self.assertGreater(head.weight.grad.abs().sum().item(),0) + + @unittest.skipIf(ISOLATED, 'Official qwen_tts not installed in isolated CPU run') + def test_official_mel_numerical_parity(self): + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + wave = np.sin(np.arange(8192, dtype=np.float32)*0.03) + expected = mel_spectrogram(torch.from_numpy(wave)[None], n_fft=1024, + num_mels=128,sampling_rate=24000,hop_size=256,win_size=1024,fmin=0,fmax=12000).transpose(1,2)[0] + np.testing.assert_allclose(prepare.extract_mel(wave), expected.numpy(), rtol=1e-6, atol=1e-6) + + +if __name__ == '__main__': + print('CPU contracts / Twinkle infrastructure: ' + ('ISOLATED DOUBLES' if ISOLATED else 'INSTALLED LIBRARY')) + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/entrypoint/test_contracts.py b/cookbook/qwen3_tts/tests/entrypoint/test_contracts.py new file mode 100644 index 00000000..03914225 --- /dev/null +++ b/cookbook/qwen3_tts/tests/entrypoint/test_contracts.py @@ -0,0 +1,110 @@ +"""CPU manifest/cursor tests; no Twinkle doubles or model imports needed.""" +import importlib.util +import json +from pathlib import Path +import random +import subprocess +import sys +import tempfile +from types import SimpleNamespace +import unittest +import numpy as np + +ROOT = Path(__file__).resolve().parents[2] + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +m = load('manifest', ROOT / 'tts/qwen3/manifest.py') +cli = load('train_cli', ROOT / 'scripts/train_qwen3_tts_lora.py') + + +class ManifestContracts(unittest.TestCase): + def test_learning_rate_decay_is_explicit_and_validated(self): + required = ['--model_path', 'model', '--train_jsonl', 'train', '--output_dir', 'out'] + self.assertEqual(cli.parse_args(required).lr_decay, 1.0) + self.assertEqual(cli.parse_args(required + ['--lr_decay', '0.8']).lr_decay, 0.8) + for value in ('0', '-1', '1.1', 'nan', 'inf'): + with self.subTest(value=value), self.assertRaises(SystemExit): + cli.parse_args(required + ['--lr_decay', value]) + + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + np.save(self.root / 'mel.npy', np.zeros((3, 128), dtype=np.float32)) + self.rows = [dict(id=str(i), text=f'text {i}', audio_codes=[[i] * 16], + ref_mel_path='mel.npy') for i in range(2)] + + def write(self): + path = self.root / 'train.jsonl' + path.write_text('\n'.join(json.dumps(r) for r in self.rows), encoding='utf-8') + return path + + def test_fingerprint_tracks_content_not_location(self): + first = m.Manifest(self.write()) + self.assertEqual(len(first), 2) + self.rows[0]['ref_mel_path'] = str(self.root / 'mel.npy') + self.assertEqual(first.fingerprint, m.Manifest(self.write()).fingerprint) + self.rows[0]['text'] += ' changed' + self.assertNotEqual(first.fingerprint, m.Manifest(self.write()).fingerprint) + second = m.Manifest(self.write()).fingerprint + np.save(self.root / 'mel.npy', np.ones((3, 128), dtype=np.float32)) + self.assertNotEqual(second, m.Manifest(self.write()).fingerprint) + + def test_rejects_empty_duplicate_and_invalid_codes(self): + original = json.loads(json.dumps(self.rows)) + for rows in ([], [original[0], original[0]], + [{**original[0], 'audio_codes': [[2048] * 16]}], + [{**original[0], 'audio_codes': [[0.5] * 16]}], + [{**original[0], 'audio_codes': [[0] * 15]}]): + with self.subTest(rows=rows), self.assertRaises(ValueError): + self.rows = rows + m.Manifest(self.write()) + + def test_rejects_invalid_mel(self): + for mel in (np.zeros((2, 80)), np.full((2, 128), np.nan)): + np.save(self.root / 'mel.npy', mel) + with self.assertRaises(ValueError): + m.Manifest(self.write()) + + def test_shuffle_is_deterministic_and_does_not_consume_rng(self): + state = random.getstate() + plan = m.sample_plan(5, 3, 42, True) + self.assertEqual(state, random.getstate()) + self.assertEqual(plan, m.sample_plan(5, 3, 42, True)) + self.assertNotEqual(plan[0], plan[1]) + for order in plan: + self.assertEqual(sorted(order), list(range(5))) + self.assertEqual(m.sample_plan(2, 2, 42, False), [[0, 1], [0, 1]]) + + def test_tail_and_resume_cursor_do_not_drop_samples(self): + schedule = m.planned_windows(m.sample_plan(5, 2, 42, False), 2) + self.assertEqual([len(w) for w in schedule], [2, 2, 1, 2, 2, 1]) + p = SimpleNamespace(optimizer_steps=3, micro_steps=5, consumed_samples=5, pending_micro_steps=0) + m.validate_cursor(p, schedule) + self.assertEqual(schedule[p.optimizer_steps:], [[0, 1], [2, 3], [4]]) + for key, value in [('micro_steps', 4), ('consumed_samples', 6), + ('pending_micro_steps', 1), ('optimizer_steps', 7)]: + with self.subTest(key=key), self.assertRaises(ValueError): + m.validate_cursor(SimpleNamespace(**{**vars(p), key: value}), schedule) + + def test_cli_help_and_model_content_fingerprint(self): + result = subprocess.run([sys.executable, str(ROOT / 'scripts/train_qwen3_tts_lora.py'), + '--help'], capture_output=True, text=True, timeout=30) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn('--resume_from', result.stdout) + weights = self.root / 'model.safetensors' + weights.write_bytes(b'first mock weight content') + first = cli.model_fingerprint(self.root) + weights.write_bytes(b'second mock weight content') + self.assertNotEqual(first, cli.model_fingerprint(self.root)) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/entrypoint/test_preparation.py b/cookbook/qwen3_tts/tests/entrypoint/test_preparation.py new file mode 100644 index 00000000..4b4665dd --- /dev/null +++ b/cookbook/qwen3_tts/tests/entrypoint/test_preparation.py @@ -0,0 +1,88 @@ +"""Split preservation and provenance; codec/mel callbacks are explicit doubles.""" +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest +import numpy as np +import soundfile as sf + +ROOT = Path(__file__).resolve().parents[2] +spec = importlib.util.spec_from_file_location('prepare_split', ROOT / 'scripts/prepare_split_manifest.py') +prep = importlib.util.module_from_spec(spec) +spec.loader.exec_module(prep) + + +class Preparation(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + for i in range(2): + sf.write(self.root / f'{i}.wav', np.sin(np.arange(2400) * (i + 1) / 20) * .1, 24000) + self.rows = [dict(id=str(i), text=f'Text {i}', group_id=str(i), speaker_id='speaker', audio=f'{i}.wav') for i in range(2)] + self.write() + + def write(self): + for name, row in zip(('train', 'dev'), self.rows): + (self.root / f'{name}.jsonl').write_text(json.dumps(row) + '\n', encoding='utf-8') + + def audit(self, reference='0'): + return prep.audit_split(self.root / 'train.jsonl', self.root / 'dev.jsonl', reference) + + def test_preserves_text_split_and_train_reference(self): + before = {p.name: p.read_bytes() for p in self.root.iterdir()} + out = self.root / 'derived' + report = prep.prepare(self.root/'train.jsonl', self.root/'dev.jsonl', out, '0', + lambda _: np.zeros((2, 16), dtype=np.int64), lambda _: np.zeros((3, 128), dtype=np.float32)) + self.assertEqual(report['counts'], {'train': 1, 'dev': 1}) + for i, split in enumerate(('train', 'dev')): + row = json.loads((out / f'{split}.jsonl').read_text()) + self.assertEqual(row['text'], self.rows[i]['text']) + self.assertEqual(row['id'], str(i)) + self.assertTrue((out / row['ref_mel_path']).is_file()) + self.assertEqual(report['mapping'][i]['reference_id'], '0') + for name, data in before.items(): + self.assertEqual((self.root/name).read_bytes(), data) + with self.assertRaises(FileExistsError): + prep.prepare(self.root/'train.jsonl', self.root/'dev.jsonl', out, '0', None, None) + + def test_rejects_overlap_for_each_identity(self): + original = dict(self.rows[1]) + for key in ('id', 'group_id', 'text', 'audio'): + self.rows[1] = {**original, key: self.rows[0][key]} + self.write() + with self.subTest(key=key), self.assertRaises(ValueError): + self.audit() + + def test_normalized_text_overlap_is_rejected_without_rewriting(self): + self.rows[1]['text'] = ' TEXT 0 ' + self.write() + with self.assertRaises(ValueError): + self.audit() + + def test_dev_reference_and_empty_dev_are_rejected(self): + with self.assertRaises(ValueError): + self.audit('1') + (self.root/'dev.jsonl').write_text('') + with self.assertRaises(ValueError): + self.audit() + + def test_reencoded_duplicate_audio_is_rejected(self): + wave, sr = sf.read(self.root/'0.wav', dtype='float32') + sf.write(self.root/'1.wav', wave, sr, subtype='FLOAT') + with self.assertRaisesRegex(ValueError, 'pcm'): + self.audit() + + def test_stereo_and_multiple_speakers_are_rejected(self): + self.rows[1]['speaker_id'] = 'other' + self.write() + with self.assertRaises(ValueError): + self.audit() + sf.write(self.root/'0.wav', np.zeros((100, 2)), 24000) + with self.assertRaises(ValueError): + prep.read_audio(self.root/'0.wav') + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/entrypoint/test_prepared_training_npu.py b/cookbook/qwen3_tts/tests/entrypoint/test_prepared_training_npu.py new file mode 100644 index 00000000..519f2eb6 --- /dev/null +++ b/cookbook/qwen3_tts/tests/entrypoint/test_prepared_training_npu.py @@ -0,0 +1,69 @@ +"""Actual preparation CLI -> validated manifests -> one real NPU training update.""" +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest +import torch_npu # register the actual remote NPU runtime + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +from tts.qwen3.manifest import Manifest + + +class PreparedTraining(unittest.TestCase): + def test_public_split_preparation_and_train_only_consumption(self): + outputs = ROOT / 'outputs' + outputs.mkdir(exist_ok=True) + model = os.environ.get('MODEL_PATH', '/models/Qwen3-TTS-12Hz-1.7B-Base') + with tempfile.TemporaryDirectory(prefix='prepared-training-', dir=outputs) as folder: + root = Path(folder) + ids, source_hashes = {}, {} + for split, fixture in [('train', ROOT/'tests/fixtures/librispeech'), + ('dev', ROOT/'tests/fixtures/librispeech/second')]: + meta = json.loads((fixture/'sample.json').read_text()) + audio = fixture/meta['audio'] + self.assertEqual(hashlib.sha256(audio.read_bytes()).hexdigest(), meta['sha256']) + source_hashes[audio] = meta['sha256'] + ids[split] = meta['id'] + # Explicit fixture-only split, not a real held-out evaluation claim. + row = dict(id=meta['id'], text=meta['text'], audio=str(audio), + group_id=meta['id'], speaker_id='1272') + (root/f'{split}.jsonl').write_text(json.dumps(row)+'\n', encoding='utf-8') + derived = root/'derived' + subprocess.run([sys.executable, '-u', str(ROOT/'scripts/prepare_split_manifest.py'), + '--train_jsonl', str(root/'train.jsonl'), '--dev_jsonl', str(root/'dev.jsonl'), + '--reference_id', ids['train'], '--model_path', model, '--output_dir', str(derived)], + check=True, timeout=600) + for split in ('train', 'dev'): + data = Manifest(derived/f'{split}.jsonl') + self.assertEqual(len(data), 1) + self.assertEqual(data.rows[0]['id'], ids[split]) + provenance = json.loads((derived/'provenance.json').read_text()) + self.assertEqual(provenance['counts'], {'train': 1, 'dev': 1}) + self.assertTrue(all(r['reference_id'] == ids['train'] for r in provenance['mapping'])) + output = root/'training' + subprocess.run([sys.executable, '-u', str(ROOT/'scripts/train_qwen3_tts_lora.py'), + '--model_path', model, '--train_jsonl', str(derived/'train.jsonl'), + '--output_dir', str(output), '--epochs', '1', '--grad_accum', '2'], + check=True, timeout=600) + trace = [json.loads(line) for line in (output/'trace.jsonl').read_text().splitlines()] + self.assertEqual(len(trace), 1) + self.assertEqual(trace[0]['sample_ids'], [ids['train']]) + result = json.loads((output/'result.json').read_text()) + self.assertEqual((result['micro_steps'], result['optimizer_steps']), (1, 1)) + self.assertTrue(result['complete']) + self.assertTrue((Path(result['checkpoint'])/'adapter_model.safetensors').is_file()) + for audio, digest in source_hashes.items(): + self.assertEqual(hashlib.sha256(audio.read_bytes()).hexdigest(), digest) + print(json.dumps(dict(stage='预编码训练联通通过 / Prepared training passed', + counts=provenance['counts'], reference_id=ids['train'], consumed_ids=trace[0]['sample_ids'], + dev_consumed=False, loss=trace[0]['losses'], visible_devices=os.environ.get('ASCEND_RT_VISIBLE_DEVICES'), + scope='public regression fixtures only; no held-out quality evaluation'), ensure_ascii=False), flush=True) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/entrypoint/test_process_resume_npu.py b/cookbook/qwen3_tts/tests/entrypoint/test_process_resume_npu.py new file mode 100644 index 00000000..89dbcbcb --- /dev/null +++ b/cookbook/qwen3_tts/tests/entrypoint/test_process_resume_npu.py @@ -0,0 +1,111 @@ +"""Actual CLI in three sequential processes, two public samples, epoch tails.""" +import gc +import json +import os +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(ROOT / 'tests/training')) +import test_resume_npu as previous +import numpy as np +import torch +from safetensors.torch import load_file +from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSTokenizer +from tts.qwen3.manifest import sample_plan, planned_windows +from tts.qwen3.speech_fixture import load_speech_fixture +from tts.qwen3.runtime import DEFAULT_MODEL_PATH + + +def compare(test, a, b, exact=False): + if isinstance(a, np.ndarray): + np.testing.assert_array_equal(a, b) + elif isinstance(a, dict): + test.assertEqual(a.keys(), b.keys()) + for key in a: + compare(test, a[key], b[key], exact) + elif isinstance(a, (list, tuple)): + test.assertEqual(len(a), len(b)) + for x, y in zip(a, b): + compare(test, x, y, exact) + else: + previous.assert_tree_equal(test, a, b, rtol=0 if exact else 1e-3, + atol=0 if exact else 1e-5) + + +class ProcessResume(unittest.TestCase): + def test_cli_multi_sample_shuffle_and_restart(self): + path = Path(os.environ.get('MODEL_PATH', DEFAULT_MODEL_PATH)).resolve(strict=True) + outputs = ROOT / 'outputs' + outputs.mkdir(exist_ok=True) + with tempfile.TemporaryDirectory(prefix='process-resume-', dir=outputs) as folder: + root = Path(folder) + codec = Qwen3TTSTokenizer.from_pretrained(str(path / 'speech_tokenizer'), + local_files_only=True, device_map='cpu', torch_dtype=torch.float32) + codec.model.float().eval() + rows = [] + for index, fixture in enumerate(('tests/fixtures/librispeech', 'tests/fixtures/librispeech/second')): + wave, text, meta = load_speech_fixture(ROOT / fixture) + with torch.no_grad(): + codes = previous.baseline.prepare.encode_audio_codes(wave, codec) + np.save(root / f'mel-{index}.npy', previous.baseline.prepare.extract_mel(wave)) + rows.append(dict(id=meta['id'], text=text, audio_codes=codes.tolist(), + ref_mel_path=f'mel-{index}.npy')) + del codec + gc.collect() + manifest = root / 'train.jsonl' + manifest.write_text('\n'.join(json.dumps(r) for r in rows), encoding='utf-8') + common = [sys.executable, '-u', str(ROOT / 'scripts/train_qwen3_tts_lora.py'), + '--model_path', str(path), '--train_jsonl', str(manifest), + '--lr_decay', '0.8', '--epochs', '3', '--grad_accum', '3', '--shuffle', '--seed', '42'] + + def run(name, extra=()): + output = root / name + print(f'启动独立进程 / Starting independent process: {name}', flush=True) + subprocess.run([*common, '--output_dir', str(output), *extra], check=True, timeout=1800) + return output, json.loads((output / 'result.json').read_text()) + + continuous, a = run('continuous') + interrupted, b = run('interrupted', ['--stop_after_updates', '1']) + resumed, c = run('resumed', ['--resume_from', b['checkpoint']]) + self.assertEqual(len({a['pid'], b['pid'], c['pid']}), 3) + self.assertTrue(a['complete'] and c['complete']) + self.assertFalse(b['complete']) + self.assertEqual((b['micro_steps'], b['optimizer_steps']), (2, 1)) + self.assertEqual((c['micro_steps'], c['optimizer_steps'], c['consumed_samples']), (6, 3, 6)) + + def trace(path): + return [json.loads(line) for line in (path / 'trace.jsonl').read_text().splitlines()] + + full, split = trace(continuous), trace(interrupted) + trace(resumed) + expected = [[rows[i]['id'] for i in w] for w in planned_windows(sample_plan(2, 3, 42, True), 3)] + self.assertEqual([r['sample_ids'] for r in full], expected) + self.assertEqual(len(full), len(split)) + for x, y in zip(full, split): + np.testing.assert_allclose(x['losses'], y['losses'], rtol=1e-3, atol=1e-5) + self.assertEqual({k: v for k, v in x.items() if k != 'losses'}, + {k: v for k, v in y.items() if k != 'losses'}) + ac, cc = Path(a['checkpoint']), Path(c['checkpoint']) + wa, wc = load_file(str(ac / 'adapter_model.safetensors')), load_file(str(cc / 'adapter_model.safetensors')) + compare(self, wa, wc) + for filename in ('optimizer.pt', 'scheduler.pt', 'rng_state.pt'): + # Only checkpoint files generated by these child processes are loaded. + compare(self, torch.load(ac / filename, map_location='cpu', weights_only=False), + torch.load(cc / filename, map_location='cpu', weights_only=False), + exact=filename != 'optimizer.pt') + for filename in ('cookbook_progress.json', 'trainer_state.json'): + self.assertEqual(json.loads((ac / filename).read_text()), json.loads((cc / filename).read_text())) + previous.baseline.report('独立进程续训通过 / Independent process resume passed', + pids=[a['pid'], b['pid'], c['pid']], sample_order=expected, + micro_steps=6, optimizer_steps=3, epoch_window_sizes=[2, 2, 2], + max_weight_difference=max((wa[k].float() - wc[k].float()).abs().max().item() for k in wa), + random_states_equal=True, rtol=1e-3, atol=1e-5, + visible_devices=os.environ.get('ASCEND_RT_VISIBLE_DEVICES')) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/fixtures/librispeech/1272-128104-0000.flac b/cookbook/qwen3_tts/tests/fixtures/librispeech/1272-128104-0000.flac new file mode 100644 index 00000000..3b0c844f Binary files /dev/null and b/cookbook/qwen3_tts/tests/fixtures/librispeech/1272-128104-0000.flac differ diff --git a/cookbook/qwen3_tts/tests/fixtures/librispeech/ATTRIBUTION.md b/cookbook/qwen3_tts/tests/fixtures/librispeech/ATTRIBUTION.md new file mode 100644 index 00000000..ba9ee34d --- /dev/null +++ b/cookbook/qwen3_tts/tests/fixtures/librispeech/ATTRIBUTION.md @@ -0,0 +1,19 @@ +# Public speech regression fixture + +Utterance: LibriSpeech dev-clean **1272-128104-0000**. +Prepared by Vassil Panayotov with assistance from Daniel Povey; derived from LibriVox audiobooks. +Corpus: https://www.openslr.org/12 — **CC BY 4.0**: https://creativecommons.org/licenses/by/4.0/ + +Obtained from the Hugging Face internal testing mirror, clean/validation row 0: +https://huggingface.co/datasets/hf-internal-testing/librispeech_asr_dummy +Mirror revision: 5be91486e11a2d616f4ec5db8d3fd248585ac07a. +The checked-in FLAC is the viewer's audio export, with no local trimming or editing. +The transcript is preserved verbatim in sample.json. The test resamples to 24 kHz at runtime. + +This public **dev-clean** example is exclusively a regression fixture. Do not include it +in customer training/evaluation datasets or claim held-out evaluation after using it to update a model. +A matching corpus transcript is not a new manual transcription or an audio quality rating. + +## Additional process-resume fixture (2026-10-26) + +`second/1272-128104-0001.flac`: row 1 of the same dataset/revision/split, same attribution and CC BY 4.0 license. SHA256 `46a9d58622b4675b29564da2d9ba73e702241c5fa969f12c387cad4aa984276a`. No local trimming or text edits; runtime resamples 16 kHz to 24 kHz. Full transcript and provenance are in `second/sample.json`. This is a regression fixture, not held-out evaluation data. diff --git a/cookbook/qwen3_tts/tests/fixtures/librispeech/sample.json b/cookbook/qwen3_tts/tests/fixtures/librispeech/sample.json new file mode 100644 index 00000000..0034e925 --- /dev/null +++ b/cookbook/qwen3_tts/tests/fixtures/librispeech/sample.json @@ -0,0 +1,15 @@ +{ + "id": "1272-128104-0000", + "text": "MISTER QUILTER IS THE APOSTLE OF THE MIDDLE CLASSES AND WE ARE GLAD TO WELCOME HIS GOSPEL", + "audio": "1272-128104-0000.flac", + "sha256": "4e25e22555cd16e90edb0a3b49fdcf1fe652b2a1250ab643634db33895c75b41", + "source": "https://huggingface.co/datasets/hf-internal-testing/librispeech_asr_dummy", + "source_revision": "5be91486e11a2d616f4ec5db8d3fd248585ac07a", + "source_split": "clean/validation", + "source_row": 0, + "corpus": "LibriSpeech dev-clean", + "license": "CC-BY-4.0", + "license_url": "https://creativecommons.org/licenses/by/4.0/", + "corpus_url": "https://www.openslr.org/12", + "changes": "HF viewer audio export; no local trimming or text edits. Runtime resamples 16 kHz to 24 kHz." +} diff --git a/cookbook/qwen3_tts/tests/fixtures/librispeech/second/1272-128104-0001.flac b/cookbook/qwen3_tts/tests/fixtures/librispeech/second/1272-128104-0001.flac new file mode 100644 index 00000000..0c74e841 Binary files /dev/null and b/cookbook/qwen3_tts/tests/fixtures/librispeech/second/1272-128104-0001.flac differ diff --git a/cookbook/qwen3_tts/tests/fixtures/librispeech/second/sample.json b/cookbook/qwen3_tts/tests/fixtures/librispeech/second/sample.json new file mode 100644 index 00000000..01ffdc1f --- /dev/null +++ b/cookbook/qwen3_tts/tests/fixtures/librispeech/second/sample.json @@ -0,0 +1,15 @@ +{ + "id": "1272-128104-0001", + "text": "NOR IS MISTER QUILTER'S MANNER LESS INTERESTING THAN HIS MATTER", + "audio": "1272-128104-0001.flac", + "sha256": "46a9d58622b4675b29564da2d9ba73e702241c5fa969f12c387cad4aa984276a", + "source": "https://huggingface.co/datasets/hf-internal-testing/librispeech_asr_dummy", + "source_revision": "5be91486e11a2d616f4ec5db8d3fd248585ac07a", + "source_split": "clean/validation", + "source_row": 1, + "corpus": "LibriSpeech dev-clean", + "license": "CC-BY-4.0", + "license_url": "https://creativecommons.org/licenses/by/4.0/", + "corpus_url": "https://www.openslr.org/12", + "changes": "HF viewer audio export; no local trimming or text edits. Runtime resamples 16 kHz to 24 kHz." +} diff --git a/cookbook/qwen3_tts/tests/lora/test_contracts.py b/cookbook/qwen3_tts/tests/lora/test_contracts.py new file mode 100644 index 00000000..cf92c162 --- /dev/null +++ b/cookbook/qwen3_tts/tests/lora/test_contracts.py @@ -0,0 +1,127 @@ +"""CPU regressions for projection and acceptance assertions; Twinkle may be doubled.""" +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / 'tests/batch1')) +import test_contracts as previous +import torch +from torch import nn +from tts.qwen3.runtime import data_config, frozen_digest, trainable_snapshot, assert_updated, assert_gradients + + +class AcceptanceContracts(unittest.TestCase): + def test_new_optimizer_groups_receive_local_tokenizer_policy(self): + module = types.ModuleType('twinkle.model.transformers.transformers') + module.OptimizerGroup = types.SimpleNamespace + wrapper = types.SimpleNamespace(tokenizer_id='/local/model', device_mesh=None) + with patch.dict(sys.modules, {module.__name__: module}), \ + patch('tts.qwen3.tokenization.local_template') as template: + for _ in range(2): # initial group and newly added adapter group + group = previous.Qwen3TTSModel._construct_default_optimizer_group(wrapper) + self.assertIs(group.template, template.return_value) + self.assertIsInstance(group.processor, previous.Qwen3TTSInputProcessor) + self.assertIsInstance(group.loss_instance, previous.Qwen3TTSLoss) + self.assertEqual(template.call_count, 2) + template.assert_called_with('/local/model') + + def cpu_batch(self, model, cfg, rows, mels): + # These are CPU regressions even on an NPU host. Fail if the processor + # consults the platform default instead of the supplied model device. + with patch('tts.qwen3.input_processor.Platform.get_local_device', + side_effect=AssertionError('CPU fixture must use model.device')): + batch = previous.Qwen3TTSInputProcessor()( + previous.collate_qwen3_batch(rows, cfg, mels), model=model) + self.assertTrue(all(value.device.type == 'cpu' for value in batch.values())) + return batch + + def test_text_projection_handles_different_width_and_receives_gradient(self): + cfg, rows, mels = previous.samples(1) + model = previous.Model() + batch = self.cpu_batch(model, cfg, rows, mels) + model.talker.model.text_embedding = nn.Embedding(64, 12) + model.talker.text_projection = nn.Linear(12, 8) + seen = [] + handle = model.talker.text_projection.register_forward_hook(lambda m, a, o: seen.append(o.shape)) + outputs = previous.qwen3_tts_forward_step(model, batch) + previous.Qwen3TTSLoss()({}, outputs)['loss'].backward() + handle.remove() + self.assertEqual(len(seen), 1) + self.assertEqual(seen[0][-1], 8) + self.assertGreater(model.talker.text_projection.weight.grad.abs().sum().item(), 0) + + def test_speaker_encoder_has_no_gradient(self): + cfg, rows, mels = previous.samples(1) + model = previous.Model() + batch = self.cpu_batch(model, cfg, rows, mels) + # Preserve the speaker encoder's [B,F,128] -> [B,H] contract. + class Speaker(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(128, 8) + def forward(self, mel): + return self.linear(mel.mean(1)) + model.speaker_encoder = Speaker() + outputs = previous.qwen3_tts_forward_step(model, batch) + previous.Qwen3TTSLoss()({}, outputs)['loss'].backward() + self.assertIsNone(model.speaker_encoder.linear.weight.grad) + self.assertFalse(outputs['target_speaker_embedding'].requires_grad) + + def test_config_uses_all_real_token_ids_not_defaults(self): + names = ['codec_pad_id', 'codec_bos_id', 'codec_eos_token_id', + 'codec_nothink_id', 'codec_think_bos_id', 'codec_think_eos_id', 'num_code_groups'] + tc = types.SimpleNamespace(**dict(zip(names, range(10, 17)))) + config = types.SimpleNamespace(talker_config=tc, tts_pad_token_id=1, tts_bos_token_id=2, tts_eos_token_id=3) + actual = data_config(config) + self.assertEqual((actual.tts_pad_token_id, actual.tts_bos_token_id, actual.tts_eos_token_id), (1, 2, 3)) + for name in names: + self.assertEqual(getattr(actual, name), getattr(tc, name)) + + def tiny_model(self): + model = nn.Module() + model.base = nn.Module() + model.base.talker = nn.Module() + for branch in ('model', 'code_predictor'): + layer = nn.Module() + layer.lora_B = nn.ParameterDict({'smoke': nn.Parameter(torch.zeros(2, 2))}) + layer.weight = nn.Parameter(torch.ones(2, 2), requires_grad=False) + setattr(model.base.talker, branch, layer) + return model + + def test_assertions_reject_no_update_and_frozen_or_missing_branch(self): + model = self.tiny_model() + before = trainable_snapshot(model) + with self.assertRaisesRegex(AssertionError, 'No parameter update'): + assert_updated(before, trainable_snapshot(model)) + model.base.talker.model.weight.requires_grad_(True) + with self.assertRaisesRegex(AssertionError, 'Only LoRA'): + trainable_snapshot(model) + model.base.talker.model.weight.requires_grad_(False) + model.base.talker.code_predictor.lora_B['smoke'].requires_grad_(False) + with self.assertRaisesRegex(AssertionError, 'Missing trainable branch'): + trainable_snapshot(model) + + def test_gradients_updates_and_frozen_digest_detect_corruption(self): + model = self.tiny_model() + before = trainable_snapshot(model) + digest = frozen_digest(model) + with self.assertRaisesRegex(AssertionError, 'Missing/nonfinite'): + assert_gradients(model) + sum(p.sum() for p in model.parameters() if p.requires_grad).backward() + self.assertEqual(len(assert_gradients(model)), 2) + torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=0.1).step() + self.assertEqual(len(assert_updated(before, trainable_snapshot(model))), 2) + self.assertEqual(digest, frozen_digest(model)) + with torch.no_grad(): + model.base.talker.model.weight[0, 0] += 1 + self.assertNotEqual(digest, frozen_digest(model)) + model.base.talker.model.lora_B['smoke'].grad.fill_(float('nan')) + with self.assertRaisesRegex(AssertionError, 'Missing/nonfinite'): + assert_gradients(model) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/lora/test_real_model.py b/cookbook/qwen3_tts/tests/lora/test_real_model.py new file mode 100644 index 00000000..ffd7f289 --- /dev/null +++ b/cookbook/qwen3_tts/tests/lora/test_real_model.py @@ -0,0 +1,181 @@ +"""Real pretrained Qwen3-TTS + Twinkle + one NPU. No doubles or skip path. + +Uses the existing local 1.7B Base directory without downloads. The synthetic sine +wave/text pair only exercises mechanics; it is not speech/quality evidence. +""" +import gc +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import sys +import tempfile +import unittest + +import numpy as np +import torch +import torch_npu # noqa: F401 + +# Also enforce local-only model loading when this test is invoked directly. +os.environ['HF_HUB_OFFLINE'] = '1' +os.environ['TRANSFORMERS_OFFLINE'] = '1' +from peft import LoraConfig +from safetensors.torch import load_file +import twinkle +from twinkle import DeviceMesh + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) +from tts.qwen3.data import build_assistant_text, build_qwen3_sample, collate_qwen3_batch +from tts.qwen3.forward import qwen3_tts_forward_step +from tts.qwen3.input_processor import Qwen3TTSInputProcessor +from tts.qwen3.loss import Qwen3TTSLoss +from tts.qwen3.tokenization import load_text_tokenizer, assert_tokenizers_match +from tts.qwen3.runtime import (DEFAULT_MODEL_PATH, ADAPTER, data_config, + load_twinkle_model, trainable_snapshot, frozen_digest, assert_gradients, assert_updated) + +spec = importlib.util.spec_from_file_location('prepare', ROOT / 'scripts/audio_features.py') +prepare = importlib.util.module_from_spec(spec) +spec.loader.exec_module(prepare) + + +def report(stage, **values): + print(json.dumps({'stage': stage, **values}, ensure_ascii=False), flush=True) + + +def evaluation(model, batch): + model.eval() + with torch.no_grad(): + result = qwen3_tts_forward_step(model, batch) + return {key: result[key].detach().float().cpu() for key in + ('talker_loss', 'sub_talker_logits', 'target_speaker_embedding')} + + +class RealModelAcceptance(unittest.TestCase): + scope = 'one synthetic sample; not speech quality or optimizer resume' + + def sample(self): + t = np.arange(24000, dtype=np.float32) / 24000 + wave = (0.08 * np.sin(2 * np.pi * 220 * t) * np.sin(np.pi * t)**2).astype(np.float32) + return wave, '这是一条训练链路测试。', {'kind': 'synthetic sine wave'} + + def test_pretrained_single_sample_lora_update_and_fresh_reload(self): + torch.manual_seed(20261023) + np.random.seed(20261023) + torch.npu.set_device(0) + mesh = DeviceMesh.from_sizes(fsdp_size=1, dp_size=1) + twinkle.initialize(mode='local', global_device_mesh=mesh, seed=20261023) + path = Path(os.environ.get('MODEL_PATH', DEFAULT_MODEL_PATH)).expanduser().resolve(strict=True) + report('加载本地模型 / Load local model', path=str(path), + source_revision='unknown: user-provided local model', + config_sha256=hashlib.sha256((path / 'config.json').read_bytes()).hexdigest(), + visible_devices=os.environ.get('ASCEND_RT_VISIBLE_DEVICES', 'unset')) + model = load_twinkle_model(path, mesh) + cfg = data_config(model.hf_config) + self.assertEqual(model.hf_config.tts_model_type, 'base') + self.assertEqual(cfg.num_code_groups, 16) + self.assertEqual(model.model.dtype, torch.bfloat16) + + # Actual codec on CPU fp32; training and speaker encoder on NPU bf16. + wave, text, sample_metadata = self.sample() + report('样本来源 / Sample provenance', **sample_metadata) + codec = model.model.speech_tokenizer + codec.model.float().eval() + with torch.no_grad(): + codes = prepare.encode_audio_codes(wave, codec) + mel = prepare.extract_mel(wave) + model.model.speech_tokenizer = None + del codec + tokenizer = load_text_tokenizer(path) + probes = [text, '你好，世界！Hello 2026.\nA B', build_assistant_text(text)] + assert_tokenizers_match(tokenizer, + [model.optimizer_group[model._get_default_group()].template.tokenizer], probes) + text_ids = tokenizer.encode(build_assistant_text(text)) + # Verify the fixed suffix on the real tokenizer rather than blindly trim. + suffix = tokenizer.encode('<|im_end|>\n<|im_start|>assistant\n', add_special_tokens=False) + self.assertEqual(len(suffix), 5) + self.assertEqual(text_ids[-len(suffix):], suffix) + row = build_qwen3_sample(np.asarray(text_ids[:-len(suffix)]), codes, cfg) + raw_batch = collate_qwen3_batch([row], cfg, [mel]) + self.assertGreater(len(codes), 0) + self.assertGreaterEqual(int(codes.min()), 0) + self.assertLess(int(codes.max()), model.hf_config.talker_config.code_predictor_config.vocab_size) + report('样本已编码 / Sample encoded', audio_sha256=hashlib.sha256(wave.tobytes()).hexdigest(), + codec_shape=list(codes.shape), mel_shape=list(mel.shape), tokens=len(text_ids)) + + model.add_adapter_to_model(ADAPTER, LoraConfig( + r=8, lora_alpha=16, target_modules=['q_proj', 'v_proj'], + lora_dropout=0.0, bias='none'), gradient_accumulation_steps=1) + assert_tokenizers_match(tokenizer, [model.optimizer_group[ADAPTER].template.tokenizer], probes) + model.set_processor(Qwen3TTSInputProcessor, adapter_name=ADAPTER) + model.set_loss(Qwen3TTSLoss(), adapter_name=ADAPTER) + model.set_optimizer('AdamW', lr=1e-3, weight_decay=0.0, adapter_name=ADAPTER) + model._lazy_wrap_model() + base = model.strategy.unwrap_model(model.model) + self.assertEqual(base.device.type, 'npu') + self.assertEqual(base.dtype, torch.bfloat16) + before = trainable_snapshot(base) + frozen_before = frozen_digest(base) + batch = model.optimizer_group[ADAPTER].processor(raw_batch, model=base) + report('执行一次更新 / Run one update', trainable_parameters=sum(v.numel() for v in before.values()), + trainable_tensors=len(before), adapter_dtype=str(next(iter(before.values())).dtype)) + outputs = model.forward_backward(inputs=batch, adapter_name=ADAPTER) + loss = torch.as_tensor(outputs['loss']) + self.assertTrue(torch.isfinite(loss).all().item()) + nonzero = assert_gradients(base) + model.clip_grad_and_step(adapter_name=ADAPTER) + torch.npu.synchronize() + after = trainable_snapshot(base) + changed = assert_updated(before, after) + self.assertEqual(frozen_before, frozen_digest(base), 'Frozen base weights changed') + self.assertEqual(model.optimizer_group[ADAPTER].cur_step, 1) + reference = evaluation(base, batch) + for value in reference.values(): + self.assertTrue(torch.isfinite(value).all().item()) + report('更新断言通过 / Update assertions passed', loss=loss.detach().float().cpu().tolist(), + nonzero_gradient_tensors=len(nonzero), changed_tensors=len(changed), frozen_sha256=frozen_before) + + output_root = ROOT / 'outputs' + output_root.mkdir(exist_ok=True) + with tempfile.TemporaryDirectory(prefix='real-lora-', dir=output_root) as output: + checkpoint = model.save('adapter', output_dir=output, adapter_name=ADAPTER, save_optimizer=False) + assert_tokenizers_match(tokenizer, [load_text_tokenizer(checkpoint)], probes) + saved = load_file(str(Path(checkpoint) / 'adapter_model.safetensors')) + expected = {name.replace('.' + ADAPTER + '.', '.'): value for name, value in after.items()} + self.assertEqual(saved.keys(), expected.keys(), 'Checkpoint must contain only this adapter') + for name in saved: + torch.testing.assert_close(saved[name], expected[name], rtol=0, atol=0) + # Remove the original model and optimizer before loading another base. + del before, after, outputs, loss, base, model, batch + gc.collect() + torch.npu.empty_cache() + report('从新底模重载 / Reload into fresh pretrained base') + reloaded = load_twinkle_model(path, mesh) + reloaded.model.speech_tokenizer = None + reloaded.add_adapter_to_model(ADAPTER, str(checkpoint), is_trainable=True) + reloaded.set_processor(Qwen3TTSInputProcessor, adapter_name=ADAPTER) + assert_tokenizers_match(tokenizer, + [reloaded.optimizer_group[ADAPTER].template.tokenizer], probes) + report('文本编码一致 / Tokenizer IDs consistent', probes=len(probes), + fix_mistral_regex=True, paths=['text loader', 'Twinkle templates', 'saved tokenizer', 'fresh reload']) + reloaded._lazy_wrap_model() + fresh = reloaded.strategy.unwrap_model(reloaded.model) + actual_weights = trainable_snapshot(fresh) + normalized = {n.replace('.' + ADAPTER + '.', '.'): v for n, v in actual_weights.items()} + self.assertEqual(normalized.keys(), saved.keys()) + for name in saved: + torch.testing.assert_close(normalized[name], saved[name], rtol=0, atol=0) + self.assertEqual(frozen_before, frozen_digest(fresh)) + reloaded_batch = reloaded.optimizer_group[ADAPTER].processor(raw_batch, model=fresh) + actual = evaluation(fresh, reloaded_batch) + differences = {} + for key in reference: + torch.testing.assert_close(actual[key], reference[key], rtol=1e-3, atol=1e-3) + differences[key] = (actual[key] - reference[key]).abs().max().item() + report('保存重载通过 / Save and fresh reload passed', max_absolute_differences=differences, + scope=self.scope) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/run.py b/cookbook/qwen3_tts/tests/run.py new file mode 100644 index 00000000..81d75b53 --- /dev/null +++ b/cookbook/qwen3_tts/tests/run.py @@ -0,0 +1,24 @@ +"""Run each contract module in a fresh process; NPU tests are explicit opt-in.""" +import argparse +import os +from pathlib import Path +import subprocess +import sys + +ROOT = Path(__file__).resolve().parents[1] +parser = argparse.ArgumentParser(description=__doc__) +group = parser.add_mutually_exclusive_group() +group.add_argument('--isolated', action='store_true', help='CPU development with explicit Twinkle doubles') +group.add_argument('--npu', action='store_true', help='Also run real local-model integrations; set MODEL_PATH') +args = parser.parse_args() +if args.npu and not os.environ.get('MODEL_PATH'): + parser.error('--npu requires MODEL_PATH pointing to an existing local model') +modules = ['batch1/test_contracts.py', 'lora/test_contracts.py', 'training/test_contracts.py', + 'entrypoint/test_contracts.py', 'entrypoint/test_preparation.py'] +if args.npu: + modules += ['lora/test_real_model.py', 'entrypoint/test_prepared_training_npu.py', + 'entrypoint/test_process_resume_npu.py'] +for module in modules: + extra = ['--isolated-twinkle'] if args.isolated and module.startswith(('batch1/', 'lora/')) else [] + print('Running: ' + module, flush=True) + subprocess.run([sys.executable, '-u', str(ROOT/'tests'/module), *extra], check=True, timeout=3600) diff --git a/cookbook/qwen3_tts/tests/training/test_contracts.py b/cookbook/qwen3_tts/tests/training/test_contracts.py new file mode 100644 index 00000000..6e11e646 --- /dev/null +++ b/cookbook/qwen3_tts/tests/training/test_contracts.py @@ -0,0 +1,122 @@ +"""CPU loop tests with real autograd/SGD and an explicit Twinkle protocol double.""" +import importlib.util +import json +from pathlib import Path +import sys +import tempfile +import types +import unittest +import torch + +ROOT = Path(__file__).resolve().parents[2] +spec = importlib.util.spec_from_file_location('training', ROOT / 'tts/qwen3/training.py') +training = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = training +spec.loader.exec_module(training) + + +class Bridge: + def __init__(self): + self.model = torch.nn.Linear(1, 1, bias=False) + self.model.weight.data.fill_(1) + self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1) + self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, 1, gamma=0.8) + self.group = types.SimpleNamespace(gradient_accumulation_steps=1, cur_step=0, + train_status=types.SimpleNamespace(num_tokens=0)) + self.optimizer_group = {'test': self.group} + self.updates = [] + + def forward_backward(self, inputs, adapter_name): + loss = (self.model.weight - inputs['target']).square().mean() + loss.backward() + self.group.cur_step += 1 + self.group.train_status.num_tokens += 1 + return {'loss': loss.item()} + + def clip_grad_and_step(self, adapter_name): + count = self.group.train_status.num_tokens + for p in self.model.parameters(): + p.grad.div_(count) + self.optimizer.step() + self.optimizer.zero_grad(set_to_none=True) + self.scheduler.step() + self.updates.append((self.group.cur_step, count)) + self.group.train_status.num_tokens = 0 + + +def batch(target): + return {'input_ids': torch.zeros(1, 4, 2), 'target': float(target)} + + +class TrainingContracts(unittest.TestCase): + def test_tail_matches_independent_mean_gradient_reference(self): + model = Bridge() + progress = training.Progress(2) + targets = [2, -1, 3, 4, -2] + reference = torch.tensor(1., requires_grad=True) + optimizer = torch.optim.SGD([reference], lr=0.1) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.8) + for chunk in [targets[:2], targets[2:4], targets[4:]]: + ((reference - torch.tensor(chunk)).square().mean()).backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + scheduler.step() + for chunk in training.windows([batch(t) for t in targets], 2): + training.train_window(model, chunk, progress, 'test') + torch.testing.assert_close(model.model.weight.squeeze(), reference) + self.assertEqual(model.updates, [(2, 2), (4, 2), (5, 1)]) + self.assertEqual((progress.micro_steps, progress.optimizer_steps, progress.consumed_samples), (5, 3, 5)) + self.assertEqual(progress.pending_micro_steps, 0) + self.assertEqual(model.scheduler.last_epoch, 3) + + def test_invalid_batch_is_rejected_before_any_backward(self): + model = Bridge() + bad = batch(1) + bad['input_ids'] = torch.zeros(2, 4, 2) + with self.assertRaisesRegex(ValueError, 'single-sample'): + training.train_window(model, [batch(1), bad], training.Progress(2), 'test') + self.assertEqual(model.group.cur_step, 0) + self.assertIsNone(model.model.weight.grad) + + def test_empty_windows_and_invalid_accumulation_do_not_step(self): + self.assertEqual(list(training.windows([], 2)), []) + with self.assertRaises(ValueError): + training.Progress(0) + with self.assertRaises(ValueError): + training.train_window(Bridge(), [], training.Progress(2), 'test') + + def test_pending_gradients_cannot_be_checkpointed(self): + model = Bridge() + progress = training.Progress(2, pending_micro_steps=1) + with self.assertRaisesRegex(ValueError, 'unstepped'): + training.save_boundary(model, progress, '.', 'unused', 'test', 'sample') + progress.pending_micro_steps = 0 + model.model.weight.grad = torch.ones_like(model.model.weight) + with self.assertRaisesRegex(ValueError, 'cleared'): + training.save_boundary(model, progress, '.', 'unused', 'test', 'sample') + + def test_missing_resume_files_never_silently_restart(self): + with tempfile.TemporaryDirectory() as temp: + with self.assertRaises(FileNotFoundError): + training.resume_boundary(Bridge(), temp, 'test', 'sample', 2) + + def test_mismatched_resume_metadata_is_rejected_before_loading(self): + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) + for name in ('adapter_config.json', 'adapter_model.safetensors', 'optimizer.pt', + 'scheduler.pt', 'rng_state.pt', 'trainer_state.json'): + (root / name).write_text('{}') + state = {'format_version': 1, 'sample_fingerprint': 'other', + 'progress': dict(accumulation_steps=2, micro_steps=2, optimizer_steps=1, + consumed_samples=2, pending_micro_steps=0)} + (root / 'cookbook_progress.json').write_text(json.dumps(state)) + with self.assertRaisesRegex(ValueError, 'sample mismatch'): + training.resume_boundary(Bridge(), root, 'test', 'sample', 2) + state['sample_fingerprint'] = 'sample' + (root / 'cookbook_progress.json').write_text(json.dumps(state)) + with self.assertRaisesRegex(ValueError, 'accumulation mismatch'): + training.resume_boundary(Bridge(), root, 'test', 'sample', 3) + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tests/training/test_resume_npu.py b/cookbook/qwen3_tts/tests/training/test_resume_npu.py new file mode 100644 index 00000000..f3439ead --- /dev/null +++ b/cookbook/qwen3_tts/tests/training/test_resume_npu.py @@ -0,0 +1,177 @@ +"""Five paired-speech micro-batches: continuous versus fresh-model resume.""" +import copy +import gc +import hashlib +import json +import os +from pathlib import Path +import random +import sys +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / 'tests/lora')) +import test_real_model as baseline # local-only imports and shared fixture utilities +import numpy as np +import torch +from peft import LoraConfig +from tts.qwen3.runtime import (ADAPTER, DEFAULT_MODEL_PATH, load_twinkle_model, data_config, + trainable_snapshot, frozen_digest, assert_updated) +from tts.qwen3.data import build_assistant_text, build_qwen3_sample, collate_qwen3_batch +from tts.qwen3.tokenization import load_text_tokenizer +from tts.qwen3.speech_fixture import load_speech_fixture +from tts.qwen3.input_processor import Qwen3TTSInputProcessor +from tts.qwen3.loss import Qwen3TTSLoss +from tts.qwen3.training import Progress, windows, train_window, save_boundary, resume_boundary + + +def cpu_tree(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().clone() + if isinstance(value, dict): + return {k: cpu_tree(v) for k, v in value.items()} + if isinstance(value, (tuple, list)): + return type(value)(cpu_tree(v) for v in value) + return copy.deepcopy(value) + + +def assert_tree_equal(test, a, b, rtol=1e-3, atol=1e-5): + if isinstance(a, torch.Tensor): + torch.testing.assert_close(a, b, rtol=rtol, atol=atol) + elif isinstance(a, dict): + test.assertEqual(a.keys(), b.keys()) + for key in a: + assert_tree_equal(test, a[key], b[key], rtol, atol) + elif isinstance(a, (tuple, list)): + test.assertEqual(len(a), len(b)) + for x, y in zip(a, b): + assert_tree_equal(test, x, y, rtol, atol) + else: + test.assertEqual(a, b) + + +def configure(model): + model.add_adapter_to_model(ADAPTER, LoraConfig(r=8, lora_alpha=16, + target_modules=['q_proj', 'v_proj'], lora_dropout=0.1, bias='none'), + gradient_accumulation_steps=1) + model.set_processor(Qwen3TTSInputProcessor, adapter_name=ADAPTER) + model.set_loss(Qwen3TTSLoss(), adapter_name=ADAPTER) + model.set_optimizer('AdamW', lr=1e-3, weight_decay=0.0, adapter_name=ADAPTER) + model.set_lr_scheduler('StepLR', step_size=1, gamma=0.8, adapter_name=ADAPTER) + model._lazy_wrap_model() + + +def rng_probe(): + return [random.random(), float(np.random.rand()), torch.rand(1).item(), + torch.rand(1, device='npu:0').cpu().item()] + + +class ResumeAcceptance(unittest.TestCase): + def test_five_micro_batches_tail_and_fresh_optimizer_resume(self): + torch.npu.set_device(0) + mesh = baseline.DeviceMesh.from_sizes(fsdp_size=1, dp_size=1) + baseline.twinkle.initialize(mode='local', global_device_mesh=mesh, seed=20261023) + path = Path(os.environ.get('MODEL_PATH', DEFAULT_MODEL_PATH)).resolve(strict=True) + model = load_twinkle_model(path, mesh) + wave, text, meta = load_speech_fixture(ROOT / 'tests/fixtures/librispeech') + codec = model.model.speech_tokenizer + codec.model.float().eval() + with torch.no_grad(): + codes = baseline.prepare.encode_audio_codes(wave, codec) + mel = baseline.prepare.extract_mel(wave) + model.model.speech_tokenizer = None + del codec + tokenizer = load_text_tokenizer(path) + ids = tokenizer.encode(build_assistant_text(text)) + suffix = tokenizer.encode('<|im_end|>\n<|im_start|>assistant\n', add_special_tokens=False) + self.assertEqual(ids[-len(suffix):], suffix) + cfg = data_config(model.hf_config) + raw = collate_qwen3_batch([build_qwen3_sample(np.asarray(ids[:-len(suffix)]), codes, cfg)], cfg, [mel]) + fingerprint = hashlib.sha256(json.dumps({'sample': meta['sha256'], 'text_ids': ids, + 'model_config': hashlib.sha256((path / 'config.json').read_bytes()).hexdigest(), + 'schedule': 'same public sample repeated five times; accumulation=2'}, sort_keys=True).encode()).hexdigest() + configure(model) + base = model.strategy.unwrap_model(model.model) + self.assertEqual(base.device.type, 'npu') + self.assertEqual(base.dtype, torch.bfloat16) + initial = trainable_snapshot(base) + frozen = frozen_digest(base) + group = model.optimizer_group[ADAPTER] + batch = group.processor(raw, model=base) + progress = Progress(accumulation_steps=2) + train_window(model, [batch, batch], progress, ADAPTER) + self.assertEqual((progress.micro_steps, progress.optimizer_steps, progress.consumed_samples), (2, 1, 2)) + at_boundary = trainable_snapshot(base) + assert_updated(initial, at_boundary) + optimizer_boundary = cpu_tree(group.optimizer.state_dict()) + scheduler_boundary = copy.deepcopy(group.lr_scheduler.state_dict()) + output_root = ROOT / 'outputs' + output_root.mkdir(exist_ok=True) + with tempfile.TemporaryDirectory(prefix='resume-', dir=output_root) as temp: + checkpoint = save_boundary(model, progress, temp, 'after-update-1', ADAPTER, fingerprint) + baseline.report('断点已保存 / Boundary saved', micro_steps=2, optimizer_steps=1, + accumulation_steps=2, dropout=0.1, visible_devices=os.environ.get('ASCEND_RT_VISIBLE_DEVICES')) + trace = [] + losses = [] + sizes = [] + for window in windows([batch] * 3, 2): + sizes.append(len(window)) + losses.extend(train_window(model, window, progress, ADAPTER, + before_micro=lambda step: trace.append((step, rng_probe())))) + expected_weights = trainable_snapshot(base) + expected_optimizer = cpu_tree(group.optimizer.state_dict()) + expected_scheduler = copy.deepcopy(group.lr_scheduler.state_dict()) + expected_progress = copy.deepcopy(progress) + self.assertEqual(sizes, [2, 1]) + self.assertEqual((progress.micro_steps, progress.optimizer_steps, progress.consumed_samples), (5, 3, 5)) + self.assertEqual(group.cur_step, 5) + self.assertEqual(group.lr_scheduler.last_epoch, 3) + self.assertEqual(group.train_status.num_tokens, 0) + self.assertTrue(all(p.grad is None for p in base.parameters())) + self.assertEqual(frozen, frozen_digest(base)) + assert_updated(at_boundary, expected_weights) + self.assertEqual(len(expected_optimizer['state']), len(initial)) + self.assertTrue(all(int(state['step'].item()) == 3 for state in expected_optimizer['state'].values())) + baseline.report('连续训练完成 / Continuous run completed', micro_steps=5, optimizer_steps=3, + window_sizes=[2, 2, 1], learning_rates=group.lr_scheduler.get_last_lr()) + del model, base, group, batch + gc.collect() + torch.npu.empty_cache() + + resumed = load_twinkle_model(path, mesh) + resumed.model.speech_tokenizer = None + configure(resumed) + fresh = resumed.strategy.unwrap_model(resumed.model) + new_group = resumed.optimizer_group[ADAPTER] + new_batch = new_group.processor(raw, model=fresh) + restored = resume_boundary(resumed, checkpoint, ADAPTER, fingerprint, 2) + restored_weights = trainable_snapshot(fresh) + for name in at_boundary: + torch.testing.assert_close(restored_weights[name], at_boundary[name], rtol=0, atol=0) + assert_tree_equal(self, optimizer_boundary, cpu_tree(new_group.optimizer.state_dict()), rtol=0, atol=0) + self.assertEqual(scheduler_boundary, new_group.lr_scheduler.state_dict()) + actual_trace = [] + actual_losses = [] + for window in windows([new_batch] * (5 - restored.consumed_samples), 2): + actual_losses.extend(train_window(resumed, window, restored, ADAPTER, + before_micro=lambda step: actual_trace.append((step, rng_probe())))) + actual_weights = trainable_snapshot(fresh) + assert_tree_equal(self, expected_weights, actual_weights) + assert_tree_equal(self, expected_optimizer, cpu_tree(new_group.optimizer.state_dict())) + self.assertEqual(expected_scheduler, new_group.lr_scheduler.state_dict()) + self.assertEqual(expected_progress, restored) + self.assertEqual(trace, actual_trace, 'Python/NumPy/CPU/NPU random streams differ') + np.testing.assert_allclose(losses, actual_losses, rtol=1e-3, atol=1e-5) + self.assertEqual(frozen, frozen_digest(fresh)) + max_difference = max((actual_weights[n].float() - expected_weights[n].float()).abs().max().item() + for n in expected_weights) + baseline.report('断点续训对照通过 / Resume comparison passed', micro_steps=restored.micro_steps, + optimizer_steps=restored.optimizer_steps, consumed_samples=restored.consumed_samples, + max_weight_difference=max_difference, random_streams_equal=True, + losses_continuous=losses, losses_resumed=actual_losses, rtol=1e-3, atol=1e-5, + scope='five repeated public-speech micro-batches; boundary-only resume, single NPU') + + +if __name__ == '__main__': + unittest.main(verbosity=2) diff --git a/cookbook/qwen3_tts/tts/__init__.py b/cookbook/qwen3_tts/tts/__init__.py new file mode 100644 index 00000000..28cc1f3b --- /dev/null +++ b/cookbook/qwen3_tts/tts/__init__.py @@ -0,0 +1 @@ +"""Local Qwen3-TTS cookbook support; no global Twinkle registration.""" diff --git a/cookbook/qwen3_tts/tts/common/__init__.py b/cookbook/qwen3_tts/tts/common/__init__.py new file mode 100644 index 00000000..00ea8e69 --- /dev/null +++ b/cookbook/qwen3_tts/tts/common/__init__.py @@ -0,0 +1 @@ +from .multi_codebook_loss import MultiCodebookLoss diff --git a/cookbook/qwen3_tts/tts/common/multi_codebook_loss.py b/cookbook/qwen3_tts/tts/common/multi_codebook_loss.py new file mode 100644 index 00000000..be7166f3 --- /dev/null +++ b/cookbook/qwen3_tts/tts/common/multi_codebook_loss.py @@ -0,0 +1,31 @@ +"""多码本 loss 契约基类。各模型负责区分时间轴 next-token 与码本轴同位监督。""" + +from typing import Optional, Sequence + +from twinkle.data_format.output import LossOutput +from twinkle.loss.base import Loss + + +class MultiCodebookLoss(Loss): + """多码本 next-token CE 基类。 + + Args: + codebook_weights: 各码本 CE 权重，None=等权（Qwen3 的残差段权重 0.3 在子类里体现）。 + ar_codebook: AR 主码本 index（该码本由主模型 next-token 预测，走标准 shift+CE）。 + """ + + require_logps = False + require_logits = True + + def __init__( + self, + codebook_weights: Optional[Sequence[float]] = None, + ar_codebook: int = 0, + ): + self.codebook_weights = codebook_weights + self.ar_codebook = ar_codebook + + def __call__(self, inputs, outputs, **kwargs) -> LossOutput: + raise NotImplementedError( + '子类实现：按模型多码本组织算 CE。Qwen3 见 tts.qwen3.loss.Qwen3TTSLoss。' + ) diff --git a/cookbook/qwen3_tts/tts/qwen3/__init__.py b/cookbook/qwen3_tts/tts/qwen3/__init__.py new file mode 100644 index 00000000..c2e53cdf --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/__init__.py @@ -0,0 +1 @@ +"""Qwen3-TTS single-card training components.""" diff --git a/cookbook/qwen3_tts/tts/qwen3/data.py b/cookbook/qwen3_tts/tts/qwen3/data.py new file mode 100644 index 00000000..c5023483 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/data.py @@ -0,0 +1,190 @@ +"""Qwen3-TTS 双通道样本构建与 batch 组织。特殊 token 应从目标模型 config 读取。""" + +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +import numpy as np + +NUM_CODE_GROUPS = 16 # Qwen3-TTS 12Hz codec. + + +@dataclass +class Qwen3TTSDataConfig: + """Token layout for the 12Hz model. + + The loader supplies actual IDs from the local model configuration; defaults + document the 1.7B Base layout and should not be assumed for other variants. + """ + + # 文本通道（Qwen3 文本词表段，151xxx，config 顶层 tts_*_token_id） + tts_pad_token_id: int = 151671 + tts_bos_token_id: int = 151672 + tts_eos_token_id: int = 151673 + # codec 通道（codec 词表段，2xxx，config.talker_config.codec_*_id） + codec_pad_id: int = 2148 + codec_bos_id: int = 2149 + codec_eos_token_id: int = 2150 + codec_nothink_id: int = 2155 + codec_think_bos_id: int = 2156 + codec_think_eos_id: int = 2157 + num_code_groups: int = NUM_CODE_GROUPS + + +def build_assistant_text(text: str) -> str: + """官方 `_build_assistant_text`：把目标文本包成 assistant 轮（ChatML）。 + + 真实 text_ids 由 Qwen3 文本 tokenizer 对该串编码、再 `[:, :-5]` 得到。 + 本模块只负责布局，text_ids 由上游（Preprocessor/Template）传入。 + """ + return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + + +def build_qwen3_sample( + text_ids: np.ndarray, + audio_codes: np.ndarray, + cfg: Qwen3TTSDataConfig, +) -> Dict[str, np.ndarray]: + """把单条 (text_ids, audio_codes) 组成双通道 InputFeature（未 pad）。 + + Args: + text_ids: 1D int array [L]，L>=4（官方 `[:, :-5]` 后的文本 token，含 assistant 标签）。 + audio_codes: 2D int array [C, num_code_groups]，目标音频多码本。 + cfg: 特殊 token id。 + + Returns: + dict，每项是长度 T=8+L+C 的 numpy 数组（input_ids/codec_ids 是 2D）。 + """ + text_ids = np.asarray(text_ids).reshape(-1) + audio_codes = np.asarray(audio_codes) + assert audio_codes.ndim == 2 and audio_codes.shape[1] == cfg.num_code_groups, ( + f"audio_codes shape 应为 [C, {cfg.num_code_groups}]，实得 {audio_codes.shape}" + ) + L = int(text_ids.shape[0]) + C = int(audio_codes.shape[0]) + assert L >= 4, f"text_ids_len L={L} 过短（需 >=4 以满足 [3:] 非空）" + + T = 8 + L + C + audio_codec_0 = audio_codes[:, 0] # [C] + ncb = cfg.num_code_groups + + input_ids = np.zeros((T, 2), dtype=np.int64) + codec_ids = np.zeros((T, ncb), dtype=np.int64) + codec_0_labels = np.full((T,), -100, dtype=np.int64) + text_embedding_mask = np.zeros((T,), dtype=bool) + codec_embedding_mask = np.zeros((T,), dtype=bool) + codec_mask = np.zeros((T,), dtype=bool) + attention_mask = np.zeros((T,), dtype=np.int64) + + # ---- 文本通道 input_ids[:,0] ---- + input_ids[0:3, 0] = text_ids[0:3] + input_ids[3:7, 0] = cfg.tts_pad_token_id + input_ids[7, 0] = cfg.tts_bos_token_id + input_ids[8 : 8 + L - 3, 0] = text_ids[3:] # L-3 个 + input_ids[8 + L - 3, 0] = cfg.tts_eos_token_id + input_ids[8 + L - 2 : 8 + L + C, 0] = cfg.tts_pad_token_id # [8+L-2 : T] + + # ---- codec 通道 input_ids[:,1] ---- + input_ids[3:8, 1] = [ + cfg.codec_nothink_id, + cfg.codec_think_bos_id, + cfg.codec_think_eos_id, + 0, # speaker emb 槽（id=0 占位，训练步注入 speaker_emb） + cfg.codec_pad_id, + ] + input_ids[8 : 8 + L - 3, 1] = cfg.codec_pad_id + input_ids[8 + L - 3, 1] = cfg.codec_pad_id + input_ids[8 + L - 2, 1] = cfg.codec_bos_id + input_ids[8 + L - 1 : 8 + L - 1 + C, 1] = audio_codec_0 # [8+L-1 : 8+L-1+C] + input_ids[8 + L - 1 + C, 1] = cfg.codec_eos_token_id # pos T-1 + + # ---- codec_0_labels（码本0 label，-100 mask） ---- + codec_0_labels[8 + L - 1 : 8 + L - 1 + C] = audio_codec_0 + codec_0_labels[8 + L - 1 + C] = cfg.codec_eos_token_id # pos T-1 + + # ---- codec_ids（全码本，喂 Code predictor） ---- + codec_ids[8 + L - 1 : 8 + L - 1 + C, :] = audio_codes + + # ---- masks ---- + text_embedding_mask[: 8 + L + C] = True # [:T] + codec_embedding_mask[3 : 8 + L + C] = True # [3:T] + codec_embedding_mask[6] = False # speaker emb 槽 + codec_mask[8 + L - 1 : 8 + L - 1 + C] = True # [8+L-1 : 8+L+C-1) = C 个目标帧 + attention_mask[: 8 + L + C] = True # [:T] + + return { + 'input_ids': input_ids, + 'codec_ids': codec_ids, + 'codec_0_labels': codec_0_labels, + 'text_embedding_mask': text_embedding_mask, + 'codec_embedding_mask': codec_embedding_mask, + 'codec_mask': codec_mask, + 'attention_mask': attention_mask, + 'length': np.int64(T), + } + + +def collate_qwen3_batch( + samples: List[Dict[str, np.ndarray]], + cfg: Qwen3TTSDataConfig, + ref_mels: Optional[List[np.ndarray]] = None, + max_length: Optional[int] = None, +) -> Dict[str, np.ndarray]: + """把多个 build_qwen3_sample 输出 pad 成一个 batch（对齐官方 collate_fn 的 padding）。 + + padding 规则（与官方一致）： + input_ids/codec_ids -> 0；masks -> False；attention_mask -> 0；codec_0_labels -> -100。 + ref_mels 单条为 [F,128] 或 [1,F,128]，同长度 stack 为 [B,F,128]。 + 不同参考长度拒绝合批；使用相同参考音频或按参考长度分批，避免无掩码补零影响池化。 + """ + b = len(samples) + if max_length is None: + max_length = int(max(int(s['length']) for s in samples)) + t = max_length + ncb = cfg.num_code_groups + + input_ids = np.zeros((b, t, 2), dtype=np.int64) + codec_ids = np.zeros((b, t, ncb), dtype=np.int64) + codec_0_labels = np.full((b, t), -100, dtype=np.int64) + text_embedding_mask = np.zeros((b, t), dtype=bool) + codec_embedding_mask = np.zeros((b, t), dtype=bool) + codec_mask = np.zeros((b, t), dtype=bool) + attention_mask = np.zeros((b, t), dtype=np.int64) + + for i, s in enumerate(samples): + Ti = int(s['length']) + input_ids[i, :Ti, :] = s['input_ids'] + codec_ids[i, :Ti, :] = s['codec_ids'] + codec_0_labels[i, :Ti] = s['codec_0_labels'] + text_embedding_mask[i, :Ti] = s['text_embedding_mask'] + codec_embedding_mask[i, :Ti] = s['codec_embedding_mask'] + codec_mask[i, :Ti] = s['codec_mask'] + attention_mask[i, :Ti] = s['attention_mask'] + + batch = { + 'input_ids': input_ids, + 'codec_ids': codec_ids, + 'codec_0_labels': codec_0_labels, + 'text_embedding_mask': text_embedding_mask[:, :, None], # [b,t,1] 对齐官方 + 'codec_embedding_mask': codec_embedding_mask[:, :, None], # [b,t,1] + 'codec_mask': codec_mask, + 'attention_mask': attention_mask, + } + if ref_mels is not None: + if len(ref_mels) != b: + raise ValueError('Each sample requires one ref_mel') + batch['ref_mels'] = np.stack(normalize_ref_mels(ref_mels), axis=0) + return batch + + +def normalize_ref_mels(ref_mels): + """Validate per-sample mel shapes without padding or merging speakers.""" + normalized = [] + for mel in ref_mels: + if mel.ndim == 3 and mel.shape[0] == 1: + mel = mel[0] + if mel.ndim != 2 or mel.shape[0] == 0 or mel.shape[1] != 128: + raise ValueError('Each ref_mel must have shape [F,128] or [1,F,128], F > 0') + normalized.append(mel) + if not normalized or any(m.shape != normalized[0].shape for m in normalized): + raise ValueError('ref_mels must have equal lengths; use a shared reference or length buckets') + return normalized diff --git a/cookbook/qwen3_tts/tts/qwen3/forward.py b/cookbook/qwen3_tts/tts/qwen3/forward.py new file mode 100644 index 00000000..58d1b1f6 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/forward.py @@ -0,0 +1,88 @@ +"""Qwen3-TTS two-stage forward. Talker 使用内部时间 shift；残差码本 logits 交自定义同位 CE。 +小型替身可测试调用契约；真实模型/NPU/LoRA 仍需单独验收。""" + +from typing import Dict, Optional +import torch + + +def qwen3_tts_forward_step( + model, + batch: Dict[str, torch.Tensor], + target_speaker_embedding: Optional[torch.Tensor] = None, +): + """Qwen3-TTS 训练一步前向（产出中间态 dict，供 Qwen3TTSLoss 算总 loss）。 + + Args: + model: `Qwen3TTSForConditionalGeneration`（底层，= `Qwen3TTSModel.model`），需含顶层 + `.speaker_encoder` + `.talker`（含 `.model.codec_embedding`、`.model.text_embedding`、 + `.text_projection`、`.codec_head`、`.code_predictor`、`.forward_sub_talker_finetune`）。 + batch: collate_qwen3_batch 输出转 torch 后的张量字典（input_ids[B,T,2]、codec_ids[B,T,16]、 + text_embedding_mask[B,T,1]、codec_embedding_mask[B,T,1]、codec_mask[B,T]、 + attention_mask[B,T]、codec_0_labels[B,T]、ref_mels[B,F,128]）。 + target_speaker_embedding: 首步算出的目标说话人 emb（官方写进 codec_embedding.weight[3000]； + 本项目走 clone 推理，可传 None，仅用于一致性记录）。 + + Returns: + dict: + - 'talker_loss': 码本0 的 mean CE（model.talker.forward 的 outputs.loss，标准时间轴 shift 正确）。 + - 'sub_talker_logits': code_predictor 输出 [N, 15, V]（N=批次目标帧总数；15=码本1-15 头数）。 + - 'sub_talker_labels': [N, 15] = talker_codec_ids[:, 1:]（码本1-15 真值，已是 head 预测目标序列）。 + - 'target_speaker_embedding': 首步缓存的 speaker emb（供外部记录）。 + 总 loss = talker_loss + 0.3 × code_predictor CE，由 Qwen3TTSLoss 从本 dict 算（不在编排里算）。 + """ + talker = model.talker + # ---- ① speaker_encoder 前置（独立前向，detach 冻结）---- + with torch.no_grad(): + speaker_embedding = model.speaker_encoder(batch['ref_mels'].to(model.dtype)) + if target_speaker_embedding is None: + target_speaker_embedding = speaker_embedding + + # ---- ② 手动组装 inputs_embeds（不走 model(input_ids=...)）---- + input_text_ids = batch['input_ids'][..., 0] # [B,T] + input_codec_ids = batch['input_ids'][..., 1] # [B,T] + text_emb_mask = batch['text_embedding_mask'].to(model.dtype) + codec_emb_mask = batch['codec_embedding_mask'].to(model.dtype) + codec_ids = batch['codec_ids'] # [B,T,16] + codec_mask = batch['codec_mask'] # [B,T] + + # 0.6B: text_hidden_size=2048, hidden_size=1024. Even equal-sized + # variants must use the pretrained projection, as official inference does. + input_text_emb = talker.text_projection(talker.model.text_embedding(input_text_ids)) * text_emb_mask + input_codec_emb = talker.model.codec_embedding(input_codec_ids) * codec_emb_mask + input_codec_emb[:, 6, :] = speaker_embedding # 注入 pos6 + input_embeddings = input_text_emb + input_codec_emb + + # 叠加码本 1-15 的输入嵌入（code_predictor.get_input_embeddings() 返回 ModuleList） + cp_input_embs = talker.code_predictor.get_input_embeddings() # ModuleList[15] + for i in range(1, talker.config.num_code_groups): + e = cp_input_embs[i - 1](codec_ids[..., i]) * codec_mask.unsqueeze(-1).to(model.dtype) + input_embeddings = input_embeddings + e + + # ---- ③ 完整时间序列交 Talker；内部 ForCausalLMLoss 仅 shift 一次。---- + outputs = talker( + inputs_embeds=input_embeddings, + attention_mask=batch['attention_mask'], + labels=batch['codec_0_labels'], + output_hidden_states=True, + use_cache=False, + ) + talker_loss = outputs.loss # 码本0 mean shift+CE（HF loss_function，时间轴正确） + + # ---- ④ 第二段 forward：Code predictor 预测码本 1-15（向量化，一次喂所有目标帧）---- + # 注：forward_sub_talker_finetune 已向量化（输入 [N,16]+[N,H] → 内部 [N,16,H] 一次前向）， + hidden = outputs.hidden_states[0][-1] # [B, T, H] + talker_hidden = hidden[codec_mask] # 每帧对应位置的 hidden，不含 EOS/pad + talker_codec_ids = codec_ids[codec_mask] # [N, 16] + sub_talker_logits, _sub_loss_ignored = talker.forward_sub_talker_finetune( + talker_codec_ids, talker_hidden + ) + # 忽略 _sub_loss_ignored（继承的 ForCausalLMLoss 对码本轴有 off-by-one，见 loss.py）； + # Qwen3TTSLoss 用 sub_talker_logits + sub_talker_labels 自己算正确的 per-head CE。 + sub_talker_labels = talker_codec_ids[:, 1:] # [N, 15] = 码本1-15 真值 + + return { + 'talker_loss': talker_loss, + 'sub_talker_logits': sub_talker_logits, # [N, 15, V] + 'sub_talker_labels': sub_talker_labels, # [N, 15] + 'target_speaker_embedding': target_speaker_embedding, + } diff --git a/cookbook/qwen3_tts/tts/qwen3/input_processor.py b/cookbook/qwen3_tts/tts/qwen3/input_processor.py new file mode 100644 index 00000000..a1399498 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/input_processor.py @@ -0,0 +1,171 @@ +"""Qwen3-TTS 单卡 padded batch processor。 + +编码样本按时间 pad 并 stack；已 collate 的 dict 直接校验并转换设备/类型。 +参考 mel 保留 [B,F,128]，不同参考长度明确拒绝；不走通用 LM 字段过滤管线。 +""" + +from typing import Dict, List + +import numpy as np +import torch + +from twinkle.data_format import InputFeature +from twinkle import Platform, remote_function +from twinkle.processor.base import InputProcessor +from twinkle.utils.torch_utils import pad_and_stack_tensors +from .data import normalize_ref_mels + +# 序列字段：强制 stack（concat=False）→ [B, maxT, *] +_SEQ_STACK_FIELDS = ( + 'input_ids', 'codec_ids', 'codec_0_labels', + 'text_embedding_mask', 'codec_embedding_mask', 'codec_mask', + 'attention_mask', 'position_ids', +) +# mask 字段 stack 后补 channel 维 → [B, maxT, 1]（对齐官方 collate_fn + forward 广播） +_MASK_CHANNEL_FIELDS = ('text_embedding_mask', 'codec_embedding_mask') + +# to_transformers_dict 白名单追加（基类 `_keys` 之外的自定义字段） +_EXTRA_WHITELIST = ( + 'codec_ids', 'codec_0_labels', + 'text_embedding_mask', 'codec_embedding_mask', 'codec_mask', + 'ref_mels', +) + + +def _to_tensor(v) -> torch.Tensor: + """numpy/list → torch；bool → uint8（规避 F.pad 对 bool 的版本兼容问题；值仍是 0/1）。""" + if isinstance(v, np.ndarray): + v = torch.from_numpy(v) + elif isinstance(v, list): + v = torch.tensor(v) + if v.dtype == torch.bool: + v = v.to(torch.uint8) + return v + + +class Qwen3TTSInputProcessor(InputProcessor): + """Qwen3-TTS 批组织：2D 强制 stack + ref_mels stack + 自定义字段白名单。""" + + padding_map = { + **InputProcessor.padding_map, + 'codec_ids': 0, + 'codec_0_labels': -100, + 'text_embedding_mask': False, + 'codec_embedding_mask': False, + 'codec_mask': False, + 'ref_mels': 0.0, + } + + @remote_function() + def __call__(self, inputs, **kwargs): + """Single-card, padded TTS: encoded samples or an already collated dict. + + Return one tensor dict. Do not run the generic LM preparation pipeline: + it iterates dict keys, changes ranks and filters out TTS fields. + """ + if self.padding_free or kwargs.get('enable_sp', False): + raise ValueError('Qwen3 TTS supports padded single-card batches only') + if isinstance(inputs, dict) and getattr(inputs['input_ids'], 'ndim', None) == 3: + batch = self.to_transformers_dict([inputs])[0] + else: + samples = [inputs] if isinstance(inputs, dict) else inputs + if not samples: + raise ValueError('Empty TTS batch') + batch = self.to_transformers_dict(self.collate_fn(samples))[0] + required = set(_EXTRA_WHITELIST) | {'input_ids', 'attention_mask'} + missing = required - batch.keys() + if missing: + raise ValueError(f'Missing TTS fields: {sorted(missing)}') + batch = {key: torch.as_tensor(value) for key, value in batch.items()} + b, t, channels = batch['input_ids'].shape + if channels != 2 or b == 0 or t == 0: + raise ValueError('input_ids must have shape [B,T,2]') + shapes = {'codec_ids': (b, t, 16), 'codec_0_labels': (b, t), + 'attention_mask': (b, t), 'codec_mask': (b, t), + 'text_embedding_mask': (b, t, 1), 'codec_embedding_mask': (b, t, 1)} + for key, shape in shapes.items(): + if tuple(batch[key].shape) != shape: + raise ValueError(f'{key} must have shape {shape}') + mel = batch['ref_mels'] + if mel.ndim != 3 or mel.shape[0] != b or mel.shape[1] == 0 or mel.shape[2] != 128: + raise ValueError('ref_mels must have shape [B,F,128], F > 0') + if kwargs.get('micro_batch_size') not in (None, b): + raise ValueError('Split TTS micro-batches in the DataLoader') + model = kwargs.get('model') + device = getattr(model, 'device', None) + if device is None: + device = Platform.get_local_device() + for key, value in batch.items(): + dtype = (torch.float32 if key == 'ref_mels' else torch.bool + if key in (*_MASK_CHANNEL_FIELDS, 'codec_mask') else torch.long) + batch[key] = value.to(device=device, dtype=dtype) + return batch + + def collate_fn(self, inputs, micro_batch_size=None, variable_seq_lengths=False, **kwargs): + seq_out: Dict[str, torch.Tensor] = {} + + for key in _SEQ_STACK_FIELDS: + vals: List[torch.Tensor] = [] + for inp in inputs: + v = inp.get(key) + if v is None: + continue + v = _to_tensor(v) + # 撤销 prepare_inputs 对 1D 字段的 unsqueeze（[1,T]→[T]）； + # 2D 字段（[T,2]/[T,16]）shape[0]=T≠1 不受影响。 + if key not in ('input_ids', 'codec_ids') and v.dim() == 2 and v.shape[0] == 1: + v = v.squeeze(0) + vals.append(v) + if not vals: + continue + pad = int(self.padding_map.get(key, 0)) + if len(vals) == 1: + stacked = vals[0].unsqueeze(0) # 保证有 batch 维 + else: + stacked = pad_and_stack_tensors(vals, pad_value=pad, concat=False) + if key in _MASK_CHANNEL_FIELDS: + stacked = stacked.unsqueeze(-1) # [B, maxT, 1] + seq_out[key] = stacked + + # 保留每条参考的 batch 归属；无长度掩码时不自动补零。 + ref_mels: List[torch.Tensor] = [] + for inp in inputs: + v = inp.get('ref_mels') + if v is None: + continue + v = _to_tensor(v) + ref_mels.append(v) + if len(ref_mels) != len(inputs): + raise ValueError('Each sample requires one ref_mel') + seq_out['ref_mels'] = torch.stack(normalize_ref_mels(ref_mels), dim=0) + + # length：标量元数据，取 batch 内最大 + lengths = [int(inp.get('length')) for inp in inputs if inp.get('length') is not None] + if lengths: + seq_out['length'] = max(lengths) + + return [InputFeature(**seq_out)] + + @staticmethod + def to_transformers_dict(inputs, **kwargs): + results = [] + # 基类白名单 + 本模型自定义字段（基类 _keys 见 processor/base.py:663-680） + _keys = { + 'input_ids', 'inputs_embeds', 'attention_mask', 'position_ids', 'labels', + 'completion_mask', 'cu_seq_lens_q', 'cu_seq_lens_k', 'cu_seqlens_q', + 'cu_seqlens_kv', 'max_length_q', 'max_length_k', 'packed_seq_params', + 'routed_experts', 'mm_token_type_ids', 'second_per_grid_ts', + } | set(InputProcessor.VLM_CONCAT_FIELDS) | set(_EXTRA_WHITELIST) + + for _input in inputs: + output = {} + for key in list(_input.keys()): + if key not in _keys: + continue + value = _input[key] + if isinstance(value, torch.Tensor) or not isinstance(value, (list, np.ndarray)): + output[key] = value + else: + output[key] = np.array(value) + results.append(InputFeature(**output)) + return results diff --git a/cookbook/qwen3_tts/tts/qwen3/loss.py b/cookbook/qwen3_tts/tts/qwen3/loss.py new file mode 100644 index 00000000..82ee8485 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/loss.py @@ -0,0 +1,74 @@ +"""Qwen3-TTS 两段 loss：Talker 时间轴 CE + 0.3 倍残差码本同位 CE。 +源码依据：Qwen3-TTS 022e286。其 residual logits 与 cb1..cb15 已对齐，再 shift 会丢 cb1 的直接监督， +并让最后 head 无直接 CE；cb15 标签会被错配给前一 head。此结论不代表官方权重的训练方式或音质收益。""" + +import torch +import torch.nn.functional as F + +from twinkle.data_format import LossOutput + +from tts.common import MultiCodebookLoss + + +def _shift_ce_like_official(logits, labels, ignore_index=-100): + """复刻 transformers ForCausalLMLoss 的 shift 语义（含 off-by-one，仅供 shift_labels=True 对照）。 + + logits [N, C, V]，labels [N, C]。 + """ + from transformers.loss.loss_utils import ForCausalLMLoss + + return ForCausalLMLoss( + logits=logits, + labels=labels, + vocab_size=logits.shape[-1], + ignore_index=ignore_index, + ) + + +class Qwen3TTSLoss(MultiCodebookLoss): + """Qwen3-TTS 2 段式 loss。 + + Args: + sub_loss_weight: code_predictor 段权重（官方 0.3）。 + shift_labels: code_predictor 段是否做标准 LM shift。 + False（默认）= 修正 off-by-one，同位 per-head CE（无 shift）； + True = 复现官方 ForCausalLMLoss 含 shift（仅用于对照）。 + ignore_index: 忽略的 label 值（非音频帧 / pad）。 + """ + + def __init__( + self, + sub_loss_weight: float = 0.3, + shift_labels: bool = False, + ignore_index: int = -100, + ar_codebook: int = 0, + ): + super().__init__(ar_codebook=ar_codebook) + self.sub_loss_weight = sub_loss_weight + self.shift_labels = shift_labels + self.ignore_index = ignore_index + + def __call__(self, inputs, outputs, **kwargs) -> LossOutput: + # talker 段（码本0）：透传编排算好的 mean CE（标准时间轴 shift 正确） + talker_loss = outputs['talker_loss'] + # code_predictor 段（码本1-15） + sub_logits = outputs['sub_talker_logits'] # [N, 15, V] + sub_labels = outputs['sub_talker_labels'] # [N, 15] + + if self.shift_labels: + sub_loss = _shift_ce_like_official(sub_logits, sub_labels, self.ignore_index) + else: + V = sub_logits.shape[-1] + # 同位 per-head CE（无 shift）：logits[:,j] 预测码本 j+1 ↔ labels[:,j]=码本 j+1 + sub_loss = F.cross_entropy( + sub_logits.reshape(-1, V), + sub_labels.reshape(-1), + ignore_index=self.ignore_index, + reduction='mean', + ) + + total = talker_loss + self.sub_loss_weight * sub_loss + # 两段均 per-token mean（镜像官方 sft_12hz.py：outputs.loss 与 sub_loss 都是 mean）。 + # num_tokens=0 走 Twinkle PER-TOKEN-MEAN 归一（grad_accum 路径）， + # 与 accelerate.backward(loss)+gradient_accumulation_steps 一致。 + return LossOutput(loss=total, num_tokens=0) diff --git a/cookbook/qwen3_tts/tts/qwen3/manifest.py b/cookbook/qwen3_tts/tts/qwen3/manifest.py new file mode 100644 index 00000000..7894f4c8 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/manifest.py @@ -0,0 +1,71 @@ +"""Validated pre-encoded JSONL and deterministic per-epoch sample plans.""" +import hashlib +import json +from pathlib import Path +import random +import numpy as np + + +class Manifest: + def __init__(self, path): + self.path = Path(path).resolve(strict=True) + self.rows = [] + ids = set() + digest = hashlib.sha256() + for line in self.path.read_text(encoding='utf-8').splitlines(): + if not line.strip(): + continue + row = json.loads(line) + key = row['id'] + if not isinstance(key, str) or not key or key in ids: + raise ValueError('Manifest IDs must be unique nonempty strings') + ids.add(key) + if not isinstance(row['text'], str) or not row['text'].strip(): + raise ValueError('Empty/non-string transcript') + codes = np.asarray(row['audio_codes']) + if codes.ndim != 2 or codes.shape[0] == 0 or codes.shape[1] != 16 or codes.dtype.kind not in 'iu': + raise ValueError('audio_codes must be nonempty integer [T,16]') + if (codes < 0).any() or (codes >= 2048).any(): + raise ValueError('12Hz audio codes must be in [0,2048)') + mel_path = (self.path.parent / row['ref_mel_path']).resolve(strict=True) + mel = np.load(mel_path, allow_pickle=False) + if mel.ndim != 2 or mel.shape[1] != 128 or mel.shape[0] == 0 or not np.isfinite(mel).all(): + raise ValueError('ref mel must be finite [F,128]') + content = {k: row[k] for k in ('id', 'text', 'audio_codes')} + content['mel_sha256'] = hashlib.sha256(mel_path.read_bytes()).hexdigest() + digest.update(json.dumps(content, sort_keys=True, ensure_ascii=False).encode('utf-8') + b'\n') + self.rows.append({**row, 'ref_mel_path': str(mel_path)}) + if not self.rows: + raise ValueError('Manifest must not be empty') + self.fingerprint = digest.hexdigest() + + def __len__(self): + return len(self.rows) + + +def sample_plan(count, epochs, seed, shuffle): + if count < 1 or epochs < 1: + raise ValueError('Dataset and epochs must be positive') + plan = [] + for epoch in range(epochs): + order = list(range(count)) + if shuffle: + random.Random(seed + epoch).shuffle(order) # does not consume training RNG + plan.append(order) + return plan + + +def planned_windows(plan, accumulation_steps): + if accumulation_steps < 1: + raise ValueError('accumulation_steps must be positive') + # Flush a short window at each epoch boundary, never drop or carry it over. + return [order[i:i + accumulation_steps] for order in plan + for i in range(0, len(order), accumulation_steps)] + + +def validate_cursor(progress, plan_windows): + if not 0 <= progress.optimizer_steps <= len(plan_windows): + raise ValueError('Optimizer cursor outside sample plan') + consumed = sum(len(w) for w in plan_windows[:progress.optimizer_steps]) + if progress.pending_micro_steps or progress.micro_steps != consumed or progress.consumed_samples != consumed: + raise ValueError('Checkpoint sample cursor does not match plan boundary') diff --git a/cookbook/qwen3_tts/tts/qwen3/model.py b/cookbook/qwen3_tts/tts/qwen3/model.py new file mode 100644 index 00000000..9104ee92 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/model.py @@ -0,0 +1,80 @@ +'Twinkle TransformersModel bridge for the Qwen3-TTS two-stage forward and loss.' + +from copy import copy + +from twinkle import remote_function +from twinkle.infra import collect_tensor_dict +from twinkle.model.transformers import TransformersModel + +from .forward import qwen3_tts_forward_step +from .input_processor import Qwen3TTSInputProcessor + + +class Qwen3TTSModel(TransformersModel): + """Override `forward` 走 Qwen3-TTS 2 段编排 + stash 中间态。 + + `calculate_loss` / `backward` / `clip_grad_norm` 等复用基类（保留梯度累积/clip/metric 抽象 + + 多 LoRA 透传）。cookbook 仍调 `model.forward_backward(inputs=batch)`，编排全在本类内。 + """ + + def _construct_default_optimizer_group(self): + # Keep current Twinkle's group/metrics lifecycle, but pass tokenizer + # policy at construction (including groups created for new adapters). + from twinkle.model.transformers.transformers import OptimizerGroup + from .tokenization import local_template + from .loss import Qwen3TTSLoss + return OptimizerGroup( + loss_instance=Qwen3TTSLoss(), + template=local_template(self.tokenizer_id), + processor=Qwen3TTSInputProcessor(device_mesh=self.device_mesh), + _device_mesh=self.device_mesh, + ) + + @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) + def forward(self, *, inputs, **kwargs): + adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + # task 不用 causal_lm，但 pop 掉避免传给下游 qwen3_tts_forward_step + kwargs.pop('task', None) + kwargs.pop('temperature', 1.0) + kwargs.pop('sampling_masks', None) + kwargs.pop('return_logits', False) + optimizer_config = self.optimizer_group[adapter_name] + self._lazy_wrap_model() + if not inputs: + raise ValueError('inputs empty, check your DataLoader outputs') + self.model.train() + + # 1. Trajectory / List[Trajectory] → InputFeature（若未编码） + if (isinstance(inputs, dict) and self._not_encoded(inputs)) or ( + isinstance(inputs, list) and self._not_encoded(inputs[0]) + ): + assert optimizer_config.template is not None, \ + 'Use set_template to add a template when trying to input `List[Trajectory]`' + if isinstance(inputs, dict): + inputs = [inputs] + inputs = optimizer_config.template.batch_encode(inputs) + + # 2. 专用 processor 保留 TTS 字段和 batch 维度。 + processor = optimizer_config.processor + if not isinstance(processor, Qwen3TTSInputProcessor): + raise TypeError('Call set_processor(Qwen3TTSInputProcessor) before TTS forward') + inputs = processor( + inputs, + sp_strategy=self.sp_strategy, + model=self.model, + hf_config=self.hf_config, + enable_sp=getattr(self, '_enable_sp', False), + ) + + # 3. Qwen3-TTS 2 段前向编排（替换基类的 `self.model(**inputs)`） + # 产出中间态 dict：talker_loss / sub_talker_logits / sub_talker_labels / + # target_speaker_embedding。 + outputs = qwen3_tts_forward_step(self.model, inputs) + + # 4. stash 中间态进 train_status.outputs，供 Qwen3TTSLoss 从 outputs 读（calculate_loss 复用基类） + status = optimizer_config.train_status + status.inputs = inputs + status.outputs = outputs + status.forward_kwargs = kwargs + status.loss_value = 0 + return copy(outputs) diff --git a/cookbook/qwen3_tts/tts/qwen3/runtime.py b/cookbook/qwen3_tts/tts/qwen3/runtime.py new file mode 100644 index 00000000..6a4d1385 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/runtime.py @@ -0,0 +1,92 @@ +'Local Qwen3-TTS loading and assertions used by the single-card integration tests.' +import hashlib +import torch + +from .data import Qwen3TTSDataConfig + +DEFAULT_MODEL_PATH = '/models/Qwen3-TTS-12Hz-1.7B-Base' +ADAPTER = 'smoke' + + +def data_config(config): + tc = config.talker_config + return Qwen3TTSDataConfig( + **{name: getattr(config, name) for name in + ('tts_pad_token_id', 'tts_bos_token_id', 'tts_eos_token_id')}, + **{name: getattr(tc, name) for name in + ('codec_pad_id', 'codec_bos_id', 'codec_eos_token_id', + 'codec_nothink_id', 'codec_think_bos_id', 'codec_think_eos_id', + 'num_code_groups')}) + + +def load_twinkle_model(path, device_mesh): + # Twinkle constructs a Template (AutoConfig/AutoProcessor) as well as the + # supplied model_cls. Passing config alone does not register these classes. + from transformers import AutoConfig, AutoModel, AutoProcessor + from qwen_tts.core.models import Qwen3TTSConfig, Qwen3TTSForConditionalGeneration, Qwen3TTSProcessor + from .model import Qwen3TTSModel + AutoConfig.register('qwen3_tts', Qwen3TTSConfig, exist_ok=True) + AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration, exist_ok=True) + AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor, exist_ok=True) + model = Qwen3TTSModel( + model_id=str(path), model_cls=Qwen3TTSForConditionalGeneration, + config=Qwen3TTSConfig.from_pretrained(path, local_files_only=True), device_mesh=device_mesh, + mixed_precision='bf16', torch_dtype=torch.bfloat16, + attn_implementation='eager', local_files_only=True) + # A single short sample does not need activation checkpointing. In + # particular, reentrant checkpointing can sever gradients from frozen inputs. + model.model.gradient_checkpointing_disable() + return model + + +def trainable_snapshot(model): + params = {n: p for n, p in model.named_parameters() if p.requires_grad} + if not params or any('.lora_' not in n for n in params): + raise AssertionError('Only LoRA parameters may be trainable') + for branch in ('.talker.model.', '.talker.code_predictor.'): + if not any(branch in n for n in params): + raise AssertionError('Missing trainable branch: ' + branch) + return {n: p.detach().cpu().clone() for n, p in params.items()} + + +def frozen_digest(model): + """Hash all frozen weights, one tensor at a time; do not retain a base copy.""" + digest = hashlib.sha256() + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + continue + digest.update(name.encode()) + value = parameter.detach().cpu().contiguous() + digest.update(str((tuple(value.shape), value.dtype)).encode()) + digest.update(value.view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def assert_gradients(model): + nonzero = [] + for name, p in model.named_parameters(): + if not p.requires_grad: + if p.grad is not None: + raise AssertionError('Frozen parameter has gradient: ' + name) + continue + if p.grad is None or not torch.isfinite(p.grad).all(): + raise AssertionError('Missing/nonfinite LoRA gradient: ' + name) + if torch.count_nonzero(p.grad).item(): + nonzero.append(name) + # Standard LoRA starts B at zero: A can legitimately have zero gradient. + for branch in ('.talker.model.', '.talker.code_predictor.'): + if not any(branch in name for name in nonzero): + raise AssertionError('No nonzero gradient in ' + branch) + return nonzero + + +def assert_updated(before, after): + if before.keys() != after.keys(): + raise AssertionError('Trainable parameter set changed') + changed = [name for name in before if not torch.equal(before[name], after[name])] + if any(not torch.isfinite(value).all() for value in after.values()): + raise AssertionError('Nonfinite updated weights') + for branch in ('.talker.model.', '.talker.code_predictor.'): + if not any(branch in name for name in changed): + raise AssertionError('No parameter update in ' + branch) + return changed diff --git a/cookbook/qwen3_tts/tts/qwen3/speech_fixture.py b/cookbook/qwen3_tts/tts/qwen3/speech_fixture.py new file mode 100644 index 00000000..5bcefbe7 --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/speech_fixture.py @@ -0,0 +1,23 @@ +"""Load a hash-verified public regression sample; no download or customer data.""" +import hashlib +import json +from pathlib import Path +import numpy as np + + +def load_speech_fixture(directory): + import soundfile as sf + from scipy.signal import resample_poly + directory = Path(directory) + metadata = json.loads((directory / 'sample.json').read_text(encoding='utf-8')) + path = directory / metadata['audio'] + if hashlib.sha256(path.read_bytes()).hexdigest() != metadata['sha256']: + raise ValueError('Speech fixture SHA256 mismatch') + wave, sr = sf.read(path, dtype='float32', always_2d=False) + if sr != 16000 or wave.ndim != 1 or len(wave) == 0 or not np.isfinite(wave).all(): + raise ValueError('Expected finite mono 16 kHz LibriSpeech fixture') + if not metadata['text'].strip(): + raise ValueError('Speech fixture transcript is empty') + # Exact rational rate conversion; preserve the full utterance and transcript. + wave = resample_poly(wave, 3, 2).astype(np.float32) + return wave, metadata['text'], metadata diff --git a/cookbook/qwen3_tts/tts/qwen3/tokenization.py b/cookbook/qwen3_tts/tts/qwen3/tokenization.py new file mode 100644 index 00000000..f68e99ab --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/tokenization.py @@ -0,0 +1,21 @@ +"""One local-only tokenizer policy for training, templates and adapter reload.""" + +def load_text_tokenizer(path): + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained(str(path), local_files_only=True, fix_mistral_regex=True) + + +def local_template(path): + from twinkle.template import Template + return Template(str(path), local_files_only=True, fix_mistral_regex=True) + + +def assert_tokenizers_match(reference, others, texts): + """Compare actual IDs, and ensure the reference preserves input text.""" + for text in texts: + ids = reference.encode(text, add_special_tokens=False) + if reference.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) != text: + raise AssertionError('Tokenizer does not preserve probe text') + for tokenizer in others: + if tokenizer.encode(text, add_special_tokens=False) != ids: + raise AssertionError('Tokenizer IDs differ across training/template/reload') diff --git a/cookbook/qwen3_tts/tts/qwen3/training.py b/cookbook/qwen3_tts/tts/qwen3/training.py new file mode 100644 index 00000000..5854dc6f --- /dev/null +++ b/cookbook/qwen3_tts/tts/qwen3/training.py @@ -0,0 +1,112 @@ +'Single-card accumulation and update-boundary checkpoints. The caller owns epoch shuffling.' +from dataclasses import asdict, dataclass +import json +import math +from pathlib import Path + + +@dataclass +class Progress: + accumulation_steps: int = 2 + micro_steps: int = 0 + optimizer_steps: int = 0 + consumed_samples: int = 0 + pending_micro_steps: int = 0 + + def __post_init__(self): + if self.accumulation_steps < 1: + raise ValueError('accumulation_steps must be positive') + + +def windows(batches, size): + if size < 1: + raise ValueError('window size must be positive') + window = [] + for batch in batches: + window.append(batch) + if len(window) == size: + yield window + window = [] + if window: + yield window + + +def train_window(model, batches, progress, adapter_name, before_micro=None): + """Exactly one optimizer/scheduler update for a full or final short window.""" + batches = list(batches) + if not 1 <= len(batches) <= progress.accumulation_steps: + raise ValueError('Expected one nonempty accumulation window') + group = model.optimizer_group[adapter_name] + if group.gradient_accumulation_steps != 1: + raise ValueError('Twinkle accumulation must be 1; cookbook owns boundaries') + if progress.pending_micro_steps or group.train_status.num_tokens: + raise ValueError('Previous accumulation window is incomplete') + if group.cur_step != progress.micro_steps: + raise ValueError('Twinkle and cookbook micro-step counts differ') + # Validate the entire window before executing any backward operation. + for batch in batches: + if batch['input_ids'].shape[0] != 1: + raise ValueError('Only single-sample micro-batches are supported') + losses = [] + for batch in batches: + if before_micro is not None: + before_micro(progress.micro_steps) + result = model.forward_backward(inputs=batch, adapter_name=adapter_name) + if not math.isfinite(float(result['loss'])): + raise FloatingPointError('Nonfinite micro-batch loss; abort this run') + losses.append(float(result['loss'])) + progress.micro_steps += 1 + progress.consumed_samples += 1 + progress.pending_micro_steps += 1 + # Each Qwen loss is a mean and reports num_tokens=0. Twinkle counts each + # as 1, so clipping normalizes by the ACTUAL window size, including tails. + if group.train_status.num_tokens != len(batches): + raise AssertionError('Unexpected loss normalization count') + model.clip_grad_and_step(adapter_name=adapter_name) + progress.optimizer_steps += 1 + progress.pending_micro_steps = 0 + return losses + + +def save_boundary(model, progress, output_dir, name, adapter_name, sample_fingerprint): + group = model.optimizer_group[adapter_name] + if progress.pending_micro_steps or group.train_status.num_tokens: + raise ValueError('Cannot save with unstepped gradients') + if group.cur_step != progress.micro_steps: + raise ValueError('Micro-step count mismatch') + if any(p.grad is not None for p in model.model.parameters()): + raise ValueError('Checkpoint requires cleared gradients') + checkpoint = Path(model.save(name, output_dir=str(output_dir), adapter_name=adapter_name, + save_optimizer=True, consumed_train_samples=progress.consumed_samples)) + (checkpoint / 'cookbook_progress.json').write_text(json.dumps({ + 'format_version': 1, 'progress': asdict(progress), + 'sample_fingerprint': sample_fingerprint}, indent=2), encoding='utf-8') + return checkpoint + + +def resume_boundary(model, checkpoint, adapter_name, sample_fingerprint, accumulation_steps): + """Restore only complete locally produced checkpoints, never silently restart.""" + checkpoint = Path(checkpoint) + for name in ('adapter_config.json', 'adapter_model.safetensors', 'optimizer.pt', + 'scheduler.pt', 'rng_state.pt', 'trainer_state.json', 'cookbook_progress.json'): + if not (checkpoint / name).is_file(): + raise FileNotFoundError(checkpoint / name) + state = json.loads((checkpoint / 'cookbook_progress.json').read_text(encoding='utf-8')) + if state['format_version'] != 1 or state['sample_fingerprint'] != sample_fingerprint: + raise ValueError('Checkpoint format/sample mismatch') + progress = Progress(**state['progress']) + if progress.pending_micro_steps or progress.accumulation_steps != accumulation_steps: + raise ValueError('Checkpoint boundary/accumulation mismatch') + trainer = json.loads((checkpoint / 'trainer_state.json').read_text(encoding='utf-8')) + if (trainer['cur_step'] != progress.micro_steps + or trainer['consumed_train_samples'] != progress.consumed_samples + or trainer['gradient_accumulation_steps'] != 1): + raise ValueError('Twinkle checkpoint progress disagrees with cookbook') + group = model.optimizer_group[adapter_name] + if getattr(group, 'optimizer', None) is None or getattr(group, 'lr_scheduler', None) is None: + raise ValueError('Create optimizer and scheduler before resuming') + model.resume_from_checkpoint(str(checkpoint), adapter_name=adapter_name) + group = model.optimizer_group[adapter_name] + if group.cur_step != progress.micro_steps or group.gradient_accumulation_steps != 1: + raise AssertionError('Twinkle did not restore micro-step state') + return progress