Skip to content

Commit 745acd2

Browse files
authored
[executorch][muse] Export vision with the CUDA batching contract
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * #24066 * #24065 * __->__ #24064 * #24063 * #24062 export_cuda_batching optionally adds vision_encoder, embed_text and a dynamic forward_from_embeddings over the same off-graph KV contract, and writes pos_embed.bin next to the program. Differential Revision: D125987292 Pull Request resolved: #24064
1 parent 6bc4128 commit 745acd2

3 files changed

Lines changed: 217 additions & 33 deletions

File tree

‎examples/models/muse-glimmer/export/export_dflash.py‎

Lines changed: 26 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -188,8 +188,6 @@ def export_dflash(
188188
mmproj loader and vision export contract to add ``vision_encoder``.
189189
"""
190190
validate_dflash_export_options(backend, use_offgraph_kv_cache)
191-
if use_offgraph_kv_cache and backend == "cuda" and mmproj is not None:
192-
raise ValueError("the CUDA batching DFlash export is text-only")
193191

194192
import executorch.extension.llm.export.gguf # noqa: F401
195193
import executorch.extension.llm.export.int4 # noqa: F401
@@ -247,6 +245,29 @@ def export_dflash(
247245
activation_dtype=activation_dtype,
248246
)
249247

248+
vision_model = None
249+
pos_embed_table = None
250+
if mmproj is not None:
251+
from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
252+
load_mmproj_vision_model,
253+
)
254+
255+
print("=" * 60)
256+
print("Loading vision encoder...")
257+
print("=" * 60)
258+
vision_model, pos_embed_table, _ = load_mmproj_vision_model(
259+
mmproj,
260+
backend=backend,
261+
activation_dtype=activation_dtype,
262+
)
263+
264+
if (
265+
draft_model is not None
266+
and use_offgraph_kv_cache
267+
and backend == "cuda"
268+
and vision_model is not None
269+
):
270+
raise ValueError("the CUDA batching DFlash export does not support vision yet")
250271
if use_offgraph_kv_cache and backend == "cuda":
251272
from executorch.examples.models.muse_glimmer.export.export_solo import (
252273
DEFAULT_FORWARD_SIZES,
@@ -265,25 +286,12 @@ def export_dflash(
265286
dflash_draft=draft_model,
266287
max_draft_owners=max_draft_owners,
267288
draft_max_cells=draft_max_cells,
289+
vision_model=vision_model,
290+
pos_embed_table=pos_embed_table,
291+
max_vision_patches=max_vision_patches,
268292
)
269293
return
270294

271-
vision_model = None
272-
pos_embed_table = None
273-
if mmproj is not None:
274-
from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
275-
load_mmproj_vision_model,
276-
)
277-
278-
print("=" * 60)
279-
print("Loading vision encoder...")
280-
print("=" * 60)
281-
vision_model, pos_embed_table, _ = load_mmproj_vision_model(
282-
mmproj,
283-
backend=backend,
284-
activation_dtype=activation_dtype,
285-
)
286-
287295
backend_export = _export_dflash_mlx if backend == "mlx" else _export_dflash_cuda
288296
backend_export(
289297
target_model,

‎examples/models/muse-glimmer/export/export_solo.py‎

Lines changed: 85 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,10 @@
1111
1212
``--use-offgraph-kv-cache`` keeps the KV cache out of the program and exports
1313
the batching runtime's contract. On MLX that is the selected-logits decoder. On
14-
CUDA it is the text-only artifact CudaExecutor (backends/cuda/batching) drives;
15-
without the flag CUDA keeps its in-graph prefill and decode methods:
14+
CUDA it is the artifact CudaExecutor (backends/cuda/batching) drives: the text
15+
forwards plus, with ``--mmproj``, ``vision_encoder``/``embed_text``/
16+
``forward_from_embeddings``; without the flag CUDA keeps its in-graph prefill
17+
and decode methods:
1618
1719
- ``forward_{N}`` for each N in ``--forward-sizes`` (default 1, 2, 4, ..., 64):
1820
T = K = N, static, each captured into a CUDA graph at runtime, its quantized
@@ -133,8 +135,8 @@ def export_and_lower(
133135
if not 1 <= max_prefill_chunk <= config.max_seq_len:
134136
raise ValueError("max_prefill_chunk must be within [1, max_seq_len]")
135137
if backend == "cuda" and use_offgraph_kv_cache:
136-
if vision_model is not None:
137-
raise ValueError("the CUDA batching export is text-only")
138+
if vision_fp32_mm != "none":
139+
raise ValueError("CUDA batching does not support vision_fp32_mm")
138140
export_cuda_batching(
139141
model,
140142
config,
@@ -143,6 +145,9 @@ def export_and_lower(
143145
max_cells=max_cells,
144146
forward_sizes=forward_sizes or DEFAULT_FORWARD_SIZES,
145147
device_sampling=sample,
148+
vision_model=vision_model,
149+
pos_embed_table=pos_embed_table,
150+
max_vision_patches=max_vision_patches,
146151
)
147152
elif backend == "cuda":
148153
_export_cuda(
@@ -257,6 +262,10 @@ def step_width_spec(method: str) -> bytes:
257262
return b"1:0"
258263

259264

265+
# Image prompts reach the batching cache through this method.
266+
EMBEDDINGS_METHOD = "forward_from_embeddings"
267+
268+
260269
def step_forward(
261270
model: MuseGlimmerModel,
262271
tokens: torch.Tensor,
@@ -268,7 +277,18 @@ def step_forward(
268277
The whole step runs through the decoder so every token's K/V reaches the
269278
cache; only the rows the batch samples from go through the LM head.
270279
"""
271-
x = model.embed_text(tokens)
280+
return embeddings_step_forward(
281+
model, model.embed_text(tokens), input_pos, logits_to_keep
282+
)
283+
284+
285+
def embeddings_step_forward(
286+
model: MuseGlimmerModel,
287+
x: torch.Tensor,
288+
input_pos: torch.Tensor,
289+
logits_to_keep: torch.Tensor,
290+
) -> torch.Tensor:
291+
"""``step_forward`` from decoder inputs [1, T, dim], e.g. spliced images."""
272292
taps = None
273293
if hasattr(model, "_dflash_target_layer_ids"):
274294
sliding_mask, global_mask = model._build_masks(input_pos)
@@ -349,9 +369,16 @@ def export_cuda_batching( # noqa: C901
349369
dflash_draft: nn.Module | None = None,
350370
max_draft_owners: int = 4,
351371
draft_max_cells: int | None = None,
372+
vision_model: nn.Module | None = None,
373+
pos_embed_table: torch.Tensor | None = None,
374+
max_vision_patches: int = 16384,
352375
):
353376
"""Exports the forward methods and writes the artifact.
354377
378+
With ``vision_model`` the token forwards are unchanged, and image prompts run
379+
through three more methods: ``vision_encoder``, ``embed_text`` and a dynamic
380+
``forward_from_embeddings`` over the same cache and weights.
381+
355382
Returns the ExecutorchProgramManager so callers can inspect it.
356383
"""
357384
import executorch.backends.cuda.quantize_op_dispatch # noqa: F401
@@ -379,6 +406,10 @@ def export_cuda_batching( # noqa: C901
379406
)
380407
from torch.export import Dim, export
381408

409+
has_vision = vision_model is not None and pos_embed_table is not None
410+
if (vision_model is None) != (pos_embed_table is None):
411+
raise ValueError("vision_model and pos_embed_table must be given together")
412+
382413
if dflash_draft is not None:
383414
from executorch.examples.models.muse_glimmer.export.export_dflash import (
384415
prepare_cuda_packed_draft,
@@ -462,10 +493,52 @@ def export_cuda_batching( # noqa: C901
462493
dynamic_shapes=({1: width}, {0: width}, {0: rows}),
463494
strict=True,
464495
)
496+
if has_vision:
497+
embeddings_width = Dim(
498+
"embeddings_width", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens
499+
)
500+
with (
501+
common.BoundMethodForward(
502+
model, types.MethodType(embeddings_step_forward, model)
503+
),
504+
torch.no_grad(),
505+
):
506+
print(
507+
f"Exporting {EMBEDDINGS_METHOD} (T in [{_CUDA_MIN_PREFILL_CHUNK}, "
508+
f"{max_step_tokens}])..."
509+
)
510+
programs[EMBEDDINGS_METHOD] = export(
511+
model,
512+
(
513+
torch.zeros((1, max_step_tokens, config.dim), dtype=torch.bfloat16),
514+
torch.arange(max_step_tokens, dtype=torch.long),
515+
torch.zeros(max_step_tokens, dtype=torch.long),
516+
),
517+
dynamic_shapes=(
518+
{1: embeddings_width},
519+
{0: embeddings_width},
520+
{0: rows},
521+
),
522+
strict=True,
523+
)
524+
with common.BoundMethodForward(model, model.embed_text), torch.no_grad():
525+
print("Exporting embed_text...")
526+
programs["embed_text"] = export(
527+
model,
528+
(torch.zeros((1, max_step_tokens), dtype=torch.long),),
529+
dynamic_shapes=({1: Dim("embed_width", min=1, max=max_step_tokens)},),
530+
strict=True,
531+
)
532+
programs["vision_encoder"] = common.export_vision_encoder(
533+
vision_model, pos_embed_table, max_vision_patches
534+
)
535+
del vision_model
465536
del model
466537
gc.collect()
467538
torch.cuda.empty_cache()
468-
forward_methods = list(programs)
539+
forward_methods = [
540+
name for name in programs if name not in ("embed_text", "vision_encoder")
541+
]
469542
if dflash_draft is not None:
470543
programs.update(draft_programs)
471544
del draft_programs, dflash_draft
@@ -520,8 +593,8 @@ def partitioner(name: str) -> CudaPartitioner:
520593
max_prefill=max_step_tokens,
521594
activation_dtype=torch.bfloat16,
522595
mutable_buffer_metadata=None,
523-
has_vision=False,
524-
max_vision_patches=0,
596+
has_vision=has_vision,
597+
max_vision_patches=max_vision_patches,
525598
)
526599
constant_methods.update(geometry)
527600
constant_methods.update(write_max_context_len(config.max_seq_len))
@@ -560,7 +633,7 @@ def partitioner(name: str) -> CudaPartitioner:
560633
del edge
561634
gc.collect()
562635

563-
common.save_pte(et_program, output_dir, None)
636+
common.save_pte(et_program, output_dir, pos_embed_table)
564637
if torch.cuda.is_available():
565638
peak_mb = torch.cuda.max_memory_allocated() / (1024**2)
566639
print(f"EXPORT_GPU_PEAK_MEMORY_MB: {peak_mb:.1f}")
@@ -1040,8 +1113,9 @@ def main() -> None:
10401113
"--use-offgraph-kv-cache",
10411114
action="store_true",
10421115
help="Allocate the KV cache at runtime and export the batching runtime's "
1043-
"contract: MLX's selected-logits decoder, or CUDA's text-only "
1044-
"forward_{N} + forward_others over a cell-layout cache. Without it CUDA "
1116+
"contract: MLX's selected-logits decoder, or CUDA's "
1117+
"forward_{N} + forward_others over a cell-layout cache (plus the vision "
1118+
"methods with --mmproj). Without it CUDA "
10451119
"exports its in-graph prefill and decode methods.",
10461120
)
10471121
parser.add_argument(
@@ -1109,10 +1183,6 @@ def main() -> None:
11091183
args = parser.parse_args()
11101184

11111185
_validate_offgraph_prefill(parser, args)
1112-
if args.backend == "cuda" and args.use_offgraph_kv_cache and args.mmproj:
1113-
parser.error(
1114-
"--use-offgraph-kv-cache on CUDA exports a text-only model; omit --mmproj."
1115-
)
11161186

11171187
if args.backend == "cuda" and not torch.cuda.is_available():
11181188
parser.error("CUDA is required for the cuda backend.")

‎examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,22 @@ def test_cell_manifest_keeps_every_layer(self) -> None:
174174
self.assertEqual(manifest["max_write"], 8)
175175
self.assertEqual(len(manifest["layers"]), TINY_CONFIG.n_layers)
176176

177+
def test_cuda_batching_rejects_vision_fp32_mm(self) -> None:
178+
from executorch.examples.models.muse_glimmer.export.export_solo import (
179+
export_and_lower,
180+
)
181+
182+
model = build_random_tiny_model()
183+
with self.assertRaisesRegex(ValueError, "vision_fp32_mm"):
184+
export_and_lower(
185+
model,
186+
TINY_CONFIG,
187+
".",
188+
backend="cuda",
189+
use_offgraph_kv_cache=True,
190+
vision_fp32_mm="triton",
191+
)
192+
177193

178194
class BatchingExportTest(unittest.TestCase):
179195
MAX_STEP = 16
@@ -184,6 +200,96 @@ def setUp(self) -> None:
184200
if not torch.cuda.is_available():
185201
self.skipTest("CUDA required")
186202

203+
def test_vision_adds_embedding_methods_over_the_same_cache(self) -> None:
204+
"""Images add vision, embed_text and one dynamic embeddings forward."""
205+
try:
206+
import gguf # noqa: F401
207+
except ImportError:
208+
self.skipTest("gguf package required")
209+
from executorch.backends.cuda.passes.lower_offgraph_kv import (
210+
OFFGRAPH_KV_COMPILE_SPEC,
211+
OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC,
212+
)
213+
from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
214+
load_mmproj_vision_model,
215+
)
216+
from executorch.examples.models.muse_glimmer.tests.test_pipeline import (
217+
_mmproj_vision_config,
218+
build_mmproj_gguf,
219+
)
220+
from executorch.exir.lowered_backend_module import get_lowered_submodules
221+
222+
with (
223+
tempfile.TemporaryDirectory() as ckpt_dir,
224+
tempfile.TemporaryDirectory() as out_dir,
225+
):
226+
save_checkpoint(ckpt_dir)
227+
model, config = load_prequantized_model(
228+
ckpt_dir, max_seq_len=TINY_CONFIG.max_seq_len
229+
)
230+
mmproj = os.path.join(ckpt_dir, "mmproj.gguf")
231+
build_mmproj_gguf(mmproj)
232+
vision, pos_table, _ = load_mmproj_vision_model(
233+
mmproj, config=_mmproj_vision_config(), backend="cuda"
234+
)
235+
program = export_cuda_batching(
236+
model,
237+
config,
238+
out_dir,
239+
max_step_tokens=self.MAX_STEP,
240+
forward_sizes=self.FORWARD_SIZES,
241+
max_cells=self.MAX_CELLS,
242+
vision_model=vision,
243+
pos_embed_table=pos_table,
244+
max_vision_patches=64,
245+
)
246+
self.assertTrue(os.path.exists(os.path.join(out_dir, "pos_embed.bin")))
247+
methods = set(program.methods)
248+
self.assertLessEqual(
249+
{"forward_from_embeddings", "embed_text", "vision_encoder"}, methods
250+
)
251+
embeddings = program.exported_program("forward_from_embeddings")
252+
lowered = get_lowered_submodules(embeddings.graph_module)
253+
self.assertEqual(len(lowered), 1)
254+
specs = {spec.key: spec.value for spec in lowered[0][1].compile_specs}
255+
# Same cache and step-width contract as the token forwards.
256+
others = get_lowered_submodules(
257+
program.exported_program("forward_others").graph_module
258+
)[0][1].compile_specs
259+
self.assertEqual(
260+
specs[OFFGRAPH_KV_COMPILE_SPEC],
261+
{s.key: s.value for s in others}[OFFGRAPH_KV_COMPILE_SPEC],
262+
)
263+
self.assertEqual(specs[OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC], b"1:0")
264+
for name in ("embed_text", "vision_encoder"):
265+
sub = get_lowered_submodules(
266+
program.exported_program(name).graph_module
267+
)
268+
self.assertEqual(len(sub), 1, name)
269+
self.assertNotIn(
270+
OFFGRAPH_KV_COMPILE_SPEC,
271+
{spec.key for spec in sub[0][1].compile_specs},
272+
name,
273+
)
274+
275+
from executorch.runtime import Runtime, Verification
276+
277+
loaded = Runtime.get().load_program(
278+
os.path.join(out_dir, "model.pte"), verification=Verification.Minimal
279+
)
280+
meta = loaded.metadata("forward_from_embeddings")
281+
self.assertEqual(meta.num_inputs(), 3)
282+
self.assertEqual(
283+
meta.input_tensor_meta(0).sizes(), (1, self.MAX_STEP, TINY_CONFIG.dim)
284+
)
285+
self.assertEqual(
286+
loaded.load_method("has_vision_encoder").execute([])[0], True
287+
)
288+
self.assertEqual(
289+
loaded.load_method("get_vision_hidden_size").execute([])[0],
290+
TINY_CONFIG.dim,
291+
)
292+
187293
def test_exports_forward_methods_over_a_cell_cache(self) -> None:
188294
from executorch.backends.cuda.passes.lower_offgraph_kv import (
189295
OFFGRAPH_KV_COMPILE_SPEC,

0 commit comments

Comments
 (0)