Skip to content

fix a position bug in decoder rope/fuse more weight" #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tomorin523 wants to merge 2 commits into
base: main
Choose a base branch
from
+240 −310
Open

fix a position bug in decoder rope/fuse more weight" #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6fa4595
fix a position bug in decoder rope/fuse more weight
Tomorin523 Apr 27, 2026
5904b85
fix: refine pi05 inference
Tomorin523 Apr 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
172 changes: 97 additions & 75 deletions convert_from_jax_pi05.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,11 @@
import pickle
import orbax.checkpoint as ocp
import math
import torch.nn.functional as F
from transformers import AutoTokenizer

def convert_weights_pi05(weights, dump_weights):
def convert_weights_pi05(weights, dump_weights, time_embeds):
num_steps = time_embeds.shape[0]
# vision encoder weights
weights['vision_patch_embedding_w'].copy_(torch.tensor(dump_weights['PaliGemma']['img']['embedding']['kernel']['value'], dtype=torch.bfloat16, device="cpu"))
weights['vision_patch_embedding_b'].copy_(torch.tensor(dump_weights['PaliGemma']['img']['embedding']['bias']['value'], dtype=torch.bfloat16, device="cpu"))
Expand Down Expand Up @@ -94,11 +96,10 @@ def convert_weights_pi05(weights, dump_weights):
w_v = dump_weights['PaliGemma']['llm']['layers']['attn']['kv_einsum_1']['w']['value'][:, 1, 0].astype('float32')
w_q = w_q.reshape((18, 1024, 8, 2, 128)).transpose((0, 1, 2, 4, 3)).reshape((18, 1024, 2048))
w_k = w_k.reshape((18, 1024, 2, 128)).transpose((0, 1, 3, 2)).reshape((18, 1024, 256))

weights['decoder_attn_qkv_w'].copy_(torch.tensor(
decoder_attn_qkv_w = torch.tensor(
np.concatenate([w_q, w_k, w_v], axis=2),
dtype=torch.bfloat16, device="cpu"
))
dtype=torch.float32, device="cpu")


w_attn_o = dump_weights['PaliGemma']['llm']['layers']['attn']['attn_vec_einsum_1']['w']['value'].reshape((18, 8 * 256, 1024)).astype('float32')
weights['decoder_attn_o_w'].copy_(torch.tensor(
Expand All @@ -109,41 +110,17 @@ def convert_weights_pi05(weights, dump_weights):
w_gate = dump_weights['PaliGemma']['llm']['layers']['mlp_1']['gating_einsum']['value'][:, 0].astype('float32')
w_up = dump_weights['PaliGemma']['llm']['layers']['mlp_1']['gating_einsum']['value'][:, 1].astype('float32')
w_down = dump_weights['PaliGemma']['llm']['layers']['mlp_1']['linear']['value'].astype('float32')
attn_adarms_kernel = torch.tensor(dump_weights['PaliGemma']['llm']['layers']['pre_attention_norm_1']['Dense_0']['kernel']['value'], dtype=torch.bfloat16, device="cpu")
attn_adarms_bias = torch.tensor(dump_weights['PaliGemma']['llm']['layers']['pre_attention_norm_1']['Dense_0']['bias']['value'], dtype=torch.bfloat16, device="cpu")

weights['decoder_ffn_gate_w'].copy_(torch.tensor(
w_gate,
dtype=torch.bfloat16, device="cpu"
))
weights['decoder_ffn_up_w'].copy_(torch.tensor(
w_up,
dtype=torch.bfloat16, device="cpu"
))
weights['decoder_ffn_down_w'].copy_(torch.tensor(
w_down,
dtype=torch.bfloat16, device="cpu"
))


attn_adarms_kernel = dump_weights['PaliGemma']['llm']['layers']['pre_attention_norm_1']['Dense_0']['kernel']['value']
attn_adarms_bias = dump_weights['PaliGemma']['llm']['layers']['pre_attention_norm_1']['Dense_0']['bias']['value']
weights['decoder_pre_attn_norm_mod_w'].copy_(torch.tensor(attn_adarms_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_pre_attn_norm_mod_b'].copy_(torch.tensor(attn_adarms_bias, dtype=torch.bfloat16, device="cpu"))

ffn_adarms_kernel = dump_weights['PaliGemma']['llm']['layers']['pre_ffw_norm_1']['Dense_0']['kernel']['value']
ffn_adarms_bias = dump_weights['PaliGemma']['llm']['layers']['pre_ffw_norm_1']['Dense_0']['bias']['value']

weights['decoder_pre_ffn_norm_mod_w'].copy_(torch.tensor(ffn_adarms_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_pre_ffn_norm_mod_b'].copy_(torch.tensor(ffn_adarms_bias, dtype=torch.bfloat16, device="cpu"))
ffn_adarms_kernel = torch.tensor(dump_weights['PaliGemma']['llm']['layers']['pre_ffw_norm_1']['Dense_0']['kernel']['value'], dtype=torch.bfloat16, device="cpu")
ffn_adarms_bias = torch.tensor(dump_weights['PaliGemma']['llm']['layers']['pre_ffw_norm_1']['Dense_0']['bias']['value'], dtype=torch.bfloat16, device="cpu")

time_mlp_in_kernel = dump_weights['time_mlp_in']['kernel']['value']
time_mlp_in_bias = dump_weights['time_mlp_in']['bias']['value']
time_mlp_out_kernel = dump_weights['time_mlp_out']['kernel']['value']
time_mlp_out_bias = dump_weights['time_mlp_out']['bias']['value']

weights['decoder_time_mlp_in_w'].copy_(torch.tensor(time_mlp_in_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_time_mlp_in_b'].copy_(torch.tensor(time_mlp_in_bias, dtype=torch.bfloat16, device="cpu"))
weights['decoder_time_mlp_out_w'].copy_(torch.tensor(time_mlp_out_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_time_mlp_out_b'].copy_(torch.tensor(time_mlp_out_bias, dtype=torch.bfloat16, device="cpu"))
time_mlp_in_kernel = torch.tensor(dump_weights['time_mlp_in']['kernel']['value'], dtype=torch.bfloat16, device="cpu")
time_mlp_in_bias = torch.tensor(dump_weights['time_mlp_in']['bias']['value'], dtype=torch.bfloat16, device="cpu")
time_mlp_out_kernel = torch.tensor(dump_weights['time_mlp_out']['kernel']['value'], dtype=torch.bfloat16, device="cpu")
time_mlp_out_bias = torch.tensor(dump_weights['time_mlp_out']['bias']['value'], dtype=torch.bfloat16, device="cpu")

action_in_proj_w = torch.tensor(
dump_weights['action_in_proj']['kernel']['value'],
Expand All @@ -157,19 +134,65 @@ def convert_weights_pi05(weights, dump_weights):
weights['decoder_action_in_proj_w'].copy_(action_in_proj_w)
weights['decoder_action_in_proj_b'].copy_(action_in_proj_b)

action_out_kernel = dump_weights['action_out_proj']['kernel']['value']
action_out_bias = dump_weights['action_out_proj']['bias']['value']
weights['decoder_action_out_proj_w'].copy_(torch.tensor(action_out_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_action_out_proj_b'].copy_(torch.tensor(action_out_bias, dtype=torch.bfloat16, device="cpu"))

action_out_kernel = torch.tensor(
dump_weights['action_out_proj']['kernel']['value'],
dtype=torch.float32, device="cpu"
)
action_out_bias = torch.tensor(
dump_weights['action_out_proj']['bias']['value'],
dtype=torch.float32, device="cpu"
)
action_out_kernel *= -1.0 / num_steps
action_out_bias *= -1.0 / num_steps
final_norm_1 = dump_weights['PaliGemma']['llm']['final_norm_1']
final_adarms_kernel = torch.tensor(final_norm_1['Dense_0']['kernel']['value'], dtype=torch.bfloat16, device="cpu")
final_adarms_bias = torch.tensor(final_norm_1['Dense_0']['bias']['value'], dtype=torch.bfloat16, device="cpu")

final_adarms_kernel = final_norm_1['Dense_0']['kernel']['value']
final_adarms_bias = final_norm_1['Dense_0']['bias']['value']
weights['decoder_final_norm_mod_w'].copy_(torch.tensor(final_adarms_kernel, dtype=torch.bfloat16, device="cpu"))
weights['decoder_final_norm_mod_b'].copy_(torch.tensor(final_adarms_bias, dtype=torch.bfloat16, device="cpu"))

decoder_time_emb = []
for step in range(num_steps):
decoder_time_emb = F.silu( F.silu((time_embeds[step].view(1,-1) @ time_mlp_in_kernel + time_mlp_in_bias)) @ time_mlp_out_kernel + time_mlp_out_bias )
for i in range(18):
decoder_style_attn = decoder_time_emb @ attn_adarms_kernel[i] + attn_adarms_bias[i]
scale, shift, weights['decoder_gate_attn'][step, i] = torch.chunk(decoder_style_attn, 3, dim=-1)
scale = scale.to(dtype=torch.float32)
shift = shift.to(dtype=torch.float32)
fused_scale = (1 + scale).view(-1, 1)
attn_w_fused = fused_scale * decoder_attn_qkv_w[i]
attn_b_fused = shift @ decoder_attn_qkv_w[i]
weights['decoder_attn_fused_qkv_w'][step, i].copy_(attn_w_fused.squeeze(0).to(dtype=torch.bfloat16))
weights['decoder_attn_fused_qkv_b'][step, i].copy_(attn_b_fused.squeeze(0).to(dtype=torch.bfloat16))

decoder_style_ffn = decoder_time_emb @ ffn_adarms_kernel[i] + ffn_adarms_bias[i]
scale, shift, weights['decoder_gate_ffn'][step, i] = torch.chunk(decoder_style_ffn, 3, dim=-1)
scale = scale.to(dtype=torch.float32)
shift = shift.to(dtype=torch.float32)
gate_w = torch.tensor(w_gate[i], dtype=torch.float32, device="cpu")
up_w = torch.tensor(w_up[i], dtype=torch.float32, device="cpu")
fused_scale = (1 + scale).view(-1, 1)
gate_w_fused = fused_scale * gate_w
up_w_fused = fused_scale * up_w
gate_b_fused = shift @ gate_w
up_b_fused = shift @ up_w
weights['decoder_ffn_fused_gate_w'][step, i].copy_(gate_w_fused.to(dtype=torch.bfloat16))
weights['decoder_ffn_fused_up_w'][step, i].copy_(up_w_fused.to(dtype=torch.bfloat16))
weights['decoder_ffn_fused_gate_b'][step, i].copy_(gate_b_fused.squeeze(0).to(dtype=torch.bfloat16))
weights['decoder_ffn_fused_up_b'][step, i].copy_(up_b_fused.squeeze(0).to(dtype=torch.bfloat16))

decoder_style_final = decoder_time_emb @ final_adarms_kernel + final_adarms_bias
scale, shift, _ = torch.chunk(decoder_style_final, 3, dim=-1)
scale = scale.to(dtype=torch.float32)
shift = shift.to(dtype=torch.float32)
fused_scale = (1 + scale).view(-1, 1)
action_w_fused = fused_scale * action_out_kernel
action_b_fused = shift @ action_out_kernel + action_out_bias
weights['decoder_action_fused_out_proj_w'][step].copy_(action_w_fused.squeeze(0).to(dtype=torch.bfloat16))
weights['decoder_action_fused_out_proj_b'][step].copy_(action_b_fused.squeeze(0).to(dtype=torch.bfloat16))

weights['decoder_ffn_down_w'].copy_(torch.tensor(
w_down,
dtype=torch.bfloat16, device="cpu"
))
def load_jax_weights(jax_path: str):
params_path = os.path.join(jax_path, "params")
print(f"Loading jax weights from {params_path}")
Expand Down Expand Up @@ -213,11 +236,11 @@ def _ensure_value_dict(node):

def prepare_adarms_cond(num_steps):
dt = -1.0 / num_steps
time = torch.tensor(1.0, dtype=torch.float32, device="cuda")
time = torch.tensor(1.0, dtype=torch.float32, device="cpu")
min_period = 4e-3
max_period = 4.0
embedding_dim = 1024
fraction = torch.linspace(0.0, 1.0, embedding_dim // 2, device="cuda")
fraction = torch.linspace(0.0, 1.0, embedding_dim // 2, device="cpu")
period = min_period * (max_period / min_period) ** fraction
time_emb = []
for step in range(num_steps):
Expand All @@ -231,21 +254,20 @@ def prepare_prompt(prompt: str, embedding_weight, tokenizer_path):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
embedding_weight_torch = torch.tensor(
embedding_weight,
dtype=torch.bfloat16, device="cuda"
dtype=torch.bfloat16, device="cpu"
)
num_embeddings, embedding_dim = embedding_weight_torch.shape
language_embedding = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim).bfloat16().cuda()
language_embedding = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim).bfloat16()
with torch.no_grad():
language_embedding.weight.copy_(embedding_weight_torch)
prompt = [prompt.strip().replace("_", " ") + "\n"]
language_tokens = tokenizer(
prompt,
max_length=48,
return_tensors="pt",
)["input_ids"].to(device="cuda").squeeze(0)
)["input_ids"].squeeze(0)
language_embeds = language_embedding(language_tokens)
language_embeds *= language_embeds.shape[-1] ** 0.5
language_embeds = language_embeds.to(device="cpu")
return language_embeds, language_embeds.shape[0]

if __name__ == "__main__":
Expand All @@ -260,10 +282,11 @@ def prepare_prompt(prompt: str, embedding_weight, tokenizer_path):
embedding_weight = dump_weights['PaliGemma']['llm']['embedder']['input_embedding']['value']
embedding_weight_torch = torch.tensor(
embedding_weight,
dtype=torch.bfloat16, device="cuda"
dtype=torch.bfloat16, device="cpu"
)
language_embeds, prompt_len = prepare_prompt(args.prompt, embedding_weight, args.tokenizer_path)
time_embeds = prepare_adarms_cond(num_steps = 10)
num_steps = 10
time_embeds = prepare_adarms_cond(num_steps = num_steps)
weights = {

"embedding_weight": torch.zeros(257152,2048, dtype=torch.bfloat16, device="cpu"),
Expand Down Expand Up @@ -295,34 +318,33 @@ def prepare_prompt(prompt: str, embedding_weight, tokenizer_path):
"encoder_ffn_up_w": torch.zeros(18, 2048, 16384, dtype=torch.bfloat16, device="cpu"),
"encoder_ffn_down_w": torch.zeros(18, 16384, 2048, dtype=torch.bfloat16, device="cpu"),


"decoder_time_embeds": torch.zeros(10, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_time_mlp_in_w": torch.zeros(1024, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_time_mlp_in_b": torch.zeros(1024, dtype=torch.bfloat16, device="cpu"),
"decoder_time_mlp_out_w": torch.zeros(1024, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_time_mlp_out_b": torch.zeros(1024, dtype=torch.bfloat16, device="cpu"),
"decoder_pre_attn_norm_mod_w": torch.zeros(18, 1024, 3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_pre_attn_norm_mod_b": torch.zeros(18, 3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_pre_ffn_norm_mod_w": torch.zeros(18, 1024, 3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_pre_ffn_norm_mod_b": torch.zeros(18, 3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_final_norm_mod_w": torch.zeros(1024, 3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_final_norm_mod_b": torch.zeros(3 * 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_attn_qkv_w": torch.zeros(18, 1024, 2560, dtype=torch.bfloat16, device="cpu"),
"decoder_attn_fused_qkv_w": torch.zeros(num_steps, 18, 1024, 2560, dtype=torch.bfloat16, device="cpu"),
"decoder_attn_fused_qkv_b": torch.zeros(num_steps, 18, 2560, dtype=torch.bfloat16, device="cpu"),

"decoder_attn_o_w": torch.zeros(18, 2048, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_gate_w": torch.zeros(18, 1024, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_up_w": torch.zeros(18, 1024, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_fused_gate_w": torch.zeros(num_steps, 18, 1024, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_fused_gate_b": torch.zeros(num_steps, 18, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_fused_up_w": torch.zeros(num_steps, 18, 1024, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_fused_up_b": torch.zeros(num_steps, 18, 4096, dtype=torch.bfloat16, device="cpu"),
"decoder_ffn_down_w": torch.zeros(18, 4096, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_action_in_proj_w": torch.zeros(32, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_action_in_proj_b": torch.zeros(1024, dtype=torch.bfloat16, device="cpu"),
"decoder_action_out_proj_w": torch.zeros(1024, 32, dtype=torch.bfloat16, device="cpu"),
"decoder_action_out_proj_b": torch.zeros(32, dtype=torch.bfloat16, device="cpu"),
"decoder_action_in_proj_w": torch.zeros(32, 1024, dtype=torch.bfloat16, device="cpu"),
"decoder_action_in_proj_b": torch.zeros(1024, dtype=torch.bfloat16, device="cpu"),

"decoder_action_fused_out_proj_w": torch.zeros(num_steps, 1024, 32, dtype = torch.bfloat16, device = "cpu"),
"decoder_action_fused_out_proj_b": torch.zeros(num_steps, 32, dtype = torch.bfloat16, device = "cpu"),


"decoder_gate_attn": torch.zeros((num_steps, 18, 1024), dtype = torch.bfloat16, device = "cpu"),

"decoder_gate_ffn": torch.zeros((num_steps, 18, 1024), dtype = torch.bfloat16, device = "cpu"),


"language_embeds": torch.zeros(prompt_len, 2048, dtype=torch.bfloat16, device="cpu"),
}

convert_weights_pi05(weights, dump_weights)
convert_weights_pi05(weights, dump_weights, time_embeds)
weights['embedding_weight'].copy_(embedding_weight_torch)
weights['language_embeds'].copy_(language_embeds)
weights['decoder_time_embeds'].copy_(time_embeds)
with open(args.output, 'wb') as f:
pickle.dump(weights, f)
print(f"Successfully converted Pi0.5 weights to {args.output}")
Loading