diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index ca548c08911..7e7b6fc2ac3 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -184,8 +184,6 @@ def export_dflash( mmproj loader and vision export contract to add ``vision_encoder``. """ validate_dflash_export_options(backend, use_offgraph_kv_cache) - if use_offgraph_kv_cache and backend == "cuda" and mmproj is not None: - raise ValueError("the CUDA batching DFlash export is text-only") import executorch.extension.llm.export.gguf # noqa: F401 import executorch.extension.llm.export.int4 # noqa: F401 @@ -243,6 +241,22 @@ def export_dflash( activation_dtype=activation_dtype, ) + vision_model = None + pos_embed_table = None + if mmproj is not None: + from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import ( + load_mmproj_vision_model, + ) + + print("=" * 60) + print("Loading vision encoder...") + print("=" * 60) + vision_model, pos_embed_table, _ = load_mmproj_vision_model( + mmproj, + backend=backend, + activation_dtype=activation_dtype, + ) + if use_offgraph_kv_cache and backend == "cuda": from executorch.examples.models.muse_glimmer.export.export_solo import ( DEFAULT_FORWARD_SIZES, @@ -261,25 +275,12 @@ def export_dflash( dflash_draft=draft_model, max_draft_owners=max_draft_owners, draft_max_cells=draft_max_cells, + vision_model=vision_model, + pos_embed_table=pos_embed_table, + max_vision_patches=max_vision_patches, ) return - vision_model = None - pos_embed_table = None - if mmproj is not None: - from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import ( - load_mmproj_vision_model, - ) - - print("=" * 60) - print("Loading vision encoder...") - print("=" * 60) - vision_model, pos_embed_table, _ = load_mmproj_vision_model( - mmproj, - backend=backend, - activation_dtype=activation_dtype, - ) - backend_export = _export_dflash_mlx if backend == "mlx" else _export_dflash_cuda backend_export( target_model, diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index f3d80a418a1..8391dbcfac0 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -133,8 +133,6 @@ def export_and_lower( if not 1 <= max_prefill_chunk <= config.max_seq_len: raise ValueError("max_prefill_chunk must be within [1, max_seq_len]") if backend == "cuda" and use_offgraph_kv_cache: - if vision_model is not None: - raise ValueError("the CUDA batching export is text-only") export_cuda_batching( model, config, @@ -143,6 +141,9 @@ def export_and_lower( max_cells=max_cells, forward_sizes=forward_sizes or DEFAULT_FORWARD_SIZES, device_sampling=sample, + vision_model=vision_model, + pos_embed_table=pos_embed_table, + max_vision_patches=max_vision_patches, ) elif backend == "cuda": _export_cuda( @@ -257,6 +258,10 @@ def step_width_spec(method: str) -> bytes: return b"1:0" +# Image prompts reach the batching cache through this method. +EMBEDDINGS_METHOD = "forward_from_embeddings" + + def step_forward( model: MuseGlimmerModel, tokens: torch.Tensor, @@ -268,7 +273,18 @@ def step_forward( The whole step runs through the decoder so every token's K/V reaches the cache; only the rows the batch samples from go through the LM head. """ - x = model.embed_text(tokens) + return embeddings_step_forward( + model, model.embed_text(tokens), input_pos, logits_to_keep + ) + + +def embeddings_step_forward( + model: MuseGlimmerModel, + x: torch.Tensor, + input_pos: torch.Tensor, + logits_to_keep: torch.Tensor, +) -> torch.Tensor: + """``step_forward`` from decoder inputs [1, T, dim], e.g. spliced images.""" taps = None if hasattr(model, "_dflash_target_layer_ids"): sliding_mask, global_mask = model._build_masks(input_pos) @@ -349,9 +365,16 @@ def export_cuda_batching( # noqa: C901 dflash_draft: nn.Module | None = None, max_draft_owners: int = 4, draft_max_cells: int | None = None, + vision_model: nn.Module | None = None, + pos_embed_table: torch.Tensor | None = None, + max_vision_patches: int = 16384, ): """Exports the forward methods and writes the artifact. + With ``vision_model`` the token forwards are unchanged, and image prompts run + through three more methods: ``vision_encoder``, ``embed_text`` and a dynamic + ``forward_from_embeddings`` over the same cache and weights. + Returns the ExecutorchProgramManager so callers can inspect it. """ import executorch.backends.cuda.quantize_op_dispatch # noqa: F401 @@ -462,10 +485,52 @@ def export_cuda_batching( # noqa: C901 dynamic_shapes=({1: width}, {0: width}, {0: rows}), strict=True, ) + if vision_model is not None: + embeddings_width = Dim( + "embeddings_width", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens + ) + with ( + common.BoundMethodForward( + model, types.MethodType(embeddings_step_forward, model) + ), + torch.no_grad(), + ): + print( + f"Exporting {EMBEDDINGS_METHOD} (T in [{_CUDA_MIN_PREFILL_CHUNK}, " + f"{max_step_tokens}])..." + ) + programs[EMBEDDINGS_METHOD] = export( + model, + ( + torch.zeros( + (1, max_step_tokens, config.dim), dtype=torch.bfloat16 + ), + torch.arange(max_step_tokens, dtype=torch.long), + torch.zeros(max_step_tokens, dtype=torch.long), + ), + dynamic_shapes=({1: embeddings_width}, {0: embeddings_width}, {0: rows}), + strict=True, + ) + with common.BoundMethodForward(model, model.embed_text), torch.no_grad(): + print("Exporting embed_text...") + programs["embed_text"] = export( + model, + (torch.zeros((1, max_step_tokens), dtype=torch.long),), + dynamic_shapes=( + {1: Dim("embed_width", min=1, max=max_step_tokens)}, + ), + strict=True, + ) + programs["vision_encoder"] = common.export_vision_encoder( + vision_model, pos_embed_table, max_vision_patches + ) + del vision_model del model gc.collect() torch.cuda.empty_cache() - forward_methods = list(programs) + forward_methods = [ + name for name in programs if name not in ("embed_text", "vision_encoder") + ] if dflash_draft is not None: programs.update(draft_programs) del draft_programs, dflash_draft @@ -524,8 +589,8 @@ def partitioner(name: str) -> CudaPartitioner: max_prefill=max_step_tokens, activation_dtype=torch.bfloat16, mutable_buffer_metadata=None, - has_vision=False, - max_vision_patches=0, + has_vision=pos_embed_table is not None, + max_vision_patches=max_vision_patches, ) constant_methods.update(geometry) constant_methods.update(write_max_context_len(config.max_seq_len)) @@ -564,7 +629,7 @@ def partitioner(name: str) -> CudaPartitioner: del edge gc.collect() - common.save_pte(et_program, output_dir, None) + common.save_pte(et_program, output_dir, pos_embed_table) if torch.cuda.is_available(): peak_mb = torch.cuda.max_memory_allocated() / (1024**2) print(f"EXPORT_GPU_PEAK_MEMORY_MB: {peak_mb:.1f}") @@ -1113,10 +1178,6 @@ def main() -> None: args = parser.parse_args() _validate_offgraph_prefill(parser, args) - if args.backend == "cuda" and args.use_offgraph_kv_cache and args.mmproj: - parser.error( - "--use-offgraph-kv-cache on CUDA exports a text-only model; omit --mmproj." - ) if args.backend == "cuda" and not torch.cuda.is_available(): parser.error("CUDA is required for the cuda backend.") diff --git a/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py b/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py index fcdc306e9d4..b181e72ee8c 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py +++ b/examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py @@ -184,6 +184,96 @@ def setUp(self) -> None: if not torch.cuda.is_available(): self.skipTest("CUDA required") + def test_vision_adds_embedding_methods_over_the_same_cache(self) -> None: + """Images add vision, embed_text and one dynamic embeddings forward.""" + try: + import gguf # noqa: F401 + except ImportError: + self.skipTest("gguf package required") + from executorch.backends.cuda.passes.lower_offgraph_kv import ( + OFFGRAPH_KV_COMPILE_SPEC, + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + ) + from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import ( + load_mmproj_vision_model, + ) + from executorch.examples.models.muse_glimmer.tests.test_pipeline import ( + _mmproj_vision_config, + build_mmproj_gguf, + ) + from executorch.exir.lowered_backend_module import get_lowered_submodules + + with ( + tempfile.TemporaryDirectory() as ckpt_dir, + tempfile.TemporaryDirectory() as out_dir, + ): + save_checkpoint(ckpt_dir) + model, config = load_prequantized_model( + ckpt_dir, max_seq_len=TINY_CONFIG.max_seq_len + ) + mmproj = os.path.join(ckpt_dir, "mmproj.gguf") + build_mmproj_gguf(mmproj) + vision, pos_table, _ = load_mmproj_vision_model( + mmproj, config=_mmproj_vision_config(), backend="cuda" + ) + program = export_cuda_batching( + model, + config, + out_dir, + max_step_tokens=self.MAX_STEP, + forward_sizes=self.FORWARD_SIZES, + max_cells=self.MAX_CELLS, + vision_model=vision, + pos_embed_table=pos_table, + max_vision_patches=64, + ) + self.assertTrue(os.path.exists(os.path.join(out_dir, "pos_embed.bin"))) + methods = set(program.methods) + self.assertLessEqual( + {"forward_from_embeddings", "embed_text", "vision_encoder"}, methods + ) + embeddings = program.exported_program("forward_from_embeddings") + lowered = get_lowered_submodules(embeddings.graph_module) + self.assertEqual(len(lowered), 1) + specs = {spec.key: spec.value for spec in lowered[0][1].compile_specs} + # Same cache and step-width contract as the token forwards. + others = get_lowered_submodules( + program.exported_program("forward_others").graph_module + )[0][1].compile_specs + self.assertEqual( + specs[OFFGRAPH_KV_COMPILE_SPEC], + {s.key: s.value for s in others}[OFFGRAPH_KV_COMPILE_SPEC], + ) + self.assertEqual(specs[OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC], b"1:0") + for name in ("embed_text", "vision_encoder"): + sub = get_lowered_submodules( + program.exported_program(name).graph_module + ) + self.assertEqual(len(sub), 1, name) + self.assertNotIn( + OFFGRAPH_KV_COMPILE_SPEC, + {spec.key for spec in sub[0][1].compile_specs}, + name, + ) + + from executorch.runtime import Runtime, Verification + + loaded = Runtime.get().load_program( + os.path.join(out_dir, "model.pte"), verification=Verification.Minimal + ) + meta = loaded.metadata("forward_from_embeddings") + self.assertEqual(meta.num_inputs(), 3) + self.assertEqual( + meta.input_tensor_meta(0).sizes(), (1, self.MAX_STEP, TINY_CONFIG.dim) + ) + self.assertEqual( + loaded.load_method("has_vision_encoder").execute([])[0], True + ) + self.assertEqual( + loaded.load_method("get_vision_hidden_size").execute([])[0], + TINY_CONFIG.dim, + ) + def test_exports_forward_methods_over_a_cell_cache(self) -> None: from executorch.backends.cuda.passes.lower_offgraph_kv import ( OFFGRAPH_KV_COMPILE_SPEC,