diff --git a/configs/acoustic.yaml b/configs/acoustic.yaml index fad75600e..e5f71c24a 100644 --- a/configs/acoustic.yaml +++ b/configs/acoustic.yaml @@ -67,6 +67,7 @@ max_beta: 0.02 enc_ffn_kernel_size: 3 use_rope: true rope_interleaved: false +rope_theta: 10000 use_stretch_embed: true use_variance_scaling: true rel_pos: true diff --git a/configs/templates/config_acoustic.yaml b/configs/templates/config_acoustic.yaml index e344fb450..02ff55b6e 100644 --- a/configs/templates/config_acoustic.yaml +++ b/configs/templates/config_acoustic.yaml @@ -75,6 +75,7 @@ diffusion_type: reflow enc_ffn_kernel_size: 3 use_rope: true rope_interleaved: false +rope_theta: 10000 use_stretch_embed: true use_variance_scaling: true use_shallow_diffusion: true diff --git a/configs/templates/config_variance.yaml b/configs/templates/config_variance.yaml index 116154ac7..d3acd51a6 100644 --- a/configs/templates/config_variance.yaml +++ b/configs/templates/config_variance.yaml @@ -66,6 +66,7 @@ tension_logit_max: 10.0 enc_ffn_kernel_size: 3 use_rope: true rope_interleaved: false +rope_theta: 10000 use_stretch_embed: false use_variance_scaling: true hidden_size: 384 diff --git a/configs/variance.yaml b/configs/variance.yaml index d4e203670..961857c34 100644 --- a/configs/variance.yaml +++ b/configs/variance.yaml @@ -37,6 +37,7 @@ predict_tension: false enc_ffn_kernel_size: 3 use_rope: true rope_interleaved: false +rope_theta: 10000 use_stretch_embed: false use_variance_scaling: true rel_pos: true diff --git a/modules/commons/rotary_embedding_torch.py b/modules/commons/rotary_embedding_torch.py index 1a1fa193e..703bf940f 100644 --- a/modules/commons/rotary_embedding_torch.py +++ b/modules/commons/rotary_embedding_torch.py @@ -1,29 +1,22 @@ import torch -from einops import rearrange, repeat -from torch import einsum, Tensor +from torch import Tensor from torch.nn import Module -def rotate_half(x: Tensor, interleaved=True) -> Tensor: - if not interleaved: - # x_half1, x_half2 = x.chunk(2, dim=-1) - # Using torch.split instead of chunk for ONNX export compatibility. - x1, x2 = torch.split(x, x.size(-1) // 2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x = rearrange(x, '... (d r) -> ... d r', r=2) - x1, x2 = x.unbind(dim=-1) - x = torch.stack((-x2, x1), dim=-1) - return rearrange(x, '... d r -> ... (d r)') - - -def apply_rotary_emb(freqs: Tensor, t: Tensor, interleaved=True) -> Tensor: - rot_dim = freqs.shape[-1] +def apply_rotary_emb(freqs_cos: Tensor, freqs_sin: Tensor, t: Tensor, interleaved=True) -> Tensor: + rot_dim = freqs_cos.shape[-1] t_to_rotate = t[..., :rot_dim] t_pass_through = t[..., rot_dim:] - t_rotated = (t_to_rotate * freqs.cos()) + (rotate_half(t_to_rotate, interleaved) * freqs.sin()) + if interleaved: + x = t_to_rotate.view(*t_to_rotate.shape[:-1], t_to_rotate.size(-1) // 2, 2) + x1, x2 = x.unbind(dim=-1) + rotated_half = torch.stack((-x2, x1), dim=-1).reshape_as(t_to_rotate) + else: + x1, x2 = torch.split(t_to_rotate, t_to_rotate.size(-1) // 2, dim=-1) + rotated_half = torch.cat((-x2, x1), dim=-1) + t_rotated = (t_to_rotate * freqs_cos) + (rotated_half * freqs_sin) return torch.cat((t_rotated, t_pass_through), dim=-1) @@ -40,24 +33,29 @@ def __init__( self.cached_freqs_seq_len = max_seq_len inv_freq = 1. / (theta ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer('inv_freq', inv_freq, persistent=False) - self.register_buffer('cached_freqs', self._precompute_cache(max_seq_len), persistent=False) + cos, sin = self._precompute_cache(max_seq_len) + self.register_buffer('cached_cos', cos, persistent=False) + self.register_buffer('cached_sin', sin, persistent=False) def _precompute_cache(self, seq_len: int): - seq = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) - freqs = einsum('i, j -> i j', seq, self.inv_freq) + # Cache fp32 cos/sin, cast only at use — fp16/bf16 training must not + # recompute trig on low-precision angles. + seq = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.einsum('i, j -> i j', seq, self.inv_freq.float()) if self.interleaved: - freqs = repeat(freqs, '... n -> ... (n r)', r=2) + freqs = torch.repeat_interleave(freqs, 2, dim=-1) else: freqs = torch.cat((freqs, freqs), dim=-1) - return freqs + return torch.cos(freqs), torch.sin(freqs) - def forward(self, seq_len: int) -> Tensor: + def forward(self, seq_len: int): if seq_len > self.cached_freqs_seq_len: raise RuntimeError("sequence exceeds RoPE max_seq_len!") - return self.cached_freqs[0: seq_len].detach() + return self.cached_cos[0: seq_len].detach(), self.cached_sin[0: seq_len].detach() def rotate_queries_or_keys(self, t: Tensor) -> Tensor: device, dtype, seq_len = t.device, t.dtype, t.shape[-2] - freqs = self.forward(seq_len=seq_len) - - return apply_rotary_emb(freqs.to(device=device, dtype=dtype), t, self.interleaved) + freqs_cos, freqs_sin = self.forward(seq_len=seq_len) + freqs_cos = freqs_cos.to(device=device, dtype=dtype) + freqs_sin = freqs_sin.to(device=device, dtype=dtype) + return apply_rotary_emb(freqs_cos, freqs_sin, t, self.interleaved) diff --git a/modules/fastspeech/acoustic_encoder.py b/modules/fastspeech/acoustic_encoder.py index 241f9871c..b960c5a25 100644 --- a/modules/fastspeech/acoustic_encoder.py +++ b/modules/fastspeech/acoustic_encoder.py @@ -58,6 +58,7 @@ def __init__(self, vocab_size): dropout=hparams['dropout'], num_heads=hparams['num_heads'], use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False), use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True), + rope_theta=hparams.get('rope_theta', 10000), mix_ln_layer=self.mix_ln_layer ) diff --git a/modules/fastspeech/tts_modules.py b/modules/fastspeech/tts_modules.py index 10f774156..2b1549956 100644 --- a/modules/fastspeech/tts_modules.py +++ b/modules/fastspeech/tts_modules.py @@ -373,7 +373,7 @@ def __init__( self, hidden_size, num_layers, ffn_kernel_size=9, ffn_act='gelu', dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True, - use_rope=False, rope_interleaved=True, mix_ln_layer=None + use_rope=False, rope_interleaved=True, rope_theta=10000, mix_ln_layer=None ): super().__init__() self.num_layers = num_layers @@ -386,7 +386,9 @@ def __init__( "RoPE requires the hidden size to be multiple of " f"num_heads * 2 = {num_heads * 2}, but got {embed_dim}." ) - rotary_embed = RotaryEmbedding(dim=embed_dim // num_heads, interleaved=rope_interleaved) + rotary_embed = RotaryEmbedding( + dim=embed_dim // num_heads, theta=rope_theta, interleaved=rope_interleaved + ) else: rotary_embed = None self.layers = nn.ModuleList([ diff --git a/modules/fastspeech/variance_encoder.py b/modules/fastspeech/variance_encoder.py index 712964846..6e0aa79a2 100644 --- a/modules/fastspeech/variance_encoder.py +++ b/modules/fastspeech/variance_encoder.py @@ -34,7 +34,8 @@ def __init__(self, vocab_size): ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'], dropout=hparams['dropout'], num_heads=hparams['num_heads'], use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False), - use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True) + use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True), + rope_theta=hparams.get('rope_theta', 10000) ) dur_hparams = hparams['dur_prediction_args'] @@ -128,7 +129,8 @@ def get_hparam(key): ffn_kernel_size=get_hparam('enc_ffn_kernel_size'), ffn_act=get_hparam('ffn_act'), dropout=get_hparam('dropout'), num_heads=get_hparam('num_heads'), use_pos_embed=get_hparam('use_pos_embed'), rel_pos=get_hparam('rel_pos'), - use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True) + use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True), + rope_theta=hparams.get('rope_theta', 10000) ) self.out_proj = Linear(hidden_size, hparams['hidden_size']) diff --git a/requirements.txt b/requirements.txt index 4645417f7..1a98f7951 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,6 @@ # See instructions at https://pytorch.org/get-started/locally/ click -einops>=0.7.0 h5py librosa<0.10.0 lightning~=2.3.0