Marlin24 Compressor

[Satrat]

Implements Marlin24 compression in SparseML. Note that this compressor does not have a decompress method, it is meant only for feeding directly into vllm.

Example Usage

Since this is a one-way conversion for vLLM, we won't automatically convert to this format when save_compressed is set in SparseML. Instead the save will have to be done explicitly like this:

from sparseml.transformers import SparseAutoModelForCausalLM, SparseAutoTokenizer
import torch

model_dir = "/network/sadkins/llama7b_sparse_24_w4a16_group128"
output_dir = "llama7b_marlin24"
model = SparseAutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16, device_map="cuda")
tokenizer = SparseAutoTokenizer.from_pretrained(model_dir)
model.save_pretrained(output_dir, quantization_format="marlin-24")
tokenizer.save_pretrained(output_dir)

Testing

Tested with a 7b Llama 2 model in vLLM, confirmed outputs for the marlin-24 format match the ouputs for the equivalent pack-quantized model. Also added some unit tests to compressed-tensors

Prompt: 'Hello, my name is', Generated text: '… Business: And I'm here to tell you about my new book!'
Prompt: 'The capital of France is', Generated text: ' Paris, but did you know that Parisians call themselves Parisins? True enough'
Prompt: 'The president of the United States is', Generated text: ', by his or her position, obliged to represent the United States internation,'
Prompt: 'The Boston Bruins are', Generated text: ' a hockey team based in Boston, Massachusetts that compete in the National Hockey League'

[dbogunowicz]

Suggested change

	prefix = name[: -(len(weight_suffix))]
	prefix = name.replace(weight_suffix)

More readable :)

[Satrat]

Yes but the replace code runs the risk of replacing ".weight" if it exists elsewhere in the string, not likely but possible

[dbogunowicz]

@bfineran can we have a license attached to this code if it's largely copied over from an open-source project (vllm), even though "our" Alex is the author?

[Satrat]

Good call, I think this got added in automatically when I ran make style, is there a way to turn that off for a certain file??

[bfineran]

you can skip it with `neuralmagic: no copyright. it's an apache license though so I think usage should be pretty permissable

[horheynm]

I like the design and how simple the implementation is. I didnt check the math and bit operations but other than that green from me

[bfineran]

you can skip it with `neuralmagic: no copyright. it's an apache license though so I think usage should be pretty permissable