How to compress t5-v1_1-xxl and Gemma-2-2B?

[mingyi456]

I managed to successfully compress Cosmos-Predict2 and Chroma, but when I tried to compress the T5 text encoder model used by Flux, I get the following error instead:

Traceback (most recent call last):
  File "F:\AI setups\Diffusers\models\compress t5.py", line 42, in <module>
    compress_model(
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\dfloat11\dfloat11.py", line 622, in compress_model
    save_file(model.state_dict(), os.path.join(save_path, 'model.safetensors'))
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\safetensors\torch.py", line 352, in save_file
    serialize_file(_flatten(tensors), filename, metadata=metadata)
                   ^^^^^^^^^^^^^^^^^
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\safetensors\torch.py", line 577, in _flatten
    raise RuntimeError(
RuntimeError:
            Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'shared.weight', 'encoder.embed_tokens.weight'}].
            A potential way to correctly save your model is to use `save_model`.
            More information at https://huggingface.co/docs/safetensors/torch_shared_tensors

Is this due to an error with my compression code, or is what I am trying to do not supported? The complete code I used, including the pattern_dict, is below:

import torch
from dfloat11 import compress_model
from transformers import T5EncoderModel
save_path = r".\t5-v1_1-xxl-DF11"
save_single_file = True
check_correctness = True
block_range = (0, 100)

text_encoder_2 = T5EncoderModel.from_pretrained(
	r"..\models\FLUX.1-dev",
	subfolder = "text_encoder_2",
	torch_dtype = torch.bfloat16,
	local_files_only = True
)

pattern_dict={
	"block\.\d+": (
		"layer.0.SelfAttention.q",
        "layer.0.SelfAttention.k",
        "layer.0.SelfAttention.v",
        "layer.0.SelfAttention.o",
        "layer.1.DenseReluDense.wi_0",
        "layer.1.DenseReluDense.wi_1",
        "layer.1.DenseReluDense.wo",
	)
}

# Compress the model using DFloat11 compression
compress_model(
	model=text_encoder_2,
	pattern_dict= pattern_dict,
	save_path = save_path,
	save_single_file = save_single_file,
	check_correctness = check_correctness,
	block_range = block_range,
)

Edit: Found an issue with the pattern_dict, block should be replaced with encoder.block, but the shared tensors issue will still stop the file from being saved after the compression process finishes.

[mingyi456]

I tried the same code with Gemma-2-2B (since it is also used as a text encoder by some diffusion models), and I get the same error. Even if I purposely mess up my pattern_dict, the error is still the same, meaning the issue lies with the shared tensors (which is what the error message says),

But if I edit the code to replace save_file() with save_model(), removing the .state_dict(), the code instantly exits, while dumping an a copy of the Gemma-2-2B weights of the exact same file size as the original model. That should probably mean my pattern_dict is not matching the correct layers and tensors properly, so the code finds nothing to compress.

[mingyi456]

I found an issue with the pattern_dict, which was why the script was refusing to compress anything. The fixed code for compressing Gemma-2 is below:

import torch
from dfloat11 import compress_model
from transformers import Gemma2ForCausalLM

save_path = r".\gemma-2-2b-DF11"
save_single_file = True
check_correctness = True
block_range = (0, 100)
text_encoder = Gemma2ForCausalLM.from_pretrained(
	r".\gemma-2-2b",
	torch_dtype = torch.bfloat16,
)
pattern_dict={
	"model.layers\.\d+": (
		"self_attn.q_proj",
        "self_attn.k_proj",
        "self_attn.v_proj",
        "self_attn.o_proj",
        "mlp.gate_proj",
        "mlp.up_proj",
        "mlp.down_proj",
	)
}
# Compress the model using DFloat11 compression
compress_model(
	model=text_encoder,
	pattern_dict= pattern_dict,
	save_path = save_path,
	save_single_file = save_single_file,
	check_correctness = check_correctness,
	block_range = block_range,
)

Now the compression starts and runs smoothly, as expected. I also made the following modification to the save_file() line, otherwise the same shared tensor error will occur, making the script refuse to save the compressed weights:

            # save_file(model.state_dict(), os.path.join(save_path, 'model.safetensors'))
            from safetensors.torch import save_model as save_model_safetensors
            save_model_safetensors(model, os.path.join(save_path, 'model.safetensors'))

But I end up with the following error instead:

Traceback (most recent call last):
  File "F:\AI setups\Diffusers\models\compress gemma-2-2b.py", line 32, in <module>
    compress_model(
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\dfloat11\dfloat11.py", line 624, in compress_model
    save_model_safetensors(model, os.path.join(save_path, 'model.safetensors'))
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\safetensors\torch.py", line 169, in save_model
    to_removes = _remove_duplicate_names(state_dict)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "F:\AI setups\Diffusers\diffusers-venv\Lib\site-packages\safetensors\torch.py", line 113, in _remove_duplicate_names
    raise RuntimeError(
RuntimeError: Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {'model.layers.0.split_positions'}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue.

@tonyzhang617 Can you help me with this?

[mingyi456]

I changed my pattern_dict for compressing Gemma-2-2B to include "model" : "embed_tokens", and the code compresses successfully, but the compressed file is 4.4GB in size, only about 0.4GB smaller than the uncompressed version, even though I checked that the compression ratio of every layer is less than 70%. This is with the default save_file() function in the compression function.

Also, when trying to load the compressed Gemma model as the text encoder for Sana-Sprint, I get a key error saying lm_head.weight is missing or something.

[mingyi456]

After changing the pattern_dict some more, and adding lm_head, I can get a 4.04GB output file, but that still fails to load correctly.
This is the pattern_dict I used:

pattern_dict={
	"model.embed_tokens" : [],
	"model.layers\.\d+": (
		"self_attn.q_proj",
		"self_attn.k_proj",
		"self_attn.v_proj",
		"self_attn.o_proj",
		"mlp.gate_proj",
		"mlp.up_proj",
		"mlp.down_proj",
		# "input_layernorm",
		# "post_attention_layernorm",
		# "pre_feedforward_layernorm",
		# "post_feedforward_layernorm"		
	),
	# "model" : ("norm",),
	"lm_head" : []
}

And this is the error I get:

Shape mismatch for model.layers.0.input_layernorm.weight: model torch.Size([3072]) vs loaded torch.Size([2304])
Shape mismatch for model.layers.0.post_attention_layernorm.weight: model torch.Size([3072]) vs loaded torch.Size([2304])
Cannot find module path for model.layers.0.post_feedforward_layernorm.weight
Cannot find module path for model.layers.0.pre_feedforward_layernorm.weight

(repeated for all layers)

[mingyi456]

For some reason, my Gemma-2-2B DF11 compressed model now works, but I need to specify the bfloat16_model parameter, or else the Cannot find module path error pops up. This is strange, because I thought that was only for diffusers models, and models that live in the transformers library do not require that.
The final model filesize is 3.30GB, while the uncompressed model is 4.86GB, which is within expectation. However, the total VRAM usage of the entire Sana-Sprint pipeline, with a DF11 compressed text encoder, is only 0.5GB less, which is weird. At least inference time is only 0.2-0.5s slower with pipe.to("cuda"), and actually ~0.5s faster with pipe.enable_model_cpu_offload().

[mingyi456]

Now, I also got compressed T5-XXL working, also only if I specify bfloat16_model, but for some reason randomly initializing the original bf16 weights for both Gemma 2 and T5-XXL leads to a different result with the final generated image. The code I used to load the 2 text encoders is:

from transformers import AutoModel, Gemma2ForCausalLM, Gemma2PreTrainedModel, Gemma2Model, AutoConfig
from transformers.modeling_utils import no_init_weights

with no_init_weights():
	text_encoder = Gemma2Model._from_config(
		AutoConfig.from_pretrained(
			gemma2_dir
		),
		torch_dtype=torch.bfloat16,
		attn_implementation = "flash_attention_2"
	).to(torch.bfloat16)
text_encoder_df11 = DFloat11Model.from_pretrained(gemma2_df11_dir, device = "cpu", bfloat16_model = text_encoder, attn_implementation = "flash_attention_2")

But if I initialize the text encoders normally, and then replace them with the DF11 weights, I get the exact same result. So maybe my compressed weights are missing some tensors?