File size: 5,164 Bytes
08c77f0 | 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Load the mixed packed Inkling 2 checkpoint with the native Transformers model."""
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import Inkling2Config, Inkling2ForConditionalGeneration
from transformers.conversion_mapping import get_model_conversion_mapping
from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key
def dequantize(weight, scale, amax=None, denominator=None):
if weight.dtype == torch.uint8:
if amax is None or denominator is None:
raise ValueError("NVFP4 weights require per-expert amax and an outer-scale denominator")
codes = torch.stack((weight & 15, weight >> 4), dim=-1).flatten(-2)
levels = weight.new_tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32)
values = levels[(codes & 7).long()] * torch.where(codes < 8, 1.0, -1.0)
return values * scale.float().repeat_interleave(16, -1) * (amax / denominator).reshape(-1, 1, 1)
if weight.dtype == torch.float8_e4m3fn:
scales = torch.pow(2.0, scale.float() - 127).repeat_interleave(32, -1)
return weight.float() * scales
raise ValueError(f"Unsupported packed expert dtype: {weight.dtype}")
class Inkling2TinyForConditionalGeneration(Inkling2ForConditionalGeneration):
@classmethod
def from_pretrained(
cls, pretrained_model_name_or_path, *model_args, config=None, dtype=torch.bfloat16,
revision=None, cache_dir=None, local_files_only=False, token=None,
attn_implementation="eager", **kwargs,
):
if model_args:
raise ValueError("Inkling 2 does not take positional construction arguments")
options = dict(revision=revision, cache_dir=cache_dir, local_files_only=local_files_only, token=token)
root = Path(pretrained_model_name_or_path)
def resolve(filename):
if root.is_dir():
return root / filename
return Path(hf_hub_download(str(pretrained_model_name_or_path), filename, **options))
import json
raw = json.loads(resolve("config.json").read_text())
# The native config otherwise assumes the all-MXFP8 sibling. This
# adapter dequantizes each tensor according to its actual stored dtype.
raw.pop("quantization_config", None)
config = Inkling2Config(**raw)
config._attn_implementation = attn_implementation
dtype = kwargs.pop("torch_dtype", dtype)
if dtype in (None, "auto"):
dtype = torch.bfloat16
elif isinstance(dtype, str):
dtype = getattr(torch, dtype)
if dtype not in (torch.float32, torch.bfloat16, torch.float16):
raise ValueError("The Transformers adapter requires a floating compute dtype")
device_map = kwargs.pop("device_map", None)
if device_map == "auto":
device_map = "cuda" if torch.cuda.is_available() else "cpu"
if isinstance(device_map, dict):
if set(device_map) != {""}:
raise ValueError("Use a single device for this tiny Transformers adapter")
device_map = device_map[""]
allowed = {"_from_auto", "_from_pipeline", "trust_remote_code", "weights_only", "adapter_kwargs", "_commit_hash"}
unexpected = set(kwargs) - allowed
if unexpected:
raise TypeError(f"Unsupported loading arguments: {sorted(unexpected)}")
model = Inkling2ForConditionalGeneration(config)
transforms = get_model_conversion_mapping(model)
renamings = [t for t in transforms if isinstance(t, WeightRenaming)]
converters = [t for t in transforms if isinstance(t, WeightConverter)]
source = load_file(str(resolve("model.safetensors")))
state = {}
for name, weight in source.items():
if name.endswith(("_scale", "_amax", "_outer_scale_denominator")):
continue
if name + "_scale" in source:
weight = dequantize(
weight, source[name + "_scale"], source.get(name + "_amax"),
source.get(name + "_outer_scale_denominator"),
)
target, _ = rename_source_key(name, renamings, converters)
if name.endswith("w13_weight"):
# The source interleaves gate and up; Transformers chunks them.
e, rows, columns = weight.shape
weight = weight.reshape(e, rows // 2, 2, columns).transpose(1, 2).reshape(e, rows, columns)
if target in state:
raise ValueError(f"Checkpoint keys map to the same parameter: {target}")
keep_float32 = source[name].dtype == torch.float32 or any(
part in target for part in model._keep_in_fp32_modules_strict
)
target_dtype = torch.float32 if keep_float32 else dtype
state[target] = weight.to(target_dtype).contiguous()
model.load_state_dict(state, strict=True, assign=True)
model.eval()
if device_map is not None:
model.to(device_map)
return model
|