From ba8b808391821ff4454d0c41ed35fd5a93ae396e Mon Sep 17 00:00:00 2001 From: Sidney Swift <158200036+sidneyswift@users.noreply.github.com> Date: Mon, 9 Nov 2026 20:23:54 -0500 Subject: [PATCH] Revert grouped GPU decoding (api#1234) In production, grouped decoding doubled full_report instead of cutting it: on the same song, wall time went from 181 s to 367 s and billed seconds from 179 to 365. generate_batch goes back to answering prompts one at a time, as verified after api#1233. This reverts merge commit 026071c1a4ddc045738533801febbce07098644e. 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, 48 insertions(+), 245 deletions(-) diff --git a/modal/README.md b/modal/README.md index bab134980..99667dcda 100644 --- a/modal/README.md +++ b/modal/README.md @@ -31,20 +31,7 @@ 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. - -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 +the whole song. 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 700fabda6..207e921a3 100644 --- a/modal/serve_music_flamingo.py +++ b/modal/serve_music_flamingo.py @@ -31,10 +31,6 @@ 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 @@ -217,74 +213,38 @@ def parse_batch_prompts(prompts, defaults: dict) -> list: return jobs -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: +def run_batch(jobs: list, audio_path: str, run_prompt) -> list: """ Answers each job (from parse_batch_prompts) against `audio_path`, or - 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. + 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. """ import os import time - answers = [None] * len(jobs) - for group in group_jobs(jobs, group_size): + answers = [] + for job in jobs: started = time.time() - clips = {} + clip_path = extract_clip(audio_path, *job["clip"]) if job["clip"] else None try: - 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 - ]) + response = run_prompt( + job["prompt"], clip_path or audio_path, job["max_new_tokens"], + job["temperature"], job["top_p"], job["do_sample"], + ) finally: - for clip_path in clips.values(): + if clip_path: os.unlink(clip_path) - 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 + 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) return answers @@ -429,7 +389,11 @@ def run_prompt(self, prompt, audio_path, max_tokens, temperature, top_p, do_samp answer. Shared by /generate and /generate_batch. """ # ---- Build the conversation ---- - conversation = build_conversation(prompt, audio_path) + # 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}] # ---- Run inference ---- inputs = self.processor.apply_chat_template( @@ -459,62 +423,6 @@ 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): """ @@ -582,9 +490,8 @@ 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 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 + downloaded (and transcoded if needed) once; each prompt then runs + sequentially as its own conversation against the same file. The API charges one base price per request, so batching saves the base on every prompt after the first. @@ -601,17 +508,13 @@ 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) - "group_size": 8 (optional, 1-8, default 8; - 1 runs prompts one at a time) + "do_sample": false (optional, default false) } Returns: { "answers": [{"id": "genre", "response": "...", "elapsed_seconds": 1.2, - "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) + "clip_start": 45.2, "clip_end": 60.4}], (clip echoed when one was cut) "elapsed_seconds": 9.87 (whole batch, download included) } @@ -638,25 +541,18 @@ 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_prompts, group_size) + answers = run_batch(jobs, tmp_path, self.run_prompt) finally: os.unlink(tmp_path) @@ -678,36 +574,3 @@ 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 5c1c79a85..2819433ac 100644 --- a/modal/test_serve_music_flamingo.py +++ b/modal/test_serve_music_flamingo.py @@ -17,7 +17,6 @@ check_generation_settings, cut_clip, extract_clip, - group_jobs, is_number, parse_batch_prompts, run_batch, @@ -165,44 +164,21 @@ def test_a_clip_window_at_infinity_or_too_large_for_a_float_is_refused(): class FakeModel: - """Stands in for the GPU: records each group it is given and what audio each prompt had.""" + """Stands in for the GPU: records each call and what audio it was given.""" def __init__(self, fail_on=None): - self.groups = [] + self.calls = [] self.fail_on = 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): + 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: raise RuntimeError("model failed") - 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]] + return f"answer to {prompt}" def test_run_batch_answers_in_order_with_each_prompts_own_settings_and_clip(tmp_path): @@ -230,38 +206,15 @@ def test_run_batch_answers_in_order_with_each_prompts_own_settings_and_clip(tmp_ assert os.path.exists(source) -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": "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): +def test_run_batch_deletes_the_clip_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) + jobs = parse_batch_prompts( + [{"id": "hook", "prompt": "Describe.", "clip_start": 0.5, "clip_end": 1.5}], DEFAULTS + ) model = FakeModel(fail_on="Describe.") with pytest.raises(RuntimeError, match="model failed"): run_batch(jobs, source, model) - assert all(not os.path.exists(call["audio_path"]) for call in model.calls) + assert not os.path.exists(model.calls[0]["audio_path"]) assert os.path.exists(source)