| """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()) |
| |
| |
| 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"): |
| |
| 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 |
|
|
