From 84d6a4d3775e0a5f39825a20b365759c535a1347 Mon Sep 17 00:00:00 2001 From: Sidney Swift <158200036+sidneyswift@users.noreply.github.com> Date: Thu, 5 Nov 2026 01:08:57 -0500 Subject: [PATCH] Modal: decode a batch's prompts together on the GPU generate_batch answered its prompts one at a time, which leaves an A100 mostly idle on an 8B model; full_report now runs its 13 sections in one batch, so the time adds up. Prompts that share temperature, top_p and do_sample are now decoded together, up to 8 per generate() call (the processor pads on the left; each answer is cut to its own max_new_tokens). - Each answer's elapsed_seconds is its share of its group's time, so the answers still add up to the time billed. - Fails safe: one prompt takes the /generate path; out of GPU memory splits the group in half and retries; any other grouped error falls back to one prompt at a time. - Optional request field group_size (1-8, default 8); 1 restores one-at-a-time decoding. - bench entrypoint compares the two on a real song and audio URL. Not yet run on a GPU (the local Modal workspace is paused); verify on production by comparing a full_report against the one-at-a-time baseline. Co-Authored-By: Claude Opus 5.5 --- modal/README.md | 15 ++- modal/serve_music_flamingo.py | 201 ++++++++++++++++++++++++----- modal/test_serve_music_flamingo.py | 77 ++++++++--- 3 files changed, 245 insertions(+), 48 deletions(-) diff --git a/modal/README.md b/modal/README.md index 99667dcda..bab134980 100644 --- a/modal/README.md +++ b/modal/README.md @@ -31,7 +31,20 @@ both or neither) to hear only that stretch, cut fresh from the one download for each prompt, and its own `temperature`, `top_p` or `do_sample` to override the batch's settings for that prompt alone. An answer from a clip echoes `clip_start`/`clip_end`, so a caller can tell it was not answered from -the whole song. The validation and the batch loop are pure Python and tested +the whole song. + +Prompts that share settings are decoded together, up to `GENERATE_GROUP_SIZE` +(8) per `generate()` call, instead of one at a time; one prompt at a time +leaves the A100 mostly idle. Each answer's `elapsed_seconds` is its share of +its group's time, so the answers still add up to the time billed. A group that +runs out of GPU memory is split in half and retried, and any other error in a +grouped call falls back to one prompt at a time. A request may pass +`group_size` (1–8); 1 restores one-at-a-time decoding. To compare the two on a +real song in a workspace you can spend GPU time in: + +```bash +modal run modal/serve_music_flamingo.py::bench --audio-url-file url.txt --prompts-file prompts.json +``` The validation and the batch loop are pure Python and tested locally (ffmpeg and ffprobe must be on the PATH): ```bash diff --git a/modal/serve_music_flamingo.py b/modal/serve_music_flamingo.py index 207e921a3..700fabda6 100644 --- a/modal/serve_music_flamingo.py +++ b/modal/serve_music_flamingo.py @@ -31,6 +31,10 @@ MAX_NEW_TOKENS = 512 # default max tokens to generate MAX_NEW_TOKENS_LIMIT = 2048 # hard cap per prompt on /generate_batch MAX_BATCH_PROMPTS = 40 # most prompts one /generate_batch request may carry +# Prompts that share temperature, top_p and do_sample are decoded together, this +# many per generate() call. One prompt at a time leaves the GPU mostly idle; a +# group that runs out of GPU memory is split in half and retried. +GENERATE_GROUP_SIZE = 8 MODEL_CACHE_DIR = "/model_cache" # where model weights live inside the container AUDIO_DOWNLOAD_TIMEOUT = 120 # seconds; an unreachable host must not hold a GPU # Formats the processor's audio loader reads directly. Anything else is @@ -213,38 +217,74 @@ def parse_batch_prompts(prompts, defaults: dict) -> list: return jobs -def run_batch(jobs: list, audio_path: str, run_prompt) -> list: +def build_conversation(prompt: str, audio_path) -> list: + """ + One user turn in Music Flamingo's chat format: the prompt, then the audio + at `audio_path` when given. + """ + content = [{"type": "text", "text": prompt}] + if audio_path: + content.append({"type": "audio", "path": audio_path}) + return [{"role": "user", "content": content}] + + +def group_jobs(jobs: list, size: int = GENERATE_GROUP_SIZE) -> list: + """ + Splits jobs into groups that can share one generate() call: the same + temperature, top_p and do_sample, at most `size` each. Returns lists of + job indexes; groups and the jobs in them keep the order jobs first appear. + """ + groups, filling = [], {} + for index, job in enumerate(jobs): + key = (job["temperature"], job["top_p"], job["do_sample"]) + group = filling.get(key) + if group is None or len(group) == size: + group = [] + groups.append(group) + filling[key] = group + group.append(index) + return groups + + +def run_batch(jobs: list, audio_path: str, run_prompts, group_size: int = GENERATE_GROUP_SIZE) -> list: """ Answers each job (from parse_batch_prompts) against `audio_path`, or - against a clip cut from it when the job has one; every clip is deleted - after its prompt, whether or not the prompt succeeds. `run_prompt(prompt, - audio_path, max_new_tokens, temperature, top_p, do_sample)` returns the - answer text. Returns the answers in job order; an answer from a clip - echoes the window so the caller can tell it came from the clip. + against a clip cut from it when the job has one. Jobs that share settings + run together, `group_size` per call to `run_prompts(requests)`, which takes + (prompt, audio_path, max_new_tokens, temperature, top_p, do_sample) tuples + with the same settings and returns their answers in order. A group's clips + are deleted after it runs, whether or not it succeeds. + + Returns the answers in job order. Each answer's `elapsed_seconds` is its + share of its group's time, so the answers add up to the time spent; an + answer from a clip echoes the window so the caller can tell where it came from. """ import os import time - answers = [] - for job in jobs: + answers = [None] * len(jobs) + for group in group_jobs(jobs, group_size): started = time.time() - clip_path = extract_clip(audio_path, *job["clip"]) if job["clip"] else None + clips = {} try: - response = run_prompt( - job["prompt"], clip_path or audio_path, job["max_new_tokens"], - job["temperature"], job["top_p"], job["do_sample"], - ) + for index in group: + if jobs[index]["clip"]: + clips[index] = extract_clip(audio_path, *jobs[index]["clip"]) + responses = run_prompts([ + (jobs[index]["prompt"], clips.get(index, audio_path), jobs[index]["max_new_tokens"], + jobs[index]["temperature"], jobs[index]["top_p"], jobs[index]["do_sample"]) + for index in group + ]) finally: - if clip_path: + for clip_path in clips.values(): os.unlink(clip_path) - answer = { - "id": job["id"], - "response": response, - "elapsed_seconds": round(time.time() - started, 2), - } - if job["clip"]: - answer["clip_start"], answer["clip_end"] = job["clip"] - answers.append(answer) + share = round((time.time() - started) / len(group), 2) + for index, response in zip(group, responses): + job = jobs[index] + answer = {"id": job["id"], "response": response, "elapsed_seconds": share} + if job["clip"]: + answer["clip_start"], answer["clip_end"] = job["clip"] + answers[index] = answer return answers @@ -389,11 +429,7 @@ def run_prompt(self, prompt, audio_path, max_tokens, temperature, top_p, do_samp answer. Shared by /generate and /generate_batch. """ # ---- Build the conversation ---- - # Music Flamingo uses a chat-style format with "user" and "assistant" roles - content = [{"type": "text", "text": prompt}] - if audio_path: - content.append({"type": "audio", "path": audio_path}) - conversation = [{"role": "user", "content": content}] + conversation = build_conversation(prompt, audio_path) # ---- Run inference ---- inputs = self.processor.apply_chat_template( @@ -423,6 +459,62 @@ def run_prompt(self, prompt, audio_path, max_tokens, temperature, top_p, do_samp ) return decoded[0] + def run_prompts(self, requests: list) -> list: + """ + Runs prompts that share temperature, top_p and do_sample as ONE + generate() call, each as its own conversation with its own audio, and + returns the answers in order, each cut to its own max_new_tokens. + `requests` holds (prompt, audio_path, max_new_tokens, temperature, + top_p, do_sample) tuples. Used by /generate_batch. + + Fails safe: a single prompt takes the same path as /generate; if the + GPU runs out of memory the group is split in half and retried; any + other error in a grouped call falls back to one prompt at a time. + """ + import torch + + if len(requests) == 1: + return [self.run_prompt(*requests[0])] + try: + return self._generate_together(requests) + except torch.cuda.OutOfMemoryError: + torch.cuda.empty_cache() + half = len(requests) // 2 + return self.run_prompts(requests[:half]) + self.run_prompts(requests[half:]) + except Exception as err: + print(f"[generate_batch] grouped generation of {len(requests)} failed ({err!r}); " + "answering one at a time") + return [self.run_prompt(*request) for request in requests] + + def _generate_together(self, requests: list) -> list: + """One generate() call for `requests` (see run_prompts); the processor pads on the left.""" + conversations = [build_conversation(prompt, audio_path) for prompt, audio_path, *_ in requests] + inputs = self.processor.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=True, + return_dict=True, + ).to(self.model.device) + if "input_features" in inputs: + inputs["input_features"] = inputs["input_features"].to(self.model.dtype) + + limits = [max_tokens for _, _, max_tokens, *_ in requests] + _, _, _, temperature, top_p, do_sample = requests[0] + outputs = self.model.generate( + **inputs, + max_new_tokens=max(limits), + temperature=temperature, + top_p=top_p, + do_sample=do_sample, + ) + + # Every row's new tokens start after the padded prompt; cut each to its own limit + start = inputs.input_ids.shape[1] + return self.processor.batch_decode( + [outputs[row, start:start + limit] for row, limit in enumerate(limits)], + skip_special_tokens=True, + ) + @modal.fastapi_endpoint(method="POST", docs=True, requires_proxy_auth=True) def generate(self, request: dict): """ @@ -490,8 +582,9 @@ def generate(self, request: dict): def generate_batch(self, request: dict): """ Many short prompts about ONE audio file in one request. The audio is - downloaded (and transcoded if needed) once; each prompt then runs - sequentially as its own conversation against the same file. The API + downloaded (and transcoded if needed) once; each prompt runs as its own + conversation against the same file, and prompts that share settings + are decoded together, up to `group_size` per generate() call. The API charges one base price per request, so batching saves the base on every prompt after the first. @@ -508,13 +601,17 @@ def generate_batch(self, request: dict): ], "temperature": 0.2, (optional, default 1.0) "top_p": 1.0, (optional, default 1.0) - "do_sample": false (optional, default false) + "do_sample": false, (optional, default false) + "group_size": 8 (optional, 1-8, default 8; + 1 runs prompts one at a time) } Returns: { "answers": [{"id": "genre", "response": "...", "elapsed_seconds": 1.2, - "clip_start": 45.2, "clip_end": 60.4}], (clip echoed when one was cut) + "clip_start": 45.2, "clip_end": 60.4}], (clip echoed when one was cut; + seconds are the answer's share + of its group's time) "elapsed_seconds": 9.87 (whole batch, download included) } @@ -541,18 +638,25 @@ def reject(detail: str): "top_p": request.get("top_p", 1.0), "do_sample": request.get("do_sample", False), } + group_size = request.get("group_size", GENERATE_GROUP_SIZE) try: jobs = parse_batch_prompts(request.get("prompts"), defaults) check_generation_settings(defaults) except ValueError as err: reject(str(err)) + if ( + isinstance(group_size, bool) + or not isinstance(group_size, int) + or not 1 <= group_size <= GENERATE_GROUP_SIZE + ): + reject(f"group_size must be 1-{GENERATE_GROUP_SIZE}") # ---- Download (and transcode) the audio ONCE ---- tmp_path = self.download_audio(audio_url) # ---- Run every prompt against the same file (or a clip cut from it) ---- try: - answers = run_batch(jobs, tmp_path, self.run_prompt) + answers = run_batch(jobs, tmp_path, self.run_prompts, group_size) finally: os.unlink(tmp_path) @@ -574,3 +678,36 @@ def main(): }) print(f"\nModel response:\n{result['response']}") print(f"Took {result['elapsed_seconds']}s") + + +@app.local_entrypoint() +def bench(audio_url_file: str, prompts_file: str, out_file: str = ""): + """ + Times one /generate_batch request with prompts one at a time (group_size + 1) and decoded together (the default), on the same audio and prompts, and + counts the answers that match. Run it in a workspace you can spend GPU + time in; it never deploys: + modal run modal/serve_music_flamingo.py::bench --audio-url-file url.txt --prompts-file prompts.json + + `prompts_file` is a JSON list of /generate_batch prompts. + """ + import json + + audio_url = open(audio_url_file).read().strip() + prompts = json.load(open(prompts_file)) + flamingo = MusicFlamingo() + # Warm the GPU so the first timed run does not pay for it + flamingo.generate_batch.remote( + {"audio_url": audio_url, "prompts": [{"id": "warm", "prompt": "Genre?", "max_new_tokens": 8}]} + ) + results = {} + for size in (1, GENERATE_GROUP_SIZE): + results[size] = flamingo.generate_batch.remote( + {"audio_url": audio_url, "prompts": prompts, "group_size": size} + ) + print(f"group_size {size}: {results[size]['elapsed_seconds']} s for {len(prompts)} prompts") + pairs = zip(results[1]["answers"], results[GENERATE_GROUP_SIZE]["answers"]) + same = sum(a["response"].strip() == b["response"].strip() for a, b in pairs) + print(f"identical answers: {same} of {len(prompts)}") + if out_file: + json.dump(results, open(out_file, "w"), indent=1) diff --git a/modal/test_serve_music_flamingo.py b/modal/test_serve_music_flamingo.py index 2819433ac..5c1c79a85 100644 --- a/modal/test_serve_music_flamingo.py +++ b/modal/test_serve_music_flamingo.py @@ -17,6 +17,7 @@ check_generation_settings, cut_clip, extract_clip, + group_jobs, is_number, parse_batch_prompts, run_batch, @@ -164,21 +165,44 @@ def test_a_clip_window_at_infinity_or_too_large_for_a_float_is_refused(): class FakeModel: - """Stands in for the GPU: records each call and what audio it was given.""" + """Stands in for the GPU: records each group it is given and what audio each prompt had.""" def __init__(self, fail_on=None): - self.calls = [] + self.groups = [] self.fail_on = fail_on - def __call__(self, prompt, audio_path, max_new_tokens, temperature, top_p, do_sample): - self.calls.append({ - "prompt": prompt, "audio_path": audio_path, "existed": os.path.exists(audio_path), - "seconds": duration(audio_path), "max_new_tokens": max_new_tokens, - "temperature": temperature, "top_p": top_p, "do_sample": do_sample, - }) - if prompt == self.fail_on: + @property + def calls(self): + return [call for group in self.groups for call in group] + + def __call__(self, requests): + group = [] + for prompt, audio_path, max_new_tokens, temperature, top_p, do_sample in requests: + group.append({ + "prompt": prompt, "audio_path": audio_path, "existed": os.path.exists(audio_path), + "seconds": duration(audio_path), "max_new_tokens": max_new_tokens, + "temperature": temperature, "top_p": top_p, "do_sample": do_sample, + }) + self.groups.append(group) + if any(call["prompt"] == self.fail_on for call in group): raise RuntimeError("model failed") - return f"answer to {prompt}" + return [f"answer to {call['prompt']}" for call in group] + + +def jobs_with(*settings): + """One plain job per settings tuple (temperature, do_sample), ids j0, j1...""" + return parse_batch_prompts([ + {"id": f"j{i}", "prompt": f"q{i}", "temperature": t, "do_sample": s} + for i, (t, s) in enumerate(settings) + ], DEFAULTS) + + +def test_group_jobs_keeps_settings_apart_and_caps_each_group(): + jobs = jobs_with((1.0, False), (0.7, True), (1.0, False), (1.0, False), (0.7, True)) + + assert group_jobs(jobs, size=2) == [[0, 2], [1, 4], [3]] + assert group_jobs(jobs, size=8) == [[0, 2, 3], [1, 4]] + assert group_jobs(jobs, size=1) == [[0], [1], [2], [3], [4]] def test_run_batch_answers_in_order_with_each_prompts_own_settings_and_clip(tmp_path): @@ -206,15 +230,38 @@ def test_run_batch_answers_in_order_with_each_prompts_own_settings_and_clip(tmp_ assert os.path.exists(source) -def test_run_batch_deletes_the_clip_when_the_model_fails(tmp_path): +def test_run_batch_decodes_prompts_that_share_settings_together(tmp_path): source = tone(str(tmp_path / "song.wav"), 3) - jobs = parse_batch_prompts( - [{"id": "hook", "prompt": "Describe.", "clip_start": 0.5, "clip_end": 1.5}], DEFAULTS - ) + jobs = parse_batch_prompts([ + {"id": "a", "prompt": "qa", "clip_start": 0, "clip_end": 1}, + {"id": "b", "prompt": "qb", "temperature": 0.7, "do_sample": True}, + {"id": "c", "prompt": "qc", "clip_start": 1, "clip_end": 2}, + {"id": "d", "prompt": "qd"}, + ], DEFAULTS) + model = FakeModel() + + answers = run_batch(jobs, source, model, group_size=8) + + assert [[c["prompt"] for c in g] for g in model.groups] == [["qa", "qc", "qd"], ["qb"]] + assert [a["id"] for a in answers] == ["a", "b", "c", "d"] + assert [a["response"] for a in answers] == ["answer to qa", "answer to qb", "answer to qc", "answer to qd"] + # The two clips in the group were separate files, both gone afterwards + clip_a, clip_c = model.groups[0][0]["audio_path"], model.groups[0][1]["audio_path"] + assert clip_a != clip_c and not os.path.exists(clip_a) and not os.path.exists(clip_c) + # A group's prompts share its time equally + assert answers[0]["elapsed_seconds"] == answers[2]["elapsed_seconds"] == answers[3]["elapsed_seconds"] + + +def test_run_batch_deletes_the_clips_when_the_model_fails(tmp_path): + source = tone(str(tmp_path / "song.wav"), 3) + jobs = parse_batch_prompts([ + {"id": "hook", "prompt": "Describe.", "clip_start": 0.5, "clip_end": 1.5}, + {"id": "lift", "prompt": "Lift.", "clip_start": 1.5, "clip_end": 2.5}, + ], DEFAULTS) model = FakeModel(fail_on="Describe.") with pytest.raises(RuntimeError, match="model failed"): run_batch(jobs, source, model) - assert not os.path.exists(model.calls[0]["audio_path"]) + assert all(not os.path.exists(call["audio_path"]) for call in model.calls) assert os.path.exists(source)