音效开启cpu offload

[zwukong]

16G跑不动，支持accelerate开启cpu offload吗

from accelerate import init_empty_weights, load_checkpoint_and_dispatch
import torch

# 1. 安装依赖（首次使用需执行）
# pip install accelerate

# 2. 修改模型加载方式
model = AutoModel.from_pretrained(
    pretrained_model_name_or_path,
    trust_remote_code=True,
    attn_implementation=attn_implementation,
    torch_dtype=dtype,
    device_map="auto",  # ← 启用自动设备映射（含 CPU offload）
    offload_folder="offload_cache",  # 可选：指定 offload 临时目录
)
# model.to(device)  # ←  移除

[YWMditto]

Of course — you can try the code below. We verified that GPU memory usage drops to 14–15 GB, with an inference speed of about 3 tokens/s.

from pathlib import Path
import importlib.util
from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal, Final, cast

import torch
import torchaudio
from transformers import AutoModel, AutoProcessor

# Disable the broken cuDNN SDPA backend
torch.backends.cuda.enable_cudnn_sdp(False)
# Keep these enabled as fallbacks
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)

pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-TTS"

# 这里的 device 只用于“输入张量放哪儿”；
# 模型本身会由 accelerate 的 device_map 自动切分（含 CPU offload）。
has_cuda = torch.cuda.is_available()
input_device = "cuda" if has_cuda else "cpu"

# 给 accelerate 一个明确的显存预算，防止 "auto" 把权重尽量塞满 GPU。
# 17GB -> 16GB以内时，建议把预算压到 14GiB，给运行时缓存留余量。
gpu_vram_budget_gib = 14
cpu_offload_budget_gib = 120

# dtype：如果有 CUDA，建议 bfloat16/float16；否则 float32
dtype = torch.bfloat16 if has_cuda else torch.float32


def resolve_attn_implementation() -> str:
    # Prefer FlashAttention 2 when package + device conditions are met.
    if (
        has_cuda
        and importlib.util.find_spec("flash_attn") is not None
        and dtype in {torch.float16, torch.bfloat16}
    ):
        major, _ = torch.cuda.get_device_capability()
        if major >= 8:
            return "flash_attention_2"

    # CUDA fallback: use PyTorch SDPA kernels.
    if has_cuda:
        return "sdpa"

    # CPU fallback.
    return "eager"


attn_implementation = resolve_attn_implementation()
print(f"[INFO] Using attn_implementation={attn_implementation}")
print(f"[INFO] has_cuda={has_cuda}, input_device={input_device}, dtype={dtype}")

if has_cuda:
    max_memory: Optional[Dict[Union[int, str], str]] = {
        i: f"{gpu_vram_budget_gib}GiB" for i in range(torch.cuda.device_count())
    }
    max_memory["cpu"] = f"{cpu_offload_budget_gib}GiB"
else:
    max_memory = None

# -------- Processor --------


processor = AutoProcessor.from_pretrained(
    pretrained_model_name_or_path,
    trust_remote_code=True,
)

# 关键点：
# - 既然你要 CPU offload，audio_tokenizer 不要强行 .to("cuda")
# - 让它保持默认（通常在 CPU），可最大化省显存
# - 如果你的 audio_tokenizer 特别慢，也可以改成放 GPU，但那会增加显存占用
processor.audio_tokenizer = processor.audio_tokenizer.to("cpu")

# -------- Inputs --------
text_1 = "亲爱的你，\n你好呀。\n\n今天，我想用最认真、最温柔的声音，对你说一些重要的话。\n这些话，像一颗小小的星星，希望能在你的心里慢慢发光。\n\n首先，我想祝你——\n每天都能平平安安、快快乐乐。\n\n希望你早上醒来的时候，\n窗外有光，屋子里很安静，\n你的心是轻轻的，没有着急，也没有害怕。\n\n希望你吃饭的时候胃口很好，\n走路的时候脚步稳稳，\n晚上睡觉的时候，能做一个又一个甜甜的梦。\n\n我希望你能一直保持好奇心。\n对世界充满问题，\n对天空、星星、花草、书本和故事感兴趣。\n当你问“为什么”的时候，\n希望总有人愿意认真地听你说话。\n\n我也希望你学会温柔。\n温柔地对待朋友，\n温柔地对待小动物，\n也温柔地对待自己。\n\n如果有一天你犯了错，\n请不要太快责怪自己，\n因为每一个认真成长的人，\n都会在路上慢慢学会更好的方法。\n\n愿你拥有勇气。\n当你站在陌生的地方时，\n当你第一次举手发言时，\n当你遇到困难、感到害怕的时候，\n希望你能轻轻地告诉自己：\n“我可以试一试。”\n\n就算没有一次成功，也没有关系。\n失败不是坏事，\n它只是告诉你，你正在努力。\n\n我希望你学会分享快乐。\n把开心的事情告诉别人，\n把笑声送给身边的人，\n因为快乐被分享的时候，\n会变得更大、更亮。\n\n如果有一天你感到难过，\n我希望你知道——\n难过并不丢脸，\n哭泣也不是软弱。\n\n愿你能找到一个安全的地方，\n慢慢把心里的话说出来，\n然后再一次抬起头，看见希望。\n\n我还希望你能拥有梦想。\n这个梦想也许很大，\n也许很小，\n也许现在还说不清楚。\n\n没关系。\n梦想会和你一起长大，\n在时间里慢慢变得清楚。\n\n最后，我想送你一个最最重要的祝福：\n\n愿你被世界温柔对待，\n也愿你成为一个温柔的人。\n\n愿你的每一天，\n都值得被记住，\n都值得被珍惜。\n\n亲爱的你，\n请记住，\n你是独一无二的，\n你已经很棒了，\n而你的未来，\n一定会慢慢变得闪闪发光。\n\n祝你健康、勇敢、幸福，\n祝你永远带着笑容向前走。"
text_2 = "We stand on the threshold of the AI era.\nArtificial intelligence is no longer just a concept in laboratories, but is entering every industry, every creative endeavor, and every decision. It has learned to see, hear, speak, and think, and is beginning to become an extension of human capabilities. AI is not about replacing humans, but about amplifying human creativity, making knowledge more equitable, more efficient, and allowing imagination to reach further. A new era, jointly shaped by humans and intelligent systems, has arrived."
text_3 = "nin2 hao3，qing3 wen4 nin2 lai2 zi4 na3 zuo4 cheng2 shi4？"
text_4 = "nin2 hao3，qing4 wen3 nin2 lai2 zi4 na4 zuo3 cheng4 shi3？"
text_5 = "您好，请问您来自哪 zuo4 cheng2 shi4？"
text_6 = "/həloʊ, meɪ aɪ æsk wɪtʃ sɪti juː ɑːr frʌm?/"

ref_audio_1 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_zh.wav"
ref_audio_2 = "https://speech-demo.oss-cn-shanghai.aliyuncs.com/moss_tts_demo/tts_readme_demo/reference_en.m4a"

conversations = [
    [processor.build_user_message(text=text_1)],
    [processor.build_user_message(text=text_2)],
    [processor.build_user_message(text=text_3)],
    [processor.build_user_message(text=text_4)],
    [processor.build_user_message(text=text_5)],
    [processor.build_user_message(text=text_6)],
    [processor.build_user_message(text=text_1, reference=[ref_audio_1])],
    [processor.build_user_message(text=text_2, reference=[ref_audio_2])],
    [processor.build_user_message(text=text_2, tokens=325)],
    [processor.build_user_message(text=text_2, tokens=600)],
]

# -------- Model with CPU offload --------
# 关键点：
# 1) device_map="auto" 触发 accelerate 自动切分
# 2) offload_folder 指定 CPU offload 临时目录
# 3) 不要再 model.to(device)
# 4) low_cpu_mem_usage=True 通常更省内存（默认也常为 True，这里显式写上）
model = AutoModel.from_pretrained(
    pretrained_model_name_or_path,
    trust_remote_code=True,
    attn_implementation=attn_implementation,
    torch_dtype=dtype,
    device_map="auto",               # ✅ 自动设备映射（含 CPU offload）
    max_memory=max_memory,           # ✅ 显存硬上限，超出的层会 offload 到 CPU
    offload_folder="offload_cache",  # ✅ offload 临时目录
    offload_state_dict=True,         # ✅ 常见配置：权重也可 offload/分片加载
    low_cpu_mem_usage=True,
)

model.eval()

# 按实际切分结果决定输入放哪张卡，避免固定 "cuda" 带来的不必要占用。
if has_cuda and hasattr(model, "hf_device_map"):
    for _, mapped_device in model.hf_device_map.items():
        if isinstance(mapped_device, int):
            input_device = f"cuda:{mapped_device}"
            break
        if isinstance(mapped_device, str) and mapped_device.startswith("cuda"):
            input_device = mapped_device
            break
print(f"[INFO] resolved_input_device={input_device}")

# -------- Inference --------
batch_size = 1
max_new_tokens = 2048  # 先保证显存稳定在 16GB 内，时长不够再逐步上调

save_dir = Path("inference_root")
save_dir.mkdir(exist_ok=True, parents=True)

sample_idx = 0
if has_cuda:
    torch.cuda.reset_peak_memory_stats()
with torch.no_grad():
    for start in range(0, len(conversations), batch_size):
        batch_conversations = conversations[start : start + batch_size]
        batch = processor(batch_conversations, mode="generation")

        # 输入张量：放到 input_device（一般 cuda），模型自己会在需要时搬运/dispatch
        input_ids = batch["input_ids"].to(input_device)
        attention_mask = batch["attention_mask"].to(input_device)

        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
        )

        for message in processor.decode(outputs):
            audio = message.audio_codes_list[0]

            # 为安全起见：torchaudio.save 期望 CPU tensor（某些版本/后端对 GPU tensor 不友好）
            audio = audio.detach().to("cpu")

            out_path = save_dir / f"sample{sample_idx}.wav"
            sample_idx += 1
            torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)

        if has_cuda:
            peak_alloc = torch.cuda.max_memory_allocated() / (1024 ** 3)
            peak_reserved = torch.cuda.max_memory_reserved() / (1024 ** 3)
            print(f"[INFO] peak_alloc={peak_alloc:.2f} GiB, peak_reserved={peak_reserved:.2f} GiB")
            torch.cuda.reset_peak_memory_stats()

print(f"[INFO] Done. Saved {sample_idx} samples to: {save_dir.resolve()}")

[zwukong]

好家伙，SoundEffect还得下一个8G的MOSS-Audio-Tokenizer， 玩不了了，太夸张了。https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer
https://huggingface.co/OpenMOSS-Team/MOSS-SoundEffect

[zwukong]

速度过于夸张，光accelerate就要处理1个小时

[INFO] Using attn_implementation=flash_attention_2
[INFO] has_cuda=True, input_device=cuda, dtype=torch.bfloat16
Loading weights: 100%|█| 1600/1600 [00:00<00:00, 1617.36it/s, Materializing param=quantizer.quantizers.31.out_proj.para
`torch_dtype` is deprecated! Use `dtype` instead!
Loading weights: 100%|███████████████████████| 431/431 [00:05<00:00, 79.07it/s, Materializing param=lm_heads.16.weight]
Some parameters are on the meta device because they were offloaded to the cpu.
[INFO] resolved_input_device=cuda:0
Generating bs1 ...:   3%|█▊                                                        | 64/2048 [02:02<1:03:12,  1.91s/it]
Traceback (most recent call last):
  File "K:\AI\MOSS-TTS\e1.py", line 154, in <module>
    outputs = model.generate(
  File "K:\AI\miniconda3\envs\moss-tts\Lib\site-packages\accelerate\utils\modeling.py", line 343, in set_module_tensor_to_device
    new_value = value.to(device, non_blocking=non_blocking)

[xiami2019]

你好 @zwukong MOSS-Audio-Tokenizer是1.6B参数。
我们后面会迭代小的音效模型~

[zwukong]

套路太深，Tokenizer是一点不提呀，知道还有这玩意就不下了，浪费时间。比其他tts一点优势也没有