cpu installation

Author: carmocca

.github/workflows/cpu-tests.yml

@@ -40,6 +40,9 @@ jobs:
 
     - name: Run tests without the package installed
       run: |
+        # install torch cpu nightly
+        python -c "with open('requirements.txt', 'r+') as fp: c = fp.read().replace('cu118', 'cpu'); fp.seek(0); fp.write(c); fp.truncate()"
+
         pip install pytest -r requirements.txt 'transformers==4.27.3'
         pip list
 

lit_stablelm/model.py

@@ -146,7 +146,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
 
         # efficient attention using Flash Attention CUDA kernels
-        y = F.scaled_dot_product_attention(q, k, v.float(), attn_mask=None, dropout_p=0.0, is_causal=True, scale=1.0 / math.sqrt(head_size))
+        y = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, scale=1.0 / math.sqrt(head_size)
+        )
 
         y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side