support bf16 training and inference

Author: Paulmzr

scripts/eval_continue.py

@@ -147,12 +147,12 @@ def main():
             audio_inputs = processor(raw_audio=audio_arrays, sampling_rate=SAMPLING_RATE, return_tensors="pt") # 'padding_mask': b,t  'input_values': b,c,t
 
 
-            encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(model.dtype).to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
+            encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
             speech_inputs_embeds = model.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t
 
             speech_attention_mask = audio_inputs["padding_mask"][..., ::STRIDE].to(device)
             assert speech_inputs_embeds.size(-1) == speech_attention_mask.size(-1)
-            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2) #b,t,d
+            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2).to(model.dtype) #b,t,d
 
 
             speech_inputs_embeds = speech_inputs_embeds[:,:FREQ * PROMPT_LEN,:]
@@ -177,7 +177,7 @@ def main():
             new_embeds = output_sequences[1]
             generated_ids = output_sequences[0][:, text_input_length:]
 
-            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
             wav_len = (generated_ids.ne(eos_token_id).sum(dim=-1) + speech_input_length) * STRIDE
 

scripts/eval_reference.py

@@ -169,12 +169,12 @@ def main():
             audio_inputs = processor(raw_audio=audio_arrays, sampling_rate=SAMPLING_RATE, return_tensors="pt") # 'padding_mask': b,t  'input_values': b,c,t
 
 
-            encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(model.dtype).to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
+            encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
             speech_inputs_embeds = model.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t
 
             speech_attention_mask = audio_inputs["padding_mask"][..., ::STRIDE].to(device)
             assert speech_inputs_embeds.size(-1) == speech_attention_mask.size(-1)
-            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2) #b,t,d
+            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2).to(model.dtype) #b,t,d
 
 
             speech_input_length = speech_inputs_embeds.shape[1]
@@ -197,7 +197,7 @@ def main():
             new_embeds = output_sequences[1]
             generated_ids = output_sequences[0][:, text_input_length:]
 
-            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
             wav_len = (generated_ids.ne(eos_token_id).sum(dim=-1) + speech_input_length) * STRIDE
 

scripts/eval_stream.py

@@ -160,7 +160,7 @@ def main():
             new_embeds = output_sequences[1]
             generated_ids = output_sequences[0]
 
-            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+            new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
             wav_len = (generated_ids.ne(eos_token_id).sum(dim=-1)) * STRIDE
 

scripts/run_offline.py

@@ -106,7 +106,7 @@ def main():
 
 
     new_embeds = output_sequences[1]
-    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
 
     output_path = "output.wav"    

scripts/run_stream.py

@@ -106,7 +106,7 @@ def main():
 
 
     new_embeds = output_sequences[1]
-    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
 
     output_path = "output.wav"    

scripts/run_voice_clone.py

@@ -125,12 +125,12 @@ def main():
     audio_inputs = processor(raw_audio=audio_arrays, sampling_rate=SAMPLING_RATE, return_tensors="pt") # 'padding_mask': b,t  'input_values': b,c,t
 
 
-    encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(model.dtype).to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
+    encoder_outputs = model.codec.encode(audio_inputs["input_values"].to(device), audio_inputs["padding_mask"].to(device), bandwidth=BANDWIDTH) #1,b,r,t, 1 due to one chunk
     speech_inputs_embeds = model.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t
 
     speech_attention_mask = audio_inputs["padding_mask"][..., ::STRIDE].to(device)
     assert speech_inputs_embeds.size(-1) == speech_attention_mask.size(-1)
-    speech_inputs_embeds = speech_inputs_embeds.transpose(1,2) #b,t,d
+    speech_inputs_embeds = speech_inputs_embeds.transpose(1,2).to(model.dtype) #b,t,d
 
     speech_input_length = speech_inputs_embeds.shape[1]
 
@@ -149,7 +149,7 @@ def main():
 
 
     new_embeds = output_sequences[1]
-    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2))
+    new_audio_values = model.codec.decoder(new_embeds.transpose(-1,-2).float())
 
 
     output_path = "output.wav"    

sled/energy_distance.py

@@ -59,7 +59,7 @@ def sample(self, z, temperature=1.0, cfg=1.0):
             z = z_1 * cfg + (1 - cfg) * z_2
 
         noise = torch.randn((z.shape[0], self.noise_channels), dtype=z.dtype, device=z.device)
-        return self.net(noise, z).float()
+        return self.net(noise, z)
 
 
 

sled/sled.py

@@ -70,7 +70,7 @@ def __init__(self, config):
             target_channels=self.token_embed_dim,
             z_channels=self.hidden_size,
             width=config.diffloss_w,
-            depth=config.diffloss_d, 
+            depth=config.diffloss_d,
             noise_channels=config.noise_channels,
         )
 
@@ -85,9 +85,9 @@ def __init__(self, config):
 
     def initialize_codec(self, model_args):
         if hasattr(model_args, "codec"):
-            self.codec = EncodecModel.from_pretrained(model_args.codec, torch_dtype=self.model.dtype)
+            self.codec = EncodecModel.from_pretrained(model_args.codec, torch_dtype=torch.float32) # keep encodec model in fp32
         else:
-            self.codec = EncodecModel.from_pretrained(model_args, torch_dtype=self.model.dtype)
+            self.codec = EncodecModel.from_pretrained(model_args, torch_dtype=torch.float32)
         for param in self.codec.parameters():
             param.requires_grad = False
 
@@ -124,12 +124,12 @@ def prepare_inputs_labels_for_multimodal(
         text_inputs_embeds = self.model.embed_tokens(input_ids)
 
         with torch.no_grad():
-            encoder_outputs = self.codec.encode(audio_inputs["input_values"].to(self.model.dtype), audio_inputs["padding_mask"], bandwidth=6) #1,b,r,t, 1 due to one chunk
-            speech_inputs_embeds = self.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t
+            encoder_outputs = self.codec.encode(audio_inputs["input_values"], audio_inputs["padding_mask"], bandwidth=6) #1,b,r,t, 1 due to one chunk
+            speech_inputs_embeds = self.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t, always fp32
 
             speech_attention_mask = audio_inputs["padding_mask"][..., ::320]
             assert speech_inputs_embeds.size(-1) == speech_attention_mask.size(-1)
-            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2) #b,t,d
+            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2).to(self.model.dtype) #b,t,d, support full bf16 training
 
         net_speech_inputs_embeds = self.z_proj(speech_inputs_embeds)
         new_inputs_embeds = torch.concat([text_inputs_embeds, net_speech_inputs_embeds], dim=1) #bsz, seq_len, hidden_size
@@ -536,7 +536,7 @@ def _sample(
 
             # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
             # (the clone itself is always small)
-            next_token_logits = outputs.logits.clone()[:, -1, :].float()
+            next_token_logits = outputs.logits.clone()[:, -1, :]
             eos_next_token_logits = next_token_logits
             if self.infer_cfg != 1.0:
                 next_token_logits_normal = next_token_logits[:real_batch_size, :]

sled/sled_stream.py

@@ -190,7 +190,7 @@ def __init__(self, config):
             target_channels=self.token_embed_dim,
             z_channels=self.hidden_size,
             width=config.diffloss_w,
-            depth=config.diffloss_d, 
+            depth=config.diffloss_d,
             noise_channels=config.noise_channels,
         )
 
@@ -207,9 +207,9 @@ def __init__(self, config):
 
     def initialize_codec(self, model_args):
         if hasattr(model_args, "codec"):
-            self.codec = EncodecModel.from_pretrained(model_args.codec, torch_dtype=self.model.dtype)
+            self.codec = EncodecModel.from_pretrained(model_args.codec, torch_dtype=torch.float32) # keep encodec model in fp32
         else:
-            self.codec = EncodecModel.from_pretrained(model_args, torch_dtype=self.model.dtype)
+            self.codec = EncodecModel.from_pretrained(model_args, torch_dtype=torch.float32)
         for param in self.codec.parameters():
             param.requires_grad = False
 
@@ -247,12 +247,12 @@ def prepare_inputs_labels_for_multimodal(
         text_inputs_embeds = self.model.embed_tokens(input_ids)
 
         with torch.no_grad():
-            encoder_outputs = self.codec.encode(audio_inputs["input_values"].to(self.model.dtype), audio_inputs["padding_mask"], bandwidth=6) #1,b,r,t, 1 due to one chunk
-            speech_inputs_embeds = self.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t
+            encoder_outputs = self.codec.encode(audio_inputs["input_values"], audio_inputs["padding_mask"], bandwidth=6) #1,b,r,t, 1 due to one chunk
+            speech_inputs_embeds = self.codec.quantizer.decode(encoder_outputs.audio_codes[0].transpose(0, 1)) #b,d,t, always fp32
 
             speech_attention_mask = audio_inputs["padding_mask"][..., ::320]
             assert speech_inputs_embeds.size(-1) == speech_attention_mask.size(-1)
-            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2) #b,t,d
+            speech_inputs_embeds = speech_inputs_embeds.transpose(1,2).to(self.model.dtype) #b,t,d, support full bf16 training
 
 
         net_speech_inputs_embeds = self.z_proj(speech_inputs_embeds)
@@ -729,10 +729,10 @@ def _sample(
 
                 # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
                 # (the clone itself is always small)
-                next_token_logits = outputs.logits.clone()[:, -1, :].float()
+                next_token_logits = outputs.logits.clone()[:, -1, :]
                 eos_next_token_logits = next_token_logits.clone()
                 if self.infer_cfg != 1.0:
-                    next_token_logits_cfg = outputs_cfg.logits.clone()[:, -1, :].float()
+                    next_token_logits_cfg = outputs_cfg.logits.clone()[:, -1, :]
                     next_token_logits = next_token_logits_cfg + self.infer_cfg * (next_token_logits - next_token_logits_cfg)
                 # Store scores, attentions and hidden_states when required
                 if return_dict_in_generate: