flux-lora-resizing / upsample_lora_rank.py

Upload upsample_lora_rank.py

d69d12f verified over 1 year ago

4.88 kB

	"""
	Usage:

	python upsample_lora_rank.py --repo_id="cocktailpeanut/optimus" \
	--filename="optimus.safetensors" \
	--new_lora_path="optimus_16.safetensors" \
	--new_rank=16
	"""

	import torch
	from huggingface_hub import hf_hub_download
	import safetensors.torch
	import fire


	def orthogonal_extension(matrix, target_rows):
	"""
	Extends the given matrix to have target_rows rows by adding orthogonal rows.

	Args:
	matrix (torch.Tensor): Original matrix of shape [original_rows, columns].
	target_rows (int): Desired number of rows.

	Returns:
	extended_matrix (torch.Tensor): Matrix of shape [target_rows, columns].
	"""
	original_rows, cols = matrix.shape
	assert target_rows >= original_rows, "Target rows must be greater than or equal to original rows."

	# Perform QR decomposition
	Q, R = torch.linalg.qr(matrix.T, mode="reduced") # Transpose to get [columns, original_rows]
	Q = Q.T # Back to [original_rows, columns]

	# Generate orthogonal vectors
	if target_rows > original_rows:
	additional_rows = target_rows - original_rows
	random_matrix = torch.randn(additional_rows, cols, dtype=matrix.dtype, device=matrix.device)
	# Orthogonalize against existing Q
	for i in range(additional_rows):
	v = random_matrix[i]
	v = v - Q.T @ (Q @ v)
	v = v / v.norm()
	Q = torch.cat([Q, v.unsqueeze(0)], dim=0)
	extended_matrix = Q
	return extended_matrix


	def increase_lora_rank_orthogonal(state_dict, target_rank=16):
	"""
	Increases the rank of all LoRA matrices in the given state dict using orthogonal extension.

	Args:
	state_dict (dict): The state dict containing LoRA matrices.
	target_rank (int): Desired higher rank.

	Returns:
	new_state_dict (dict): State dict with increased-rank LoRA matrices.
	"""
	new_state_dict = state_dict.copy()
	for key in state_dict.keys():
	if "lora_A.weight" in key:
	lora_A_key = key
	lora_B_key = key.replace("lora_A.weight", "lora_B.weight")
	if lora_B_key in state_dict:
	lora_A = state_dict[lora_A_key]
	dtype = lora_A.dtype
	lora_A = lora_A.to("cuda", torch.float32)
	lora_B = state_dict[lora_B_key]
	lora_B = lora_B.to("cuda", torch.float32)

	original_rank = lora_A.shape[0]

	# Extend lora_A and lora_B
	lora_A_new = orthogonal_extension(lora_A, target_rank).to(dtype)
	lora_B_new = orthogonal_extension(lora_B.T, target_rank).T.to(dtype) # Transpose to match dimensions

	# Update the state dict
	new_state_dict[lora_A_key] = lora_A_new
	new_state_dict[lora_B_key] = lora_B_new

	print(
	f"Increased rank of {lora_A_key} and {lora_B_key} from {original_rank} to {target_rank} using orthogonal extension"
	)

	return new_state_dict


	def compare_approximation_error(orig_state_dict, new_state_dict):
	for key in orig_state_dict:
	if "lora_A.weight" in key:
	lora_A_key = key
	lora_B_key = key.replace("lora_A.weight", "lora_B.weight")
	lora_A_old = orig_state_dict[lora_A_key]
	lora_B_old = orig_state_dict[lora_B_key]
	lora_A_new = new_state_dict[lora_A_key]
	lora_B_new = new_state_dict[lora_B_key]

	# Original delta_W
	delta_W_old = (lora_B_old @ lora_A_old).to("cuda")

	# Approximated delta_W
	delta_W_new = lora_B_new @ lora_A_new

	# Compute the approximation error
	error = torch.norm(delta_W_old - delta_W_new, p="fro") / torch.norm(delta_W_old, p="fro")
	print(f"Relative error for {lora_A_key}: {error.item():.6f}")


	def main(
	repo_id: str,
	filename: str,
	new_rank: int,
	check_error: bool = False,
	new_lora_path: str = None,
	):
	# ckpt_path = hf_hub_download(repo_id="TheLastBen/The_Hound", filename="sandor_clegane_single_layer.safetensors")
	if new_lora_path is None:
	raise ValueError("Please provide a path to serialize the converted state dict.")

	ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
	original_state_dict = safetensors.torch.load_file(ckpt_path)
	new_state_dict = increase_lora_rank_orthogonal(original_state_dict, target_rank=new_rank)

	if check_error:
	compare_approximation_error(original_state_dict, new_state_dict)

	new_state_dict = {k: v.to("cpu").contiguous() for k, v in new_state_dict.items()}
	# safetensors.torch.save_file(new_state_dict, "sandor_clegane_single_layer_32.safetensors")
	safetensors.torch.save_file(new_state_dict, new_lora_path)


	if __name__ == "__main__":
	fire.Fire(main)