adding depformer norm (#379)

Author: adefossez

moshi/moshi/models/lm.py

@@ -95,6 +95,7 @@ def __init__(
         depformer_weights_per_step_schedule: list[int] | None = None,
         depformer_low_rank_embeddings: int | None = None,
         depformer_pos_emb: str = "sin",
+        depformer_norm: str | None = None,
         existing_text_padding_id: int = 3,
         existing_text_end_padding_id: int = 0,
         extra_heads_num_heads: int = 0,
@@ -193,6 +194,11 @@ def __init__(
                 depformer_dim,
                 demux_second_stream=demux_second_text_stream,
             )
+            if depformer_norm is None:
+                self.depformer_norms = nn.ModuleList([nn.Identity() for _ in range(dep_q)])
+            else:
+                self.depformer_norms = nn.ModuleList(
+                    [create_norm_fn(depformer_norm, depformer_dim) for _ in range(dep_q)])
             if depformer_dim_feedforward is None:
                 depformer_dim_feedforward = int(hidden_scale * depformer_dim)
             self.depformer = StreamingTransformer(
@@ -435,7 +441,7 @@ def forward_depformer_training(
         depformer_output = self.depformer(depformer_input)
         all_logits = []
         for cb_index in range(Ka):
-            logits = self.linears[cb_index](depformer_output[:, cb_index])
+            logits = self.linears[cb_index](self.depformer_norms[cb_index](depformer_output[:, cb_index]))
             all_logits.append(logits.view(B, T, -1))
         logits = torch.stack(all_logits, 1)
         assert logits.dim() == 4, logits.shape  # [B, Ka, T, card]
@@ -481,7 +487,7 @@ def forward_depformer(
         # depformer_input is [B, 1, depformer_dim].
         # The streaming state of the depformer ensures that the proper layer is run.
         dep_output = self.depformer(depformer_input)
-        logits = self.linears[depformer_cb_index](dep_output)
+        logits = self.linears[depformer_cb_index](self.depformer_norms[depformer_cb_index](dep_output))
         logits = logits[:, None]
         assert logits.dim() == 4, logits.shape  # [B, Ka, S, card]
         return logits

scripts/import_pytorch.py

@@ -43,7 +43,7 @@ def import_model(
         'dim', 'text_card', 'existing_text_padding_id', 'num_heads', 'num_layers', 'hidden_scale', 'causal',
         'layer_scale', 'context', 'max_period', 'gating', 'norm', 'positional_embedding',
         'depformer_dim', 'depformer_num_heads', 'depformer_num_layers', 'depformer_dim_feedforward',
-        'depformer_layer_scale', 'depformer_multi_linear',
+        'depformer_layer_scale', 'depformer_multi_linear', 'depformer_norm',
         'depformer_max_period', 'depformer_gating', 'depformer_pos_emb', 'depformer_weights_per_step',
         'depformer_low_rank_embeddings', 'demux_second_stream',
         'text_card_out']