Fix device mismatch when running hf.generate (#1486)

llmfoundry/models/layers/blocks.py

@@ -204,7 +204,8 @@ def forward(
         n = self.apply_ffn(attention_mask, m)
         # In the following line we move the `x` tensor to the same devices as the output of ffn layer. This operation should be a no-op during training.
         # This is done to fix pipeline parallel generation using hf.generate. Please see this comment for details: https://github.com/mosaicml/llm-foundry/pull/1332#issue-2386827204
-        x = x.to(device=n.device) + self.resid_ffn_dropout(n)
+        x = x.to(device=n.device,
+                ) + self.resid_ffn_dropout(n).to(device=n.device,)
         return x, attn_weights, past_key_value
 
     def apply_ffn(

llmfoundry/models/layers/ffn.py

@@ -214,7 +214,9 @@ def __init__(
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(
+            self.act(self.gate_proj(x)).to(device=x.device) * self.up_proj(x),
+        )
 
 
 def build_mptglu(

llmfoundry/models/mpt/modeling_mpt.py

@@ -161,10 +161,7 @@ def gen_rotary_embedding(
             hidden_size=d_model,
             num_attention_heads=n_heads,
         )
-        if rope_hf_config['type'] == 'no_scaling':
-            return LlamaRotaryEmbeddingFoundry(config=partial_llama_config)
-        elif rope_hf_config['type'] in {'llama3', 'linear', 'dynamic'}:
-            return LlamaRotaryEmbedding(config=partial_llama_config)
+        return LlamaRotaryEmbeddingFoundry(config=partial_llama_config)
     raise ValueError('rope_impl needs to be either dail or hf')