Adding support for unsloth 4bit compressed models

[charlesmartin14]

It would be great to add support for UNsloth 4bit compressed files such as the model.safetensors file contained in

https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B-unsloth-bnb-4bit

[charlesmartin14]

It seems we can not decompress the weight matrices directly from the unsloth model.safetensors files directly because the unsloth safetensors metadata does not contain the shape information for the weight matrices

all it contains is

model.layers.8.mlp.up_proj.weight
  model.layers.8.mlp.up_proj.weight.absmax
  model.layers.8.mlp.up_proj.weight.nested_absmax
  model.layers.8.mlp.up_proj.weight.nested_quant_map
  model.layers.8.mlp.up_proj.weight.quant_map
  model.layers.8.mlp.up_proj.weight.quant_state.bitsandbytes__nf4

This means it is necessary to obtain this information using the unsloth library code to either:

1. read the entire safetensors model file into memory directly using
   FastLanguageModel.from_pretrained(

2. read the config file, and then figure out the weight dimensions for the specific model type supported by unsloth:
   LLama Qwen Gemma etc

Neither of these solutions belongs in weightwatcher because it forces weightwatcher to be very strongly dependent on unsloth

[charlesmartin14]

The current solution is to decompress the Unsloth 4bit file directly outside of weightwatcher

Note: this will convert a single model.safetensors --> 4 separate model.###-###.safetensors files:

model.safetensors -->

model-00001-of-00004.safetensors
model-00002-of-00004.safetensors
model-00003-of-00004.safetensors
model-00004-of-00004.safetensors

and

unsloth requires access to a GPU

Code

def load_and_unquantize(
    model_name="DeepSeek-R1-Distill-Llama-8B-unsloth-bnb-4bit",
    load_in_4bit=True,
    save_path="my_uncompressed_model",
    save_method="merged_16bit",
    max_seq_length=2048,
    dtype=None,
):
    """
    Loads a 4-bit Unsloth model and saves it in merged FP16 format.
      - model_name: name or path for the 4-bit model
      - load_in_4bit: bool, whether to load in 4-bit
      - save_path: where to store the uncompressed model
      - save_method: e.g., 'merged_16bit'
      - max_seq_length: rope scaling (2048 default)
      - dtype: torch dtype, often None is fine
    """
    from unsloth import FastLanguageModel
    import os

    # 1) Load the 4-bit model + tokenizer
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_name,
        max_seq_length=max_seq_length,
        dtype=dtype,
        load_in_4bit=load_in_4bit,
    )

    # 2) Merge and save as FP16
    model.save_pretrained_merged(save_path, tokenizer, save_method)
    print(f"Uncompressed model saved to: {os.path.abspath(save_path)}")

# Call the function directly (adjust arguments as needed)
load_and_unquantize(
    model_name="DeepSeek-R1-Distill-Llama-8B-unsloth-bnb-4bit",
    load_in_4bit=True,
    save_path="my_uncompressed_model",
    save_method="merged_16bit",
    max_seq_length=2048,
    dtype=None
)