Views
No views yet
1from vllm_omni.entrypoints.omni import Omni
2
3omni = Omni(
4 model="zhengyuansu/bagel-tiny-random",
5 stage_configs_path="path/to/bagel_sharedmemory_2gpu_ci.yaml",
6 custom_pipeline_args={
7 "pipeline_class": "examples.flowgrpo_trainer.vllm_omni.pipeline_bagel.BagelPipelineWithLogProb"
8 },
9)
10
11params_list = omni.default_sampling_params_list
12params_list[1].num_inference_steps = 10
13params_list[1].extra_args = {"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}
14
15outputs = list(omni.generate(
16 prompts=[{"prompt": "a cute cat", "modalities": ["image"]}],
17 sampling_params_list=params_list,
18))1"""Create a tiny-random BAGEL model for CI testing.
2
3Reads real BAGEL-7B-MoT checkpoint weight names, creates matching tiny random
4tensors with scaled-down dimensions. VAE architecture is hardcoded in vllm-omni
5and cannot be shrunk, so VAE weights are kept at full size.
6
7Usage:
8 python scripts/create_tiny_bagel.py --source ByteDance-Seed/BAGEL-7B-MoT
9"""
10
11import argparse
12import json
13import os
14import re
15import shutil
16
17import torch
18from safetensors import safe_open
19from safetensors.torch import save_file
20
21# LLM/ViT dimension shrinkage
22EMA_DIM_MAP = {
23 3584: 64, # LLM hidden_size
24 18944: 128, # LLM intermediate_size
25 1152: 64, # ViT hidden_size
26 4304: 128, # ViT intermediate_size
27 128: 32, # head_dim
28 512: 64, # kv_proj dim
29}
30
31# VAE: keep original dims (architecture is hardcoded in vllm-omni)
32VAE_DIM_MAP = {}
33
34MAX_LLM_LAYERS = 1
35MAX_VIT_LAYERS = 1
36
37
38def shrink_dims(shape, dim_map):
39 return [dim_map.get(d, d) for d in shape]
40
41
42def create_tiny_configs(source_dir, output_dir):
43 with open(os.path.join(source_dir, "config.json")) as f:
44 config = json.load(f)
45
46 llm = config["llm_config"]
47 llm["hidden_size"] = 64
48 llm["num_hidden_layers"] = MAX_LLM_LAYERS
49 llm["num_attention_heads"] = 2
50 llm["num_key_value_heads"] = 2
51 llm["intermediate_size"] = 128
52 llm["max_position_embeddings"] = 4096
53 llm["max_window_layers"] = MAX_LLM_LAYERS
54
55 vit = config["vit_config"]
56 vit["hidden_size"] = 64
57 vit["num_hidden_layers"] = MAX_VIT_LAYERS
58 vit["num_attention_heads"] = 2
59 vit["intermediate_size"] = 128
60
61 with open(os.path.join(output_dir, "config.json"), "w") as f:
62 json.dump(config, f, indent=4)
63
64 llm_standalone = dict(llm)
65 llm_standalone["qk_norm"] = True
66 llm_standalone["tie_word_embeddings"] = False
67 with open(os.path.join(output_dir, "llm_config.json"), "w") as f:
68 json.dump(llm_standalone, f, indent=4)
69
70 with open(os.path.join(output_dir, "vit_config.json"), "w") as f:
71 json.dump(dict(vit), f, indent=4)
72
73 return config
74
75
76def create_tiny_weights(source_path, dim_map, max_layers, seed=42):
77 gen = torch.Generator().manual_seed(seed)
78 weights = {}
79 is_vae = "ae" in os.path.basename(source_path).lower()
80 dtype = torch.float32 if is_vae else torch.bfloat16
81
82 with safe_open(source_path, framework="pt") as f:
83 for name in f.keys():
84 m = re.search(r"\.layers\.(\d+)\.", name)
85 if m:
86 idx = int(m.group(1))
87 for pattern, limit in max_layers.items():
88 if pattern in name and idx >= limit:
89 break
90 else:
91 pass
92 if m and any(p in name for p in max_layers) and idx >= max_layers.get(
93 next((p for p in max_layers if p in name), ""), 999
94 ):
95 continue
96
97 real_shape = list(f.get_tensor(name).shape)
98 tiny_shape = shrink_dims(real_shape, dim_map)
99
100 if "norm" in name and len(tiny_shape) == 1:
101 weights[name] = torch.ones(tiny_shape, dtype=dtype)
102 else:
103 weights[name] = torch.randn(tiny_shape, generator=gen, dtype=dtype) * 0.02
104
105 return weights
106
107
108def main():
109 parser = argparse.ArgumentParser()
110 parser.add_argument("--source", default="ByteDance-Seed/BAGEL-7B-MoT")
111 parser.add_argument("--output", default=os.path.expanduser("~/models/tiny-random/BAGEL-7B-MoT"))
112 args = parser.parse_args()
113
114 source_dir = args.source
115 if not os.path.exists(os.path.join(source_dir, "config.json")):
116 from huggingface_hub import snapshot_download
117 source_dir = snapshot_download(source_dir)
118
119 output_dir = args.output
120 os.makedirs(output_dir, exist_ok=True)
121
122 create_tiny_configs(source_dir, output_dir)
123
124 for fname in ["generation_config.json", "preprocessor_config.json", "tokenizer.json",
125 "tokenizer_config.json", "vocab.json", "merges.txt"]:
126 src = os.path.join(source_dir, fname)
127 if os.path.exists(src):
128 shutil.copy2(src, os.path.join(output_dir, fname))
129
130 ema = create_tiny_weights(
131 os.path.join(source_dir, "ema.safetensors"),
132 dim_map=EMA_DIM_MAP,
133 max_layers={"language_model": MAX_LLM_LAYERS, "vit_model": MAX_VIT_LAYERS},
134 seed=42,
135 )
136 save_file(ema, os.path.join(output_dir, "ema.safetensors"))
137
138 vae = create_tiny_weights(
139 os.path.join(source_dir, "ae.safetensors"),
140 dim_map=VAE_DIM_MAP,
141 max_layers={},
142 seed=43,
143 )
144 save_file(vae, os.path.join(output_dir, "ae.safetensors"))
145
146 weight_map = {k: "ema.safetensors" for k in ema}
147 weight_map.update({k: "ae.safetensors" for k in vae})
148 total_size = sum(t.numel() * t.element_size() for t in ema.values())
149 total_size += sum(t.numel() * t.element_size() for t in vae.values())
150 with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
151 json.dump({"metadata": {"total_size": total_size}, "weight_map": weight_map}, f, indent=4)
152
153
154if __name__ == "__main__":
155 main()