File size: 1,837 Bytes
e5b8f61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """CUDA BF16 label readout for the merged Jet model."""
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, Qwen3_5ForCausalLM
def install_linear_attention():
"""Use FLA's differentiable chunk kernel; leave convolution on native PyTorch."""
from fla.ops.gated_delta_rule import chunk_gated_delta_rule
from transformers.models.qwen3_5 import modeling_qwen3_5 as impl
def chunk(query, key, value, g, beta, chunk_size=64, initial_state=None,
output_final_state=False, use_qk_l2norm_in_kernel=False, **kwargs):
return chunk_gated_delta_rule(q=query,k=key,v=value,g=g,beta=beta,
initial_state=initial_state,output_final_state=output_final_state,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel)
impl.torch_chunk_gated_delta_rule = chunk
def load_model(path):
install_linear_attention()
tokenizer = AutoTokenizer.from_pretrained(path)
model, info = Qwen3_5ForCausalLM.from_pretrained(path, dtype=torch.bfloat16,
device_map={'': 'cuda'}, attn_implementation='sdpa', output_loading_info=True)
if info.get('missing_keys') or info.get('mismatched_keys') or info.get('unexpected_keys'):
raise RuntimeError(f'Incomplete model load: {info}')
model.config.use_cache = False
model.eval()
return model, tokenizer
def label_logits(model, example):
base=model.get_base_model() if hasattr(model,'get_base_model') else model
ids=torch.tensor([example['ids']],device='cuda',dtype=torch.long)
hidden=base.model(input_ids=ids,use_cache=False).last_hidden_state[:, -1, :]
labels=torch.tensor(example['labels'],device='cuda',dtype=torch.long)
# Only materialize the requested rows of the frozen vocabulary head.
return F.linear(hidden,base.lm_head.weight.index_select(0,labels))[0].float()
|