diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index d5a69aa67f..1ece87c080 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -45,9 +45,10 @@ pass def geglu_exact_forward_kernel(gate, up): batch, seq_len, hd = gate.shape n_elements = gate.numel() - out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda:0") + device = gate.device + out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = device) grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - with torch_cuda_device(gate.device): + with torch_cuda_device(device): _exact_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) return out pass @@ -139,9 +140,10 @@ pass def geglu_approx_forward_kernel(gate, up): batch, seq_len, hd = gate.shape n_elements = gate.numel() - out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda:0") + device = gate.device + out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = device) grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - with torch_cuda_device(gate.device): + with torch_cuda_device(device): _approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) return out pass diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index eb3a2e38cc..2c4edf334b 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -460,7 +460,9 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None): else: reshape = False pass - + + if X.device != W.device: + print(X.device, W.device, torch.cuda.current_device()) out = torch_matmul(X, W, out = out) if W_quant is not None: del W diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index fe0627f8d7..7f475869c8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -385,6 +385,7 @@ def LlamaAttention_fast_forward( head_dim = self.head_dim assert(n_kv_heads * n_groups == n_heads) + print(hidden_states.device, torch.cuda.current_device()) Q, K, V = self.apply_qkv(self, hidden_states) Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2) K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)