Skip to content

Revert grouped GPU decoding (api#1234) #1241

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sswiftdrtv merged 1 commit into from
Nov 10, 2026
Merged

Revert grouped GPU decoding (api#1234) #1241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 1 addition & 14 deletions modal/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
201 changes: 32 additions & 169 deletions modal/serve_music_flamingo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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.

Expand All @@ -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)
}

Expand All @@ -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)

Expand All @@ -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)
Loading
Loading