NeoX style RoPE (#277)

* refactor RoPE

refactor RoPE

* NeoX style RoPE

* fix export ONNX model before RoPE refactor

Author: KakaruHayate

configs/acoustic.yaml

@@ -64,6 +64,7 @@ timesteps: 1000
 max_beta: 0.02
 enc_ffn_kernel_size: 3
 use_rope: true
+rope_interleaved: false
 use_stretch_embed: true
 use_variance_scaling: true
 rel_pos: true

configs/templates/config_acoustic.yaml

@@ -71,6 +71,7 @@ augmentation_args:
 diffusion_type: reflow
 enc_ffn_kernel_size: 3
 use_rope: true
+rope_interleaved: false
 use_stretch_embed: true
 use_variance_scaling: true
 use_shallow_diffusion: true

configs/templates/config_variance.yaml

@@ -65,6 +65,7 @@ tension_logit_max: 10.0
 
 enc_ffn_kernel_size: 3
 use_rope: true
+rope_interleaved: false
 use_stretch_embed: false
 use_variance_scaling: true
 hidden_size: 384

configs/variance.yaml

@@ -36,6 +36,7 @@ predict_tension: false
 
 enc_ffn_kernel_size: 3
 use_rope: true
+rope_interleaved: false
 use_stretch_embed: false
 use_variance_scaling: true
 rel_pos: true

deployment/exporters/acoustic_exporter.py

@@ -1,6 +1,7 @@
 import json
 from pathlib import Path
 from typing import List, Union, Tuple, Dict
+import warnings
 
 import onnx
 import onnxsim
@@ -78,6 +79,7 @@ def __init__(
                     self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
             if self.freeze_spk is not None:
                 self.model.fs2.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
+        self.rope_interleaved = hparams.get('rope_interleaved', None)
 
     def build_model(self) -> DiffSingerAcousticONNX:
         model = DiffSingerAcousticONNX(
@@ -88,8 +90,21 @@ def build_model(self) -> DiffSingerAcousticONNX:
                 for p in self.phoneme_dictionary.cross_lingual_phonemes
             })
         ).eval().to(self.device)
+        if self.rope_interleaved is None:
+            warnings.warn(
+                "After RoPE is refactored, the checkpoint no longer contains relevant parameters. "
+                "(https://github.com/openvpi/DiffSinger/pull/276)"
+                "In order to export ONNX with behavior compatible with past checkpoints, "
+                "it will be set to 'strict=False', which will no longer check the validity of the checkpoint. "
+                "Please understand what you are doing.",
+                UserWarning,
+                stacklevel=2
+            )
+            strict=False
+        else:
+            strict=True
         load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
-                  prefix_in_ckpt='model', strict=True, device=self.device)
+                  prefix_in_ckpt='model', strict=strict, device=self.device)
         return model
 
     def export(self, path: Path):

deployment/exporters/variance_exporter.py

@@ -1,6 +1,7 @@
 import json
 from pathlib import Path
 from typing import Union, List, Tuple, Dict
+import warnings
 
 import onnx
 import onnxsim
@@ -81,6 +82,7 @@ def __init__(
                     self.export_spk = [(name, {name: 1.0}) for name in self.spk_map.keys()]
             if self.freeze_spk is not None:
                 self.model.register_buffer('frozen_spk_embed', self._perform_spk_mix(self.freeze_spk[1]))
+        self.rope_interleaved = hparams.get('rope_interleaved', None)
 
     def build_model(self) -> DiffSingerVarianceONNX:
         model = DiffSingerVarianceONNX(
@@ -90,6 +92,19 @@ def build_model(self) -> DiffSingerVarianceONNX:
                 for p in self.phoneme_dictionary.cross_lingual_phonemes
             })
         ).eval().to(self.device)
+        if self.rope_interleaved is None:
+            warnings.warn(
+                "After RoPE is refactored, the checkpoint no longer contains relevant parameters. "
+                "(https://github.com/openvpi/DiffSinger/pull/276)"
+                "In order to export ONNX with behavior compatible with past checkpoints, "
+                "it will be set to 'strict=False', which will no longer check the validity of the checkpoint. "
+                "Please understand what you are doing.",
+                UserWarning,
+                stacklevel=2
+            )
+            strict=False
+        else:
+            strict=True
         load_ckpt(model, hparams['work_dir'], ckpt_steps=self.ckpt_steps,
                   prefix_in_ckpt='model', strict=True, device=self.device)
         model.build_smooth_op(self.device)

modules/fastspeech/acoustic_encoder.py

@@ -38,7 +38,7 @@ 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)
+            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True)
         )
 
         self.pitch_embed = Linear(1, hparams['hidden_size'])

modules/fastspeech/tts_modules.py

@@ -369,14 +369,14 @@ def mel2ph_to_dur(mel2ph, T_txt, max_dur=None):
 class FastSpeech2Encoder(nn.Module):
     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):
+                 dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True, use_rope=False, rope_interleaved=True):
         super().__init__()
         self.num_layers = num_layers
         embed_dim = self.hidden_size = hidden_size
         self.dropout = dropout
         self.use_pos_embed = use_pos_embed
         if use_pos_embed and use_rope:
-            rotary_embed = RotaryEmbedding(dim = embed_dim // num_heads)
+            rotary_embed = RotaryEmbedding(dim = embed_dim // num_heads, interleaved = rope_interleaved)
         else:
             rotary_embed = None
         self.layers = nn.ModuleList([

modules/fastspeech/variance_encoder.py

@@ -33,7 +33,7 @@ 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)
+            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True)
         )
 
         dur_hparams = hparams['dur_prediction_args']
@@ -127,7 +127,7 @@ 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')
+            use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True)
         )
         self.out_proj = Linear(hidden_size, hparams['hidden_size'])