| from __future__ import annotations |
|
|
| import gc |
| import os |
| from pathlib import Path |
| from typing import Iterator |
| import torch |
| from safetensors import safe_open |
|
|
|
|
| def work_dir() -> Path: |
| env = os.environ.get("MAYA_FAST_WORK") |
| if env: |
| return Path(env) |
| return Path.home() / ".cache" / "maya_fast" |
|
|
|
|
| def ckpt_paths() -> dict: |
| from maya3 import hub |
| return hub.resolve("maya-research/maya2-native") |
|
|
|
|
| def iter_ckpt_tensors( |
| skip_prefixes: tuple[str, ...] = (), |
| files: list[str] | None = None, |
| ) -> Iterator[tuple[str, torch.Tensor]]: |
| """Yield one tensor at a time from checkpoint safetensors, stripping model. prefix. |
| If files is None, uses ckpt_paths()["ckpt"]. |
| """ |
| from tts.model import load_s3 as L |
|
|
| if files is None: |
| paths = ckpt_paths() |
| src = os.path.abspath(str(paths["ckpt"])) |
| _, files, _ = L._safetensors_files(src) |
|
|
| seen = set() |
| for f in files: |
| with safe_open(f, "pt") as fh: |
| for k in fh.keys(): |
| name = k[len(L.MODEL_PREFIX):] if k.startswith(L.MODEL_PREFIX) else k |
| if any(name.startswith(p) for p in skip_prefixes): |
| continue |
| if name in seen: |
| continue |
| seen.add(name) |
| t = fh.get_tensor(k) |
| yield name, t |
| del t |
|
|
|
|
| def _stub_lm(path: str, dtype: torch.dtype, attn: str): |
| from transformers import AutoConfig, AutoModelForCausalLM |
| cfg = AutoConfig.from_pretrained(path) |
| cfg.num_hidden_layers = 0 |
| cfg.layer_types = [] |
| with torch.device("meta"): |
| lm = AutoModelForCausalLM.from_config(cfg, dtype=torch.float32, attn_implementation="sdpa") |
| for m in lm.modules(): |
| if hasattr(m, "inv_freq") and hasattr(m, "rope_type"): |
| from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS |
| fn = m.rope_init_fn if getattr(m, "rope_init_fn", None) is not None else ( |
| m.compute_default_rope_parameters if m.rope_type == "default" else ROPE_INIT_FUNCTIONS[m.rope_type]) |
| with torch.device("cpu"): |
| inv, scale = fn(m.config, torch.device("cpu")) |
| m.inv_freq = inv |
| if hasattr(m, "original_inv_freq"): |
| m.original_inv_freq = inv |
| if hasattr(m, "attention_scaling"): |
| m.attention_scaling = scale |
| lm.tie_weights() |
| return lm |
|
|
|
|
| class _UpcastFp32(torch.nn.Module): |
| """Parametrization: compute on an fp32 copy of a bf16 weight, made per call.""" |
|
|
| def forward(self, w: torch.Tensor) -> torch.Tensor: |
| return w.float() |
|
|
|
|
| class _Bf16Embedding(torch.nn.Embedding): |
| def forward(self, idx: torch.Tensor) -> torch.Tensor: |
| |
| return torch.nn.functional.embedding(idx, self.weight).float() |
|
|
|
|
| def _keep_bf16_params_fp32_compute(model: torch.nn.Module) -> None: |
| """Parameters stored as bf16 are upcast to fp32 at use, so weights stay half-size in RAM. |
| Embedding tables are read by row. Other modules get a parametrization that upcasts per call. |
| Upcasting bf16 is exact, so results match an fp32 copy of the same weights.""" |
| from torch.nn.utils import parametrize |
|
|
| embed_weights = set() |
| for mod in model.modules(): |
| if isinstance(mod, torch.nn.Embedding) and mod.weight.dtype == torch.bfloat16: |
| embed_weights.add(id(mod.weight)) |
| mod.__class__ = _Bf16Embedding |
| for mod in list(model.modules()): |
| for pname, p in list(mod._parameters.items()): |
| if p is None or p.dtype != torch.bfloat16 or id(p) in embed_weights: |
| continue |
| if isinstance(mod, torch.nn.Embedding): |
| continue |
| parametrize.register_parametrization(mod, pname, _UpcastFp32(), unsafe=True) |
|
|
|
|
| def build_light_model(): |
| """Build light TTSModelV2 on meta device with 0 backbone decoder layers, |
| and populate non-backbone layer weights via iter_ckpt_tensors.""" |
| import tts.model.backbone as BB |
| import tts.model.load_s3 as L |
|
|
| orig_stub = BB._lm_from_config |
| try: |
| BB._lm_from_config = _stub_lm |
| paths = ckpt_paths() |
| src = str(paths["ckpt"]) |
| model_cfg, cfg_path = L.resolve_model_cfg(src) |
| model_cfg["attn"] = "sdpa" |
|
|
| anchor_cfg = model_cfg.get("speaker_anchor") |
| depth_grouped = bool(model_cfg.get("depth_grouped", True)) |
|
|
| with torch.device("meta"): |
| model = L.build_eval_model( |
| model_cfg, |
| anchor_cfg=anchor_cfg, |
| amortize_frac=1.0, |
| depth_grouped=depth_grouped, |
| ) |
|
|
| mods = dict(model.named_modules()) |
| for name, t in iter_ckpt_tensors(skip_prefixes=("backbone.lm.model.layers.",)): |
| parent_name, _, leaf = name.rpartition(".") |
| parent = mods.get(parent_name) |
| if parent is None or (leaf not in parent._parameters and leaf not in parent._buffers): |
| continue |
| is_param = leaf in parent._parameters |
| if is_param: |
| |
| parent._parameters[leaf] = torch.nn.Parameter(t, requires_grad=False) |
| else: |
| parent._buffers[leaf] = t.float() |
| del t |
|
|
| model.backbone.lm.tie_weights() |
| _keep_bf16_params_fp32_compute(model) |
| gc.collect() |
| return model.eval() |
| finally: |
| BB._lm_from_config = orig_stub |
|
|
