File size: 6,336 Bytes
2ca6bc7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Export pinned Decision 3.0 text inference; vision export is separate."""
import argparse
import json
from pathlib import Path
import torch
import torch.nn.functional as F
from coreai_torch import TorchConverter, get_decomp_table
from safetensors.torch import load_file
from transformers import Qwen3_5Model
from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb_vision
from transformers.vision_utils import (
get_vision_interpolation_indices_and_weights,
get_vision_position_ids,
)
class DecisionText(torch.nn.Module):
def __init__(self, backbone, readout):
super().__init__()
self.language_model = backbone.language_model
self.register_buffer("readout", readout.float())
def forward(
self, input_ids, attention_mask, position_ids, visual_embeddings, visual_mask
):
embeddings = self.language_model.embed_tokens(input_ids)
embeddings = torch.where(
visual_mask, visual_embeddings.to(embeddings.dtype), embeddings
)
hidden = self.language_model(
inputs_embeds=embeddings,
attention_mask={
"full_attention": attention_mask[:, None, None, :].bool(),
"linear_attention": attention_mask,
},
position_ids=position_ids,
use_cache=False,
).last_hidden_state
return hidden[:, -1].float() @ self.readout.T
class DecisionVision(torch.nn.Module):
"""The unchanged vision weights at a fixed one-image 224px grid."""
def __init__(self, visual, side=224):
super().__init__()
self.visual = visual
grid = torch.tensor(
[[1, side // visual.config.patch_size, side // visual.config.patch_size]]
)
indices, weights = get_vision_interpolation_indices_and_weights(
grid,
num_grid_per_side=visual.num_grid_per_side,
mode=visual.interpolation_mode,
align_corners=visual.interpolation_align_corners,
spatial_merge_size=visual.config.spatial_merge_size,
kwargs={},
)
self.register_buffer("indices", indices)
self.register_buffer("weights", weights)
self.register_buffer(
"positions",
get_vision_position_ids(grid, visual.spatial_merge_size, kwargs={}),
)
def forward(self, pixel_values):
v = self.visual
hidden = F.linear(
pixel_values.to(v.patch_embed.proj.weight.dtype),
v.patch_embed.proj.weight.flatten(1),
v.patch_embed.proj.bias,
)
hidden = hidden + (v.pos_embed(self.indices) * self.weights[:, :, None]).sum(
1
).to(hidden.dtype)
cos, sin = v.rotary_pos_emb(hidden, self.positions)
for block in v.blocks:
x = block.norm1(hidden)
attn = block.attn
q, k, value = (
attn.qkv(x)
.reshape(x.shape[0], 3, attn.num_heads, -1)
.permute(1, 0, 2, 3)
.unbind(0)
)
q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
y = F.scaled_dot_product_attention(
q.transpose(0, 1).unsqueeze(0),
k.transpose(0, 1).unsqueeze(0),
value.transpose(0, 1).unsqueeze(0),
scale=attn.scaling,
)
hidden = hidden + attn.proj(
y.squeeze(0).transpose(0, 1).reshape(x.shape[0], -1)
)
hidden = hidden + block.mlp(block.norm2(hidden))
return v.merger(hidden).float()
def save(model, sample, path):
with torch.no_grad():
exported = torch.export.export(model, sample, strict=False)
print("Decomposing", flush=True)
exported = exported.run_decompositions(get_decomp_table())
print("Converting", flush=True)
program = TorchConverter().add_exported_program(exported).to_coreai()
program.save_asset(path)
return exported.graph_signature.user_outputs[0]
def main():
p = argparse.ArgumentParser()
p.add_argument("--source", default="upstream/d3-edge")
p.add_argument("--output", default="bundles/d3-edge-CoreAI")
p.add_argument("--length", type=int, default=512)
p.add_argument("--precision", choices=["bfloat16", "float32"], default="bfloat16")
a = p.parse_args()
torch.set_num_threads(4)
source = Path(a.source)
output = Path(a.output)
output.mkdir(parents=True, exist_ok=True)
print("Loading backbone", flush=True)
backbone = Qwen3_5Model.from_pretrained(
source, dtype=getattr(torch, a.precision), attn_implementation="sdpa"
).eval()
model = DecisionText(
backbone, load_file(str(source / "readout.safetensors"))["weight"]
).eval()
hidden = backbone.config.text_config.hidden_size
sample = (
torch.ones(1, a.length, dtype=torch.int32),
torch.ones(1, a.length, dtype=torch.int32),
torch.arange(a.length, dtype=torch.int32).view(1, 1, -1).expand(4, 1, -1),
torch.zeros(1, a.length, hidden),
torch.zeros(1, a.length, 1, dtype=torch.bool),
)
print("Exporting", flush=True)
out = save(model, sample, output / "decision.aimodel")
vision = DecisionVision(backbone.visual).eval()
patches = (224 // backbone.config.vision_config.patch_size) ** 2
width = (
3
* backbone.config.vision_config.temporal_patch_size
* backbone.config.vision_config.patch_size**2
)
vision_out = save(vision, (torch.zeros(patches, width),), output / "vision.aimodel")
(output / "manifest.json").write_text(
json.dumps(
{
"source": source.name,
"sequence_length": a.length,
"hidden_size": hidden,
"inputs": [
"input_ids",
"attention_mask",
"position_ids",
"visual_embeddings",
"visual_mask",
],
"output": out,
"vision_output": vision_out,
"image_side": 224,
"precision": a.precision,
},
indent=2,
)
+ "\n"
)
print("Saved", flush=True)
if __name__ == "__main__":
main()
|