File size: 3,482 Bytes
c5c8412 | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | # Modified implementation for Kev; see LICENSE and NOTICE.
"""Qwen3/Llama backbone adapter: the parent's decoder (all layers, final norm, embedding matrix as `embed_tokens`; no
lm_head, no vocabulary projection), the attention backend, and prefix-cache replication for batched question branches.
Attention backend (verified against transformers 5.17.0 source): Qwen3Model hands its attention_mask to
masking_utils.create_causal_mask, whose _preprocess_mask_arguments returns any 4D tensor unchanged ("If the mask is
already 4D, simply return as-is"); integrations/sdpa_attention.py then passes it to
torch.nn.functional.scaled_dot_product_attention as `attn_mask` (is_causal is forced off when a mask is given), and the
eager path adds it to the attention scores. So the packed form's additive float block-causal mask is honoured by both:
SDPA on CUDA, eager elsewhere (the default). Row and cached passes give a 2D padding mask and let transformers build
the causal mask; passing it (or a cache) also disables the position_ids "packed sequence" detection, which only runs
when both are None.
The Llama path has also been exercised on Nesso2-0.4B with Transformers 5.17.0
and ROCm SDPA: packed/rows/cached parity, isolation, gradients and save/reload.
See experiments/lumi/README.md for the measured scope and limitations.
"""
from pathlib import Path
import torch
from transformers import AutoModel, DynamicCache
def default_attn(device) -> str:
return "sdpa" if str(device).startswith("cuda") else "eager"
def load_backbone(name_or_path, *, revision=None, dtype=torch.float32, attn=None, device="cpu", local_files_only=False):
"""Supported decoder from a parent or artifact. Reject missing backbone weights (silent random initialization)."""
lm, info = AutoModel.from_pretrained(str(name_or_path), revision=revision, dtype=dtype, attn_implementation=attn or default_attn(device),
local_files_only=local_files_only, output_loading_info=True)
supported = {"qwen3": "Qwen3Model", "llama": "LlamaModel"}
if supported.get(lm.config.model_type) != type(lm).__name__:
raise ValueError(f"{name_or_path}: expected a Qwen3 or Llama decoder, got {type(lm).__name__} ({lm.config.model_type})")
if info["missing_keys"]:
raise ValueError(f"{name_or_path}: {len(info['missing_keys'])} backbone weights missing from the checkpoint, e.g. {sorted(info['missing_keys'])[:3]}")
return lm.to(device)
def save_backbone(lm, out_dir, dtype):
"""Cast the live merged decoder's parameters to `dtype` and write it as an HF directory."""
cast_parameters(lm, dtype)
lm.save_pretrained(Path(out_dir))
def cast_parameters(module, dtype):
"""Cast parameters only. nn.Module.to(dtype) would also cast the rotary `inv_freq` buffer, which from_pretrained keeps
in fp32 at any load dtype; a bf16 inv_freq moves bf16 probabilities by ~2e-2 on identical weights (measured)."""
for p in module.parameters():
p.data = p.data.to(dtype)
return module
def replicate_cache(cache: DynamicCache, n: int, config) -> DynamicCache:
"""A new DynamicCache holding `n` batch copies of a batch-1 prefix cache. DynamicLayer.update concatenates into new
tensors, so running branches on the replica never writes to `cache` (the caller's prefix stays reusable)."""
return DynamicCache(ddp_cache_data=[(k.expand(n, -1, -1, -1), v.expand(n, -1, -1, -1)) for k, v, *_ in cache], config=config)
|