Ensure that the other ranks do not hit EOS

mosaicml · Oct 7, 2023 · e758d7d · e758d7d
1 parent 655fc75
commit e758d7d
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_hf_mpt_gen.py b/tests/test_hf_mpt_gen.py
@@ -102,9 +102,13 @@ def forward(
         inputs_embeds: Optional[torch.FloatTensor] = None,
     ):
         result = super().forward(input_ids, past_key_values, attention_mask, prefix_mask, sequence_id, labels, return_dict, output_attentions, output_hidden_states, use_cache, inputs_embeds)
-        # Rank 0 should hit EOS immediately.
+        # Modify the logits to select the next token.
         if dist.get_global_rank() == 0:
+            # Rank 0 hits EOS immediately.
             result.logits[:, :, EOS_TOKEN_ID] = torch.inf
+        else:
+            # Other ranks do not hit EOS.
+            result.logits[:, :, EOS_TOKEN_ID] = -torch.inf
         return result
 
 def mock_from_config(config: MPTConfig, **_):