解决内存占用较高问题 / Fix High Memory Usage Issue

简介 / Introduction

当前VoxCPM2在运行时内存占用过高，可能影响在资源受限环境下的部署和运行效率。

Currently, VoxCPM2 has a high memory footprint during runtime, which may affect deployment and operational efficiency in resource-constrained environments.

解决方案 / Solution

voxcpm.model.voxcpm2_low_memeory.VoxCPM2ModelLowMemory.from_local

@classmethod
    def from_local(
            cls,
            path: str,
            optimize: bool = True,
            training: bool = False,
            device: str | None = None,
            lora_config: LoRAConfig = None,
    ):
        with open(os.path.join(path, "config.json"), "r", encoding="utf-8") as _cfg_f:
            config = VoxCPMConfig.model_validate_json(_cfg_f.read())
        tokenizer = LlamaTokenizerFast.from_pretrained(path)
        audio_vae_config = getattr(config, "audio_vae_config", None)

        with torch.device('cuda'):
            audio_vae = AudioVAEV2(config=audio_vae_config) if audio_vae_config else AudioVAEV2()
            torch.set_default_dtype(torch.bfloat16)
            model = cls(config, tokenizer, audio_vae, lora_config, device=device)

        # Try to load AudioVAE from safetensors first, fallback to pytorch
        audiovae_safetensors_path = os.path.join(path, "audiovae.safetensors")
        audiovae_pth_path = os.path.join(path, "audiovae.pth")
        if os.path.exists(audiovae_safetensors_path) and SAFETENSORS_AVAILABLE:
            print(f"Loading AudioVAE from safetensors: {audiovae_safetensors_path}", file=sys.stderr)
            vae_state_dict = load_file(audiovae_safetensors_path, device="cpu")
        elif os.path.exists(audiovae_pth_path):
            print(f"Loading AudioVAE from pytorch: {audiovae_pth_path}", file=sys.stderr)
            checkpoint = torch.load(
                audiovae_pth_path,
                map_location="cpu",
                weights_only=True,
            )
            vae_state_dict = checkpoint.get("state_dict", checkpoint)
        else:
            raise FileNotFoundError(
                f"AudioVAE checkpoint not found. Expected either {audiovae_safetensors_path} or {audiovae_pth_path}"
            )
        if not training:
            lm_dtype = get_dtype(model.config.dtype)
            model = model.to(lm_dtype)
        else:  # training mode
            for name, param in model.named_parameters():
                if "audio_vae" in name:  # freeze VAE weights
                    param.requires_grad = False
                    continue
                if lora_config is not None:
                    if "lora" not in name:  # freeze non-LoRA weights
                        param.requires_grad = False
        model.audio_vae = model.audio_vae.to(torch.float32)

        # Try to load from safetensors first, fallback to pytorch_model.bin
        safetensors_path = os.path.join(path, "model.safetensors")
        pytorch_model_path = os.path.join(path, "pytorch_model.bin")

        if os.path.exists(safetensors_path) and SAFETENSORS_AVAILABLE:
            print(f"Loading model from safetensors: {safetensors_path}", file=sys.stderr)
            model_state_dict = load_file(safetensors_path)
        elif os.path.exists(pytorch_model_path):
            print(f"Loading model from pytorch_model.bin: {pytorch_model_path}", file=sys.stderr)
            checkpoint = torch.load(
                pytorch_model_path,
                map_location="cpu",
                weights_only=True,
            )
            model_state_dict = checkpoint.get("state_dict", checkpoint)
        else:
            raise FileNotFoundError(f"Model file not found. Expected either {safetensors_path} or {pytorch_model_path}")

        for kw, val in vae_state_dict.items():
            model_state_dict[f"audio_vae.{kw}"] = val

        # LoRALinear keeps weight/bias compatible with nn.Linear but adds
        # lora_A/lora_B, which are absent from base pretrained checkpoints.
        model.load_state_dict(model_state_dict, strict=False, assign=True)
        if training:
            return model

        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.synchronize()  # 等待 GPU 完成所有操作
            gc.collect()  # 再次触发 GC，确保 CPU 侧的旧参数被释放

        return model.to(model.device).eval().optimize(disable=True)

Start with voxcpm.model.voxcpm2_low_memeory.VoxCPM2ModelLowMemory.from_local and trace how the AudioVAE and model checkpoints are loaded, assigned, converted, and moved to the device. Measure runtime memory in the affected loading path and verify that the completed change reduces the footprint while preserving model loading and inference behavior.