| """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) |
| |
| return F.linear(hidden,base.lm_head.weight.index_select(0,labels))[0].float() |
|
|
|
|
