Update modeling_quiet.py

modeling_quiet.py

@@ -1836,7 +1836,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 elif ahead_idx >= self.n_ahead - 1:
                     if labels is not None:  # we're in the talk phase
                         cur_talk_n = ahead_idx - (self.n_ahead - 1) + 1
-                        # print("Setting rm to labels", cur_talk_n, "during", ahead_idx)
                         shift_labels = labels[..., cur_talk_n:].contiguous().to(probabilities_2d.device)
                         padding = torch.full_like(
                             labels[..., :cur_talk_n],
@@ -1848,44 +1847,50 @@ class QuietForCausalLM(QuietPreTrainedModel):
                             [shift_labels, padding],
                             dim=-1
                         )
-                        
-                        # print((new_rm_tokens > self.vocab_size - 1).any().item())
                         new_rm_tokens = torch.clamp(new_rm_tokens, 0, self.vocab_size - 1)
-
-                        # Now safely convert rm tokens to one-hot
                         probabilities_2d = F.one_hot(new_rm_tokens, num_classes=self.vocab_size).reshape(-1, self.vocab_size).to(probabilities_2d.dtype)
+                        skip_sampling = True
                     else:
                         continue
+
                 temperature = self.gumbel_temperature if self.training else 0.001
                 prev_sample_probs = sample_probs
                 sample_probs = probabilities_2d
+
                 if ahead_idx < self.n_ahead - 1 and not skip_sampling:
                     probabilities_2d = F.gumbel_softmax(sample_probs, tau=temperature, hard=True, dim=-1)
                     if self.gumbel_detach:
                         probabilities_2d = probabilities_2d.detach()
-                sampled_token_history.append(probabilities_2d.argmax(dim=-1).detach().cpu())
+                    sampled_token_history.append(probabilities_2d.argmax(dim=-1).detach().cpu())
+
                 # convert rm logits directly to embeddings
                 contains_start = self.use_start_thought_token and (probabilities_2d[..., self.start_token_id].sum() > 0)
                 contains_end = self.use_end_thought_token and (probabilities_2d[..., self.end_token_id].sum() > 0)
                 contains_thought = contains_start or contains_end
 
-                if not contains_thought:
-                    with torch.set_grad_enabled(not self.train_only_thinking_embedding):
-                        inputs_embeds = probabilities_2d @ (self.model.embed_tokens.weight.to(probabilities.device).to(probabilities.dtype))
-                else:
-                    thought_id = self.start_token_id if contains_start else self.end_token_id
-                    cur_thought_embedding = start_embedding if contains_start else end_embedding
-                    if self.use_reparam_for_thought_embeddings:
-                        inputs_embeds = torch.randn(batch_size, seq_len, self.model.config.hidden_size, device=input_ids.device, dtype=cur_thought_embedding.dtype)
-                        inputs_embeds = inputs_embeds * torch.exp(cur_thought_embedding[1]) + cur_thought_embedding[0]
-                        if contains_start:
-                            sampled_start = inputs_embeds.clone().detach()
+                # Flash Attention modification
+                if self._attn_implementation == "flash_attention_2":
+                    probabilities_2d = probabilities_2d.view(batch_size, seq_len, -1)
+                    
+                    if contains_thought:
+                        thought_id = self.start_token_id if contains_start else self.end_token_id
+                        cur_thought_embedding = start_embedding if contains_start else end_embedding
+                        
+                        if self.use_reparam_for_thought_embeddings:
+                            inputs_embeds = torch.randn(batch_size, seq_len, self.model.config.hidden_size, device=input_ids.device, dtype=cur_thought_embedding.dtype)
+                            inputs_embeds = inputs_embeds * torch.exp(cur_thought_embedding[1]) + cur_thought_embedding[0]
+                            if contains_start:
+                                sampled_start = inputs_embeds.clone().detach()
+                            else:
+                                sampled_end = inputs_embeds.clone().detach()
                         else:
-                            sampled_end = inputs_embeds.clone().detach()
+                            inputs_embeds = cur_thought_embedding.unsqueeze(0).repeat(batch_size, seq_len, 1)
+                            inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
                     else:
-                        inputs_embeds = cur_thought_embedding.unsqueeze(0).repeat(batch_size, seq_len, 1)
-                        inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
-                inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
+                        with torch.set_grad_enabled(not self.train_only_thinking_embedding):
+                            inputs_embeds = probabilities_2d @ (self.model.embed_tokens.weight.to(probabilities.device).to(probabilities.dtype))
+                            
+                    inputs_embeds = inputs_embeds.view(probabilities.size(0), probabilities.size(1), -1).to(self.model.embed_tokens.weight.dtype)
 
                 if len(attention_mask.shape) == 2:
                     breakpoint()