glm 4v 1st sdp for vision (#12904)
* glm4v 1st sdp * update glm4v example * meet code review * fix style
This commit is contained in:
		
							parent
							
								
									5c100ac105
								
							
						
					
					
						commit
						e946127613
					
				
					 2 changed files with 14 additions and 22 deletions
				
			
		| 
						 | 
					@ -61,7 +61,7 @@ if __name__ == '__main__':
 | 
				
			||||||
                                                 trust_remote_code=True,
 | 
					                                                 trust_remote_code=True,
 | 
				
			||||||
                                                 use_cache=True,
 | 
					                                                 use_cache=True,
 | 
				
			||||||
                                                 model_hub=model_hub)
 | 
					                                                 model_hub=model_hub)
 | 
				
			||||||
    model = model.half().to('xpu')
 | 
					    model = model.to('xpu')
 | 
				
			||||||
    
 | 
					    
 | 
				
			||||||
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 | 
					    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -19,7 +19,7 @@
 | 
				
			||||||
 | 
					
 | 
				
			||||||
import torch
 | 
					import torch
 | 
				
			||||||
from typing import Optional, Tuple, Union
 | 
					from typing import Optional, Tuple, Union
 | 
				
			||||||
from ipex_llm.transformers.models.common import merge_qkv_base
 | 
					from ipex_llm.transformers.models.common import merge_qkv_base, padding_qkv_hd
 | 
				
			||||||
from ipex_llm.transformers.models.common import scaled_dot_product_attention
 | 
					from ipex_llm.transformers.models.common import scaled_dot_product_attention
 | 
				
			||||||
from ipex_llm.transformers.models.utils import update_past_key_value
 | 
					from ipex_llm.transformers.models.utils import update_past_key_value
 | 
				
			||||||
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp
 | 
					from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp
 | 
				
			||||||
| 
						 | 
					@ -265,26 +265,18 @@ def visual_attention_forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
 | 
				
			||||||
    q, k, v = qkv[0], qkv[1], qkv[2]
 | 
					    q, k, v = qkv[0], qkv[1], qkv[2]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    bsz, q_len, kv_seq_len, head_dim = q.shape
 | 
					    bsz, q_len, kv_seq_len, head_dim = q.shape
 | 
				
			||||||
    if use_sdp(q_len, kv_seq_len, head_dim, q):
 | 
					    q, k, v = padding_qkv_hd(
 | 
				
			||||||
        import xe_addons
 | 
					        q, k, v,
 | 
				
			||||||
        out = xe_addons.sdp(q, k, v, None)
 | 
					        head_dim, 128
 | 
				
			||||||
    elif q.device.type == "cpu":
 | 
					    )
 | 
				
			||||||
        out = torch.nn.functional.scaled_dot_product_attention(q, k, v,
 | 
					
 | 
				
			||||||
                                                               attn_mask=None,
 | 
					    attn_weights = None
 | 
				
			||||||
                                                               dropout_p=0.,
 | 
					    attn_output = scaled_dot_product_attention(
 | 
				
			||||||
                                                               is_causal=False)
 | 
					        q, k.contiguous(), v.contiguous(),
 | 
				
			||||||
    else:
 | 
					        None, False, 1 / math.sqrt(head_dim)
 | 
				
			||||||
        attn_weights = torch.matmul(q / math.sqrt(head_dim),
 | 
					    )
 | 
				
			||||||
                                    k.transpose(2, 3)).to(v.dtype)
 | 
					
 | 
				
			||||||
        if kv_seq_len >= 2048 or bsz >= 64:
 | 
					    out = attn_output[:, :, :, :head_dim]
 | 
				
			||||||
            # for memory considerations, do not upcast attention to fp32
 | 
					 | 
				
			||||||
            # for long sequences or large batches
 | 
					 | 
				
			||||||
            attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
 | 
					 | 
				
			||||||
        else:
 | 
					 | 
				
			||||||
            # upcast attention to fp32
 | 
					 | 
				
			||||||
            attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1,
 | 
					 | 
				
			||||||
                                                       dtype=torch.float32).to(v.dtype)
 | 
					 | 
				
			||||||
        out = torch.matmul(attn_weights, v)
 | 
					 | 
				
			||||||
    output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
 | 
					    output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
 | 
				
			||||||
    output = self.output_dropout(output)
 | 
					    output = self.output_dropout(output)
 | 
				
			||||||
    return output
 | 
					    return output
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
		Reference in a new issue