Add batch infer feature and gpu allocation feature#136
Open
ZTWHHH wants to merge 1 commit into
Open
Conversation
`mmeval/run.py` evaluates a single model on a single dataset. This adds a thin scheduler on top of it so a whole set of (model x dataset) combinations can run from one command, with GPUs allocated automatically from each model's estimated VRAM. The framework's single-task path is untouched: `mmeval/run.py`, `mmeval/registry.py`, and `mmeval/utils/argparser.py` are unchanged. The feature is purely additive. - mmeval/batch_infer.py: orchestrator. Selects models (--models / --model-series, with type/quantization/api/modality filters and --sort-by), expands dataset specs (mmeval_hf@/evalkit@/local@/tsv@), and invokes run.py once per task. Estimates VRAM, picks the fewest GPUs that fit each model and fills the rest with parallel shards, skips models that don't fit, and prints a summary. Unrecognized flags are forwarded to run.py, so all generation and inference options work unchanged. --list-models previews the selection. - mmeval/utils/gpu.py: GPU discovery (CUDA_VISIBLE_DEVICES / nvidia-smi, incl. MIG), a parameter+KV-cache VRAM estimate, and the allocation policy. - mmeval/model_metadata.json: companion to registry.py, a JSON object keyed by model name adding the size/architecture/modality fields the scheduler needs (registry.py stays the source of truth for which models exist and how to run them). - scripts/update_model_metadata.py: refreshes the auto-derived metadata fields from HuggingFace and cross-checks the file against registry.py. - docs: a Batch Inference section in USAGE.md and a README feature note. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds a new batch inference workflow to run many model×dataset evaluations via mmeval/run.py, including VRAM-based GPU sizing/allocation driven by a curated model metadata file.
Changes:
- Added
mmeval/batch_infer.pyscheduler CLI that filters/sorts models, expands dataset specs, and invokesrun.pyper task with computed GPU parallelism. - Added GPU discovery + VRAM estimation/allocation helpers (
mmeval/utils/gpu.py) and a curated metadata file (mmeval/model_metadata.json), plus a refresh script (scripts/update_model_metadata.py). - Updated documentation (
docs/en/USAGE.md,README.md) to describe batch inference usage and metadata maintenance.
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/update_model_metadata.py | New script to refresh auto-derivable fields in model_metadata.json from HuggingFace and cross-check against the registry. |
| README.md | Adds a batch inference mention/link under the “Scalable” bullet. |
| mmeval/utils/gpu.py | New GPU visibility/VRAM detection and VRAM-based allocation utilities for batch scheduling. |
| mmeval/model_metadata.json | New curated model metadata used for filtering and VRAM estimation in batch inference. |
| mmeval/batch_infer.py | New batch scheduler CLI wrapping run.py, with filtering/sorting, dataset spec parsing, and GPU allocation. |
| docs/en/USAGE.md | Adds detailed “Batch Inference” documentation and examples. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+144
to
+152
| parsed = [float(v) for v in value] | ||
| if any(v < 0 for v in parsed): | ||
| raise ValueError("--gpu-memory values must be >= 0") | ||
| if len(parsed) < visible_gpu_count: | ||
| parsed.extend([0.0] * (visible_gpu_count - len(parsed))) | ||
| elif len(parsed) > visible_gpu_count: | ||
| print(f"Warning: --gpu-memory has {len(parsed)} values but only " | ||
| f"{visible_gpu_count} GPU(s) visible; dropping extras: {parsed[visible_gpu_count:]}") | ||
| return parsed[:visible_gpu_count] |
|
|
||
|
|
||
| def sort_models(models: list[dict], key: str | None) -> list[dict]: | ||
| """Order models for scheduling. None keeps registry order.""" |
Comment on lines
+294
to
+295
| p.add_argument("--sort-by", choices=["newest", "oldest", "largest", "smallest"], | ||
| help="Order models for scheduling (default: registry order).") |
| | `--quantization` | `false` | `true`, `false`, or `all` | | ||
| | `--api` | `false` | `false` = local only, `true` = API only, `all` = both | | ||
| | `--modalities` | -- | keep models supporting at least one listed modality | | ||
| | `--sort-by` | registry order | `newest`, `oldest`, `largest`, `smallest` | |
| "SmolVLM-Instruct": {"hf_path": "HuggingFaceTB/SmolVLM-Instruct", "series": "smolvlm", "model_type": "instruct", "api_model": false, "quantization": false, "modalities": ["multi_image_start", "text"], "parameter_count": 2246272880, "created_at": "2024-11-18", "torch_dtype": "bfloat16", "num_hidden_layers": 24, "hidden_size": 2048, "num_attention_heads": 32, "num_key_value_heads": 32, "max_position_embeddings": 16384, "image_tokens": 729, "max_dynamic_patch": null, "required_gpu_count": 1}, | ||
| "SmolVLM-Instruct-Base": {"hf_path": "HuggingFaceTB/SmolVLM-Base", "series": "smolvlm", "model_type": "pretrain", "api_model": false, "quantization": false, "modalities": ["multi_image_start", "text"], "parameter_count": 2246272880, "created_at": "2024-11-22", "torch_dtype": "bfloat16", "num_hidden_layers": 24, "hidden_size": 2048, "num_attention_heads": 32, "num_key_value_heads": 32, "max_position_embeddings": 16384, "image_tokens": 729, "max_dynamic_patch": null, "required_gpu_count": 1}, | ||
| "SmolVLM-Instruct-DPO": {"hf_path": "HuggingFaceTB/SmolVLM-Instruct-DPO", "series": "smolvlm", "model_type": "instruct", "api_model": false, "quantization": false, "modalities": ["multi_image_start", "text"], "parameter_count": null, "created_at": "2024-11-26", "torch_dtype": null, "num_hidden_layers": null, "hidden_size": null, "num_attention_heads": null, "num_key_value_heads": null, "max_position_embeddings": null, "image_tokens": null, "max_dynamic_patch": null, "required_gpu_count": 1}, | ||
| "SmolVLM-Sythetic": {"hf_path": "HuggingFaceTB/SmolVLM-Synthetic", "series": "smolvlm", "model_type": "instruct", "api_model": false, "quantization": false, "modalities": ["multi_image_start", "text"], "parameter_count": 2246272880, "created_at": "2024-11-22", "torch_dtype": "bfloat16", "num_hidden_layers": 24, "hidden_size": 2048, "num_attention_heads": 32, "num_key_value_heads": 32, "max_position_embeddings": 16384, "image_tokens": 729, "max_dynamic_patch": null, "required_gpu_count": 1}, |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
What does this PR do?
Type of Change
Testing
Test command used:
# Command you used to testChecklist
mmeval/registery.py(if new model)