mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-21 14:00:32 +00:00
f7da293881
* Load compressed-tensors packed INT4 checkpoints straight into bitsandbytes 4-bit, and train remote DeepSeek-V3 MoE code A checkpoint published in the compressed-tensors pack-quantized integer format (W4A16 / W8A16, unsloth/Kimi-K2.7-Code among them) could not be trained: the loader kept it packed, PEFT found no weight to attach a LoRA to, and the only route was a 16-bit copy on disk (2 TB for the 595 GB Kimi checkpoint). load_in_4bit=True now decompresses each packed weight while the checkpoint streams in and hands it to the bitsandbytes quantize op in the same loading step, through a WeightConverter the bitsandbytes quantizer subclass adds when the model config carries a plan. Native MoE expert merges keep working (the packed sources are added to the expert converters), packed tensors keep their storage dtype through materialisation, the checkpoint's own quantization config is dropped from the root and every sub-config, and older transformers keep the previous behaviour. The DeepSeek-V3 modeling file Kimi ships only has an inference path (the gate asserts not training, the MoE block computes under no_grad); a training forward is added on the checkpoint's classes, and a remote composition whose inner language model supports gradient checkpointing now declares it on the outer class too. * Hand the planner the armed config of a re-quantized compressed-tensors checkpoint check_and_disable_bitsandbytes_loading drops the checkpoint's quantization config from the config object when it arms the on-the-fly re-quantization, but the device map planner rebuilt the repo's config from the model name, saw compressed-tensors again, and merge_quantization_configs refused the bitsandbytes flags. Kimi-K2.7-Code fell back to sequential, which spilled to CPU, and bitsandbytes refused the load. resolve_unsloth_device_map takes prepared_config and forwards it as the planner's config when this unsloth_zoo's build_meta_model accepts one (read off the signature); an older unsloth_zoo declines the plan as before. The two call sites pass the armed config and nothing otherwise. Two tests: the prepared config is exactly the armed one, and the planner receives it, or is not called at all on an unsloth_zoo that cannot take it. * Arm the packed INT4 re-quantization only when the transformers 4-bit load runs, and stop when the quantizer cannot be installed check_and_disable_bitsandbytes_loading armed the compressed-tensors to bitsandbytes conversion for any 4-bit request. Under fast_inference the weights never go through the transformers loader (vLLM reads the packed checkpoint itself) and under full_finetuning 4-bit is switched off right after the call, so in both cases the checkpoint's own quantization config was stripped with nothing left to consume the plan and the weight_packed tensors loaded as unmatched keys. The call now takes requantize_packed, which the two loaders set from fast_inference and full_finetuning. arm_compressed_tensors_bnb_loading also ignored the False that install_compressed_tensors_bnb_quantizer returns when a foreign class owns the bitsandbytes quantizer slot, and stripped the config anyway; it now leaves the config alone and loads the checkpoint as published. Tests: both; fail on the branch before this commit. * Keep the re-quantization plan on the config until the weights are loaded The device-map planner preprocesses a meta model built from the same config object the load uses, and the quantizer took the plan off the config in that first pass. The real load then found no plan, added no converters, and every packed expert was reported missing, re-initialised and quantized from random values: Kimi-K2.7-Code on four GPUs loaded in four seconds with 69120 expert Linears of noise. A single-GPU load never planned, which is why the tiny fixtures passed. The plan is read in _process_model_before_weight_loading and removed in _process_model_after_weight_loading. Test: two preprocess passes from one config, the second still produces the packed-weight converter. * Rebind accelerate hooks to the remote MoE training shims A device_map load wraps every dispatched module's forward with an accelerate hook that keeps the bound original as _old_forward, so the class-level gate and MoE shims were never reached from module(...) on a multi-GPU model: Kimi-K2.7-Code on four cards loaded and then hit the gate's assert not self.training on the first training step. Single-GPU loads have no hook, which is why the tiny fixture trained. prepare_remote_moe_for_training now rebinds _old_forward on every shimmed instance, including modules whose class was shimmed by an earlier call. Test: a block behind ModelHook raises before, matches eval after, and a second model hooked after the classes were patched is rebound too. * Keep a converter's own ops on their declared source contract, and arm the re-quantization for in-process classification loads The rebuilt expert converter handed its original ops the packed source patterns, and transformers' ErnieFuseAndSplitTextVisionExperts requires every one of its source patterns in the collected dict, so a packed Ernie4.5-VL-MoE checkpoint failed on the consumed scale and shape patterns. Each of the converter's own ops now runs behind an adapter that renames every decompressed bucket to its plain weight pattern and passes the original source_patterns; MergeModulelist keeps the pre-stacked tensor, every other op gets the per-expert list it expects. The vLLM gate for the re-quantization now uses _vllm_will_load_weights, so a num_labels classification load under fast_inference, which transformers still serves, keeps its 4-bit path. Tests: a packed two-module-list bucket through the many-to-many op matches the unpacked reference and the un-adapted op fails; the classification gate. * Add the SPDX line the other test files carry * Decline the packed INT4 plan when compressed-tensors cannot parse it, and drop the retired actorder spelling from the fixture The quantization config was stripped from the model config before transformers parsed the plan, so a checkpoint carrying a field that the installed compressed-tensors has since removed (actorder = "group" left in 0.19.0) died inside from_pretrained with nothing to fall back to. Parse the plan up front and load the checkpoint as published when it fails. The GPU fixture used that same retired spelling; "weight" exercises the same g_idx path and validates on every compressed-tensors release. * Keep the wrapped many-to-many op an instance of its own class transformers' WeightConverter allows a many-to-many mapping only when operations holds an instance of its internal Ernie ops. The adapter that feeds such an op its original source contract is now a subclass of the op's class, so the replacement converter still constructs. * Adapt only the op that receives the decompressed buckets, and refuse pickled shards A later op in a converter chain (Concatenate after MergeModulelist) gets the previous op's output under that op's own contract, so only the first op is given the original source patterns. A .bin shard is read whole in the model dtype before any converter runs, which would cast the packed words to bf16 garbage; the plan builder refuses it with an instruction. * Shim only the inference-only remote MoE, keep load-only converters out of the save, and refuse ambiguous buckets The training-capable DeepSeek-V2/V3 implementation has its own training dispatch and a gate that returns an auxiliary loss, so only a forward that computes its output under 'if not self.training' and nothing else is shimmed. The decompression converters are dropped from model._weight_conversions after loading, so save_pretrained does not reverse them into packed names without recreating the packed metadata. A bucket with some experts packed and others not, or projections quantized under different config groups, is refused with an instruction instead of being decompressed with the wrong scheme. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Patch a remote gate only when its MoE block is the inference-only port A training-capable block's gate computes its auxiliary routing loss in train mode; running it in eval mode dropped that loss. * Arm packed re-quantization only when an explicit quantization_config selects bitsandbytes 4-bit * Gate packed re-quantization on the caller's quantizer in the Llama loader too * Match the inference-only remote MoE port by structure, so renamed ports like sarvam train * Format the remote MoE shim tests * Resolve the expert quantization scheme per concrete layer * Accept the compressed_tensors and sparseml spellings at the requantize gate * Leave remote MoE blocks alone when either self.training branch has its own dispatch * Arm packed re-quantization for an explicit bitsandbytes 4-bit quantization_config * Run every op of a chained expert converter under the original sources; hand a stack-only MergeModulelist the per-expert list * Resolve compressed-tensors groups by exact module path first, then regex, then class * Accept expert source patterns that start with the separator Mixtral and PhiMoE list their expert sources as .experts.*.w1.weight, so the per-layer scheme lookup never found the layer path and fell back to the default group. * Check exact and regex targets before compressed-tensors' generic match With the module in hand the generic is_match loop returned the first matching group, so a broad Linear group listed ahead of an exact-path override won. * Let explicit non-bitsandbytes quantizers clear the local 4-bit flags again Only an explicit bitsandbytes 4-bit config needs the caller's flags kept. A GPTQ, AWQ or FP8 config with the default load_in_4bit = True takes the checker's answer, as on main, so the 4-bit branch does not replace it with a BitsAndBytesConfig. * Keep a caller's BitsAndBytesConfig instead of the default NF4 one With fast_inference the loader forwards load_in_4bit = True alongside an explicit config, so the 4-bit branch replaced the caller's quant type, double quant and skip list with Unsloth's NF4 defaults. * Recognise self.training inside compound predicates of remote MoE forwards if self.training and self.ep_size == 1: dispatches its own experts in training, but the check only looked at a bare (or negated) self.training test, so such a forward was taken for an inference-only port and replaced. * Evaluate the whole training predicate before shimming a remote MoE forward Grouped or negated conditions such as (self.training and self.ep_size == 1) or use_fused still hid a training branch. The test is now evaluated with self.training True and every other term unknown: a branch counts as reachable in training unless the predicate is known to rule it out. * Chain a composed bitsandbytes quantizer's conversion hook on armed loads * Gate the re-quantization tests on update_weight_conversions (transformers 5.8+) * Say when transformers lacks the quantizer conversion hook instead of blaming another library * Checkpoint remote decoder layers whose loop never calls a checkpoint function Kimi-K2.7's DeepSeek-V3 port declares gradient checkpointing support and keeps a gradient_checkpointing flag, but its decoder loop calls every layer directly, so enabling it kept every activation. Wrap the layers of such remote backbones so the flag checkpoints through the function gradient_checkpointing_enable() installed. Native models, loops that already checkpoint, and calls carrying a cache are left alone. Also name 5.8 in the old-transformers message. * Pass keyword tensors that need a gradient positionally through the remote layer checkpoint A reentrant checkpoint holds keyword arguments by closure, so a tensor carried from layer to layer as a keyword (Kimi-K3's block_residual) stayed tied to the earlier layers' graph. torch's checkpoint then failed with a second backward through the graph, and Unsloth's offloaded checkpoint re-ran the earlier layers' backward inside every recompute. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re-quantize compressed-tensors MXFP4 checkpoints to bitsandbytes 4-bit on the fly, and shim the 5.x model API for 4.x remote code * Train Kimi-K3's latent MoE and attention residuals in 4-bit Run the inference-only remote MoE port's own forward in train mode, with the block and its gate reporting eval and moe_infer swapped for the differentiable dispatch on the instance, so ports that wrap the expert mix (Kimi-K3's latent routed_expert_down_proj / norm / up_proj) keep those projections. Keep the Linear(hidden, 1) attention-residual projections Kimi-K3 reads as score vectors out of bitsandbytes, and convert a remote list _tied_weights_keys only on the transformers 5 loader that needs the mapping. * Keep Kimi-K3's MXFP4 routed experts packed and dequantize them on the fly A compressed-tensors mxfp4-pack-quantized checkpoint whose routed experts are remote-code w1 / w2 / w3 Linears now loads each MoE layer's experts into one unsloth_zoo Mxfp4StackedExperts: the per-expert packed bytes and e8m0 scales are stacked by load-time converters, never decompressed or re-quantized to NF4. The remote MoE shim dispatches the block to the stacks in train and eval. Expert LoRA stays opt in (w1 / w3 -> gate_up_proj, w2 -> down_proj). UNSLOTH_MXFP4_KEEP_PACKED=0 restores the NF4 path. * Keep packed MXFP4 experts out of the default LoRA targets in FastLanguageModel too * Keep every MXFP4 Linear packed on the compressed-tensors route too, and merge it safely An MXFP4 checkpoint now stays MXFP4 on every load route, all or nothing: - One plan decides what stays packed: remote-code expert ModuleLists become unsloth_zoo Mxfp4StackedExperts, every other packed Linear becomes an Mxfp4PackedLinear holding compressed-tensors' own weight_packed / weight_scale. A layer packed only in part, or a packed module with no home, keeps the existing route for the whole checkpoint. - 16-bit loads, and 4-bit loads on transformers without the converter hook, go through compressed-tensors' own quantizer: its MXFP4 modules are adopted in place (fixing the remote _init_weights AttributeError on Kimi-K3) and the experts are stacked once the weights are in. Full finetuning and an explicit dequantize request keep the stock loader. - Mxfp4PackedLinear decodes in the model's dtype. The fused LoRA MLP / QKV / O kernels are not installed on packed bases, since they keep the 16-bit weight until backward; get_lora_parameters gets a transient decode. - PEFT merge on a packed base used to write the delta into a throwaway decode, silently keeping the base. Merging now densifies the module first, unmerging restores the packed bytes exactly, and merge_and_unload leaves a dense Linear. - UNSLOTH_MXFP4_KEEP_PACKED=0 restores the NF4 conversion. * Keep packed expert LoRA out when finetune_mlp_modules or finetune_language_layers is off The packed-expert target step read the caller's unscoped target list, so w1 / w2 / w3 in it re-added experts.gate_up_proj / down_proj after the scoped selection had dropped them. It now reads the same scoped selection as the rest of the MoE detection. * Keep the stock compressed-tensors route on an explicit run_compressed=False transformers 5.4 asks for load-time decompression with run_compressed=False, not dequantize=True. The adoption ran anyway, then the stock after-load step decompressed the adopted modules and removed their packed bytes, so the first forward failed. * Keep a packed MXFP4 Linear's decode dtype in step with module casts and the load dtype compute_dtype is a plain attribute, so .to(dtype) / .half() / .float() left it behind: the unmerged forward still ran (it decodes in the input dtype) but a PEFT merge wrote a weight in the old dtype and the merged forward failed with a dtype mismatch. A cast now updates it, and unmerge restores the dtype the dense weight has by then. transformers 5.x does not pass the load dtype to the quantizer, so the 16-bit route adopted every module as bf16 (and returned bf16-on-disk Linears in the default dtype) on a float16 load; it now falls back to the model's dtype, as the bitsandbytes route does. * Run the decompress-request test on the GPU when there is one An earlier FastLanguageModel load in the same process patches Llama's forward with kernels that need CUDA tensors. * Adopt compressed-tensors MXFP4 modules only when the packed forward can stand in for theirs - A scheme that also quantizes activations was adopted, and its quantize-dequantize forward replaced by one that only dequantizes the weight: the activation quantization silently went away. The bitsandbytes route already declined such schemes. - A compressed module in another layout (FP8 keeps its compressed weight under weight, not weight_packed) was not counted, so adoption removed the model-wide decompress hook that module still needed and the first forward failed on an FP8 weight. * Refuse PEFT initialisers that rewrite the base weight on a packed MXFP4 Linear PiSSA, OLoRA, CorDA, LoftQ and LoRA-GA subtract the initial adapter from the base weight.data. On a packed base that is a throwaway decode, so the adapter kept its value over an unchanged base and the model silently changed. They now raise. * Fix the input of the packed-merge test so its bf16 tolerance cannot flake The input came from the global RNG and the weights reached 2**23, so the merged vs unmerged bf16 comparison cancelled badly for some inputs (one element in 96, seen by the reviewers and here). The exact assertions are unchanged. * Keep MXFP4 weights packed only when unsloth_zoo can save them A full save of a kept-packed model writes loadable checkpoints only through unsloth_zoo's save patch (packed modules and Linears a merge made dense written dense under the checkpoint's names, the MXFP4 config dropped or the bitsandbytes skip list widened). With an unsloth_zoo that lacks it, a save wrote packed bytes or merged dense weights under a config that described neither. Both routes now keep the previous path in that case. * Stack packed MXFP4 experts only in MoE blocks the remote MoE shim dispatches A block with its own training branch, or with expert parallelism, keeps its original forward, which indexes or iterates the experts itself: with an expert LoRA ParamWrapper around the stack that failed in train and eval ('ParamWrapper' object is not iterable). The planner now uses the shim's own predicate; such experts stay packed one Linear each. * Keep nothing packed when a packed weight has no scale in the checkpoint The expert scan required both tensors, but a per-Linear packed module was planned from its weight_packed alone: its weight_scale then stayed uninitialised bytes and the load only reported it missing. The previous route raised on such a checkpoint. * Carry an accelerate hook's saved forward across a packed Linear's merge and unmerge A dispatched model keeps each module's original forward as _old_forward. The merge swaps the class to nn.Linear, but the hook still called the packed forward, which read the packed bytes the merge had set aside. * Plan only all-MXFP4 checkpoints on the compressed-tensors route, and match their keys once The hook read and matched every key of any compressed-tensors checkpoint before looking at its format (about a second on INT4 Kimi-K2.7), then matched them a second time for adoption. It now gates on the same format test as the bitsandbytes route and adopts the planned modules. * Stack gate_up before allocating the down stack on the compressed-tensors route All four stacks of a layer were allocated while every per-expert byte was still held, so the transient was the whole MoE layer. Filling gate_up first frees w1 and w3 before the down stack exists, leaving the gate_up stack (two thirds of the layer) as the transient. * Drop the test-only expert swap helper, use the single zoo stacking helper, rename the marker swap_in_packed_mxfp4_experts had no caller outside the tests; they now swap the stacks the load-time plan picks, as the loader does. The per-Linear packed class marker is _unsloth_mxfp4_packed_linear, matching unsloth_zoo. * Treat a self.training branch that only raises as no dispatch, so Kimi-K3's port is still shimmed #11537 now leaves a remote MoE block alone when either self.training branch has its own dispatch. Kimi-K3's port refuses training with else: raise NotImplementedError(...), which counted as a dispatch, so the block kept its inference-only forward and the gate's assert not self.training fired on the first training step. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Map only the output embedding of a 4.x tied-weights list onto the input embedding A 4.x _tied_weights_keys list only named keys that may be tied; the tie 4.x made from it was the output embedding onto the input embedding. Mapping every listed key onto the input embedding would tie a projection the remote code shares itself to the wrong tensor. * Honor a regex target string that names the per-expert Linears of packed MXFP4 experts A string target_modules reached the packed-expert opt-in as None, so .*experts.*w1 dropped the stacks' gate_up_proj it would have trained on the per-expert modules. The regex is now matched against each stack's original per-expert module names, as PEFT's fullmatch would. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Install the MXFP4 keep-packed hook for every compressed-tensors quant_method spelling The re-quantization branch accepts compressed-tensors, compressed_tensors and the legacy sparseml spelling, but a 16-bit load never reaches it, and the keep-packed hook was only installed for compressed-tensors. The other spellings fell back to the stock loader. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Import has_mxfp4_base in llama.py directly instead of re-exporting it from unsloth.kernels The import-hoist lint refuses a name added to a package __init__ that nothing there uses and __all__ does not list. * Gate the CUDA-only tests on the real accelerator probe, not the spoofable torch.cuda check * Gate the CUDA-only tests on the real accelerator probe, not the spoofable torch.cuda check * Skip the kept-packed MXFP4 tests when unsloth-zoo lacks the packed save * Skip the MXFP4 tests on a compressed-tensors without the MXFP4 format compressed-tensors 0.12.x, which vllm 0.11.2 pins, has no mxfp4-pack-quantized format and rejects the MXFP4 quantization config the fixtures build, so ten tests failed on transformers 4.57.6 with a pydantic ValidationError instead of skipping. * Give the new MXFP4 packed Linear module the AGPL-3.0 header New files in unsloth carry the AGPL-3.0 header, as the other modules added to unsloth/models this cycle do. * Honor list targets that name packed experts by dotted path PEFT matches a list entry against a module name exactly or as a dotted suffix, so experts.3.w1 or a per-expert Linear's full path named it before the experts were stacked. Only bare w1 / w2 / w3 were recognised, so an expert-only list failed with NoMatchingPeftModuleError and a mixed list silently dropped the expert adapters. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Match regex expert targets against every expert index The regex check only tried expert 0, so a target like .*experts\.3\.w1 that named a later expert before stacking no longer opted its stack in. * Keep packed experts in the default and all-linear LoRA targets The auto target regex is built after the per-expert Linears were stacked, so it no longer names w1 / w2 / w3 and the default or all-linear request dropped the experts, which it trained before stacking (and still trains with UNSLOTH_MXFP4_KEEP_PACKED=0). Opt both stacks in on the auto path while the MLP and language families are in scope. * Keep a layer-scoped expert target scoped to its layers PEFT suffix-matches target_parameters, so the generic experts.gate_up_proj wraps every MoE layer. When the request names experts in only some layers, emit each chosen stack's qualified path instead. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update tests for the prepared config now passed as planner_config The merge with main routes the re-quantized compressed-tensors config through main's planner_config argument of resolve_unsloth_device_map instead of a separate prepared_config argument, so the planner test now passes it that way and checks the older-planner fallback through a planner without a config parameter. It installs its own planner module in sys.modules, since other tests leave fakes there. The ModelOpt source check now matches the config branch condition that also covers a re-quantized compressed-tensors checkpoint. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments in the compressed-tensors int4 path * Use the AGPL-3.0 header on the new compressed-tensors and remote MoE modules, as other new model files do * Trim comments in the Kimi-K3 MXFP4 packed path * Keep the converter wrapper copyable: native MoE compressed-tensors loads recursed transformers' loader deep-copies each WeightConverter; copy builds the wrapper without __init__ and probes __setstate__, which __getattr__ forwarded to the unset self.op, recursing until RecursionError on every native-MoE INT4 checkpoint. * Keep compressed-tensors INT4/INT8 packed with exact Triton decode instead of re-quantizing to NF4 Int4PackedLinear loads the checkpoint's weight_packed / weight_scale / weight_zero_point / weight_g_idx as stored and exposes the packed words with an Int4QuantState, so Unsloth's LoRA kernels decode them exactly per use (Triton dequant + cuBLAS; a fused CUDA-core GEMV for decode). The NF4 fallback (UNSLOTH_COMPRESSED_TENSORS_INT4=nf4, or no CUDA / ROCm / unsupported scheme) now decompresses with the same exact Triton decode. * Packed INT4 route: keep routers and renamed modules on the per-module decompress converter; cheaper decode launches Non-Linear packed modules (MoE routers) no longer push the whole model back to NF4: only they keep the decompress converter, scoped by a lookbehind on their exact names. The GEMV and dequant kernels reuse their compiled Triton kernel per specialization and cache launch arguments on the quant state. * MXFP4 packed Linear: fused GEMM for skinny inputs, fast LoRA kernels on packed bases Inputs up to 64 rows (decode, small batches) run unsloth_zoo's fused MXFP4 GEMM and never build the 16-bit weight; wider ones decode once for cuBLAS. Inference skips the autograd node. get_lora_parameters now hands the fused LoRA kernels weight_packed plus a decoder quant state, so apply_lora_mlp / qkv / o run on packed bases without holding a 16-bit weight until backward; fast_dequantize, matmul_lora and fast_linear_forward understand that state. * Format the packed INT4 route with the repo's ruff wrapper; drop an unused alias and constexpr * MXFP4 packed Linear: fused GEMM only for the forward; dX stays on decode + cuBLAS * Packed INT Triton decode: INT schemes only, MXFP4 groups keep their own decode * Tighten comments in compressed-tensors and loader paths * Trim comments in the MXFP4 packed layer * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the compressed-tensors MXFP4 oracle tests when its decompress is unsupported * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Load the packed-base helpers alongside get_lora_parameters in the AST-extracted fp8 block size tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to the JIT launch on Triton < 3.7 in the INT4 packed kernels CompiledKernel[grid](*runtime_args) only matches the launcher signature from Triton 3.7 (torch 2.13); older launchers also take the constexprs and raise TypeError, so every packed INT4 call failed on torch <= 2.12. * Route packed INT matmuls to the fused GEMV only while it is faster The GEMV decodes the whole weight once per activation row, while decode + cuBLAS is flat in rows. On B200 it wins or ties only up to ~2**24 multiply-adds for 2 / 4-bit weights (1024 x 4096 at 4 rows) and is 2-4x slower past that at 2-4 rows; for 8-bit it never wins. Bound the GEMV by work and bit width instead of a fixed 4 rows, and test both paths. * Keep the fused INT GEMV for one row, decode + cuBLAS from two The work bound in the previous commit sent single-row decode on wide layers to decode + cuBLAS, which is 2x slower than the GEMV on an RTX PRO 6000 (-7.5% Qwen3-8B W4A16 decode tok/s). One row keeps the GEMV on every GPU; two to four rows go to decode + cuBLAS (2-4x faster per B200 layer, +50% batch-4 decode on RTX PRO 6000). * Trim comments in the Kimi-K3 packed MXFP4 code * Trim comments in the compressed-tensors INT4 loading code * Skip main's load-time compressed-tensors decompress when MXFP4 is kept packed * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore the Llama classes after in-process FastLanguageModel loads in the compressed-tensors tests * Restore the Llama classes after the in-process loads in the MXFP4 compressed Linear tests * Skip the explicit bnb 4-bit arming test without compressed-tensors, like its neighbours * Skip activation-ordered int4 cases when compressed-tensors ignores weight_g_idx (0.19+) * Gate the keep-packed load tests on the transformers 5 loader, like their neighbours * Size packed INT8 compressed-tensors loads at 8 bits in the device-map planner; read the {load_in_4bit: True} shorthand as bitsandbytes * Run decode-sized packed INT4 matmuls through vLLM's Marlin when available; skip autograd for grad-free packed linears * Save a packed compressed-tensors load with the checkpoint's own config so the directory reloads; merged layers go under ignore * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use torch's tinygemm for decode-sized packed INT4 matmuls when Marlin is unavailable * Keep the runtime config on saves with stacked experts; keep caller bnb flags only while the check kept 4-bit (fast_inference) * Return empty packed INT4 matmuls without launching; drop the unmerge stash once a merge is unloaded * Trim comments; AGPL header body on the compressed-tensors bnb test to match its SPDX line * Trim comments; AGPL header bodies on the new tests to match their SPDX lines * Repack packed INT4 weights once at load instead of keeping a second copy The tinygemm and Marlin fast paths used to build their own repacked copy of each weight next to the checkpoint words, which cost about 3.3 GB extra on Qwen3-8B. Eligible layers are now repacked in place at load into Marlin's layout when vLLM is installed, else tinygemm's (bf16). The dequantize kernel reads either layout directly, so prefill and training need no second copy. Each layer is checked to dequantize bit-identically before its words are replaced, and state_dict / save converts back to the checkpoint layout. Training always takes the exact dequantize + matmul path. Set UNSLOTH_INT4_REPACK=0 to keep the checkpoint layout. * Give the auto-target test's tiny Kimi layer an attention Linear, as real checkpoints have * Exact dequant for every training forward; size packed INT8 at 8 bits under an explicit bnb config; honour the bnb skip list when adopting packed layers; refuse full saves of mixed packed layouts * Take exact dequant after a dtype cast of a repacked layer; skip unpacking repacked base words for adapter saves * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
148 lines
5.9 KiB
Python
148 lines
5.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
"""Transformers 4.x remote code (Kimi-K3) against the transformers 5 model API, offline on CPU."""
|
|
|
|
import inspect
|
|
import linecache
|
|
import sys
|
|
import types
|
|
|
|
import torch
|
|
|
|
import unsloth # noqa: F401
|
|
from transformers import PreTrainedModel
|
|
|
|
_SRC = """
|
|
import torch
|
|
from torch import nn
|
|
from transformers import PretrainedConfig, PreTrainedModel
|
|
from transformers.utils.generic import OutputRecorder
|
|
|
|
class TinyRemoteConfig(PretrainedConfig):
|
|
model_type = "tiny_remote_k3"
|
|
def __init__(self, vocab_size=32, hidden_size=8, tie_word_embeddings=True, **kwargs):
|
|
self.vocab_size = vocab_size; self.hidden_size = hidden_size
|
|
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
|
|
|
|
class TinyRemoteModel(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
|
|
|
class TinyRemoteForCausalLM(PreTrainedModel):
|
|
config_class = TinyRemoteConfig
|
|
_tied_weights_keys = ["lm_head.weight"]
|
|
tie_calls = 0
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.model = TinyRemoteModel(config)
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|
self.post_init()
|
|
def get_input_embeddings(self):
|
|
return self.model.embed_tokens
|
|
def get_output_embeddings(self):
|
|
return self.lm_head
|
|
def tie_weights(self):
|
|
type(self).tie_calls += 1
|
|
if self.config.tie_word_embeddings:
|
|
self.lm_head.weight = self.model.embed_tokens.weight
|
|
def forward(self, input_ids):
|
|
return self.lm_head(self.model.embed_tokens(input_ids))
|
|
"""
|
|
|
|
|
|
def _remote_module(name):
|
|
filename = f"<{name}>"
|
|
linecache.cache[filename] = (len(_SRC), None, _SRC.splitlines(True), filename)
|
|
mod = types.ModuleType(name)
|
|
sys.modules[name] = mod
|
|
exec(compile(_SRC, filename, "exec"), mod.__dict__)
|
|
for cls in (mod.TinyRemoteConfig, mod.TinyRemoteModel, mod.TinyRemoteForCausalLM):
|
|
cls.__module__ = name
|
|
return mod
|
|
|
|
|
|
def _is_transformers5():
|
|
return "recompute_mapping" in inspect.signature(PreTrainedModel.tie_weights).parameters
|
|
|
|
|
|
def test_output_recorder_imports_from_the_4x_location():
|
|
from transformers.utils.generic import OutputRecorder # noqa: F401
|
|
mod = _remote_module("transformers_modules.tiny_remote_a.modeling_tiny")
|
|
assert mod.OutputRecorder is not None
|
|
|
|
|
|
def test_a_4x_remote_model_builds_and_ties_on_this_transformers():
|
|
mod = _remote_module("transformers_modules.tiny_remote_b.modeling_tiny")
|
|
model = mod.TinyRemoteForCausalLM(mod.TinyRemoteConfig())
|
|
assert model.lm_head.weight is model.model.embed_tokens.weight
|
|
# transformers 5 calls tie_weights with keywords the 4.x override does not take.
|
|
kwargs = {"recompute_mapping": False} if _is_transformers5() else {}
|
|
model.tie_weights(**kwargs)
|
|
assert model.lm_head.weight is model.model.embed_tokens.weight
|
|
assert mod.TinyRemoteForCausalLM.tie_calls >= 1
|
|
if _is_transformers5():
|
|
assert model._tied_weights_keys == {"lm_head.weight": "model.embed_tokens.weight"}
|
|
else:
|
|
# 4.x reads the list itself and must keep it.
|
|
assert model._tied_weights_keys == ["lm_head.weight"]
|
|
ids = torch.tensor([[1, 2, 3]])
|
|
assert model(ids).shape == (1, 3, 32)
|
|
|
|
|
|
def test_an_untied_remote_model_keeps_its_own_head():
|
|
mod = _remote_module("transformers_modules.tiny_remote_c.modeling_tiny")
|
|
model = mod.TinyRemoteForCausalLM(mod.TinyRemoteConfig(tie_word_embeddings = False))
|
|
assert model.lm_head.weight is not model.model.embed_tokens.weight
|
|
|
|
|
|
def test_native_classes_are_left_alone():
|
|
"""Only `transformers_modules` classes get the tie_weights wrapper or the mapping."""
|
|
from transformers import LlamaConfig, LlamaForCausalLM
|
|
|
|
cfg = LlamaConfig(
|
|
hidden_size = 16,
|
|
intermediate_size = 32,
|
|
num_hidden_layers = 1,
|
|
num_attention_heads = 2,
|
|
num_key_value_heads = 2,
|
|
vocab_size = 32,
|
|
)
|
|
before = LlamaForCausalLM.__dict__.get("tie_weights")
|
|
keys = LlamaForCausalLM._tied_weights_keys
|
|
LlamaForCausalLM(cfg)
|
|
assert LlamaForCausalLM.__dict__.get("tie_weights") is before
|
|
assert LlamaForCausalLM._tied_weights_keys == keys
|
|
|
|
|
|
def test_only_the_output_embedding_is_mapped_onto_the_input_embedding():
|
|
if not _is_transformers5():
|
|
return
|
|
mod = _remote_module("transformers_modules.tiny_remote_d.modeling_tiny")
|
|
|
|
class SharedProjection(mod.TinyRemoteForCausalLM):
|
|
_tied_weights_keys = ["lm_head.weight", "proj.weight"]
|
|
|
|
def __init__(self, config):
|
|
PreTrainedModel.__init__(self, config)
|
|
self.model = mod.TinyRemoteModel(config)
|
|
self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias = False)
|
|
self.proj = torch.nn.Linear(config.hidden_size, config.hidden_size, bias = False)
|
|
self.post_init()
|
|
|
|
SharedProjection.__module__ = mod.__name__
|
|
model = SharedProjection(mod.TinyRemoteConfig())
|
|
assert model._tied_weights_keys == {"lm_head.weight": "model.embed_tokens.weight"}
|
|
assert model.proj.weight is not model.model.embed_tokens.weight
|
|
assert model.lm_head.weight is model.model.embed_tokens.weight
|