Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion examples/puzzletron/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ dependency-ordered campaign and whole-campaign resume.
Candidate evaluation can be part of a post-MIP campaign flow, where metrics,
selection, materialization, and report lineage remain connected. Configure
that route with the
[post-MIP pipeline guide](docs/post_mip_pipeline.md#downstream-evaluation).
[post-MIP pipeline guide](docs/post_mip_pipeline.md#evaluate-saved-checkpoints).

For standalone evaluation in the Puzzletron worker environment, choose the
route that matches the checkpoint and task:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,16 +44,39 @@ post_mip:
materialized:
type: materialize
input: best_lm
checkpoint_eval:
type: downstream_evaluation
input: materialized
failure_policy: strict
config: &ifeval_smoke_config
tasks: [ifeval]
limit: 2
batch_size: 1
timeout_seconds: 600
reasoning_parser: qwen3
dtype: bfloat16
gpu_memory_utilization: 0.5
max_model_len: 4096
topology:
tensor_parallel_size: 1
pipeline_parallel_size: 1
data_parallel_size: 1
prefill_context_parallel_size: 1
decode_context_parallel_size: 1
enable_expert_parallel: false
distributed_executor_backend: mp
gpu_group_size: 1
serving:
type: aiperf
input: materialized
input: checkpoint_eval
config:
input_tokens: 128
output_tokens: 32
concurrency: [1]
request_count: 4
use_server_token_count: true
benchmark_timeout: 900
readiness_timeout: 300
benchmark_timeout: 300
allow_aiperf_v011_online_tokenizer_resolution: true
topology:
tensor_parallel_size: 1
Expand Down Expand Up @@ -84,9 +107,14 @@ post_mip:
global_batch_size: 1
local_batch_size: 1
checkpoint_every_steps: 2
post_kd_checkpoint_eval:
type: downstream_evaluation
input: short_kd
failure_policy: strict
config: *ifeval_smoke_config
final_eval:
type: evaluation
input: short_kd
input: post_kd_checkpoint_eval
config:
eval_samples: 2
block_size: 512
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,30 @@ post_mip:
materialized:
type: materialize
input: best_vlm_loss
checkpoint_eval:
type: downstream_evaluation
input: materialized
failure_policy: strict
config: &realworldqa_smoke_config
profile: qwen35_vlm_realworldqa
batch_size: 1
timeout_seconds: 600
dtype: bfloat16
gpu_memory_utilization: 0.5
max_model_len: 4096
limit_mm_per_prompt: {image: 1}
topology:
tensor_parallel_size: 1
pipeline_parallel_size: 1
data_parallel_size: 1
prefill_context_parallel_size: 1
decode_context_parallel_size: 1
enable_expert_parallel: false
distributed_executor_backend: mp
gpu_group_size: 1
vlm_serving:
type: aiperf
input: materialized
input: checkpoint_eval
config:
endpoint_type: chat
input_tokens: 100
Expand All @@ -65,7 +86,8 @@ post_mip:
request_count: 1
extra_inputs: {min_tokens: 80}
use_server_token_count: true
benchmark_timeout: 900
readiness_timeout: 300
benchmark_timeout: 300
topology:
tensor_parallel_size: 1
pipeline_parallel_size: 1
Expand Down Expand Up @@ -100,9 +122,14 @@ post_mip:
global_batch_size: 1
local_batch_size: 1
checkpoint_every_steps: 2
post_kd_checkpoint_eval:
type: downstream_evaluation
input: short_vlm_kd
failure_policy: strict
config: *realworldqa_smoke_config
final_image_eval:
type: evaluation
input: short_vlm_kd
input: post_kd_checkpoint_eval
config:
eval_samples: 2
block_size: 512
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,10 @@ execution:
post.params-90.online_eval: {strategy: sharded, instances: 1}
post.params-90.best_lm: {strategy: single, instances: 1}
post.params-90.materialized: {strategy: sharded, instances: 1}
post.params-90.checkpoint_eval: {strategy: sharded, instances: 1}
post.params-90.serving: {strategy: sharded, instances: 1}
post.params-90.fastest: {strategy: single, instances: 1}
post.params-90.short_kd: {strategy: sharded, instances: 1}
post.params-90.post_kd_checkpoint_eval: {strategy: sharded, instances: 1}
post.params-90.final_eval: {strategy: sharded, instances: 1}
post.params-90.best: {strategy: single, instances: 1}
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,10 @@ execution:
post.params-90.image_eval: {strategy: sharded, instances: 1}
post.params-90.best_vlm_loss: {strategy: single, instances: 1}
post.params-90.materialized: {strategy: sharded, instances: 1}
post.params-90.checkpoint_eval: {strategy: sharded, instances: 1}
post.params-90.vlm_serving: {strategy: sharded, instances: 1}
post.params-90.fastest_vlm: {strategy: single, instances: 1}
post.params-90.short_vlm_kd: {strategy: sharded, instances: 1}
post.params-90.post_kd_checkpoint_eval: {strategy: sharded, instances: 1}
post.params-90.final_image_eval: {strategy: sharded, instances: 1}
post.params-90.best: {strategy: single, instances: 1}
2 changes: 1 addition & 1 deletion examples/puzzletron/docs/checkpoint_evaluation.md
Original file line number Diff line number Diff line change
Expand Up @@ -117,4 +117,4 @@ identity when applicable. For `lmms-eval`, retain `command.json`,
generated YAML and launcher result artifacts.

To evaluate candidates as part of a pruning campaign, use
[downstream evaluation](post_mip_pipeline.md#downstream-evaluation) instead.
[saved-checkpoint evaluation](post_mip_pipeline.md#evaluate-saved-checkpoints) instead.
27 changes: 23 additions & 4 deletions examples/puzzletron/docs/post_mip_pipeline.md
Original file line number Diff line number Diff line change
Expand Up @@ -178,15 +178,34 @@ Selection still follows `input`; `model_source` only chooses the artifact operat
on. This supports a long KD run selected using short-KD/PTQ results but restarted
from the original candidate.

## Downstream evaluation
## Evaluate saved checkpoints

`downstream_evaluation` adapts the generic
[checkpoint evaluator](checkpoint_evaluation.md) to materialized campaign
candidates and publishes their task metrics. Add it after a `materialize` node;
the linked example config shows the complete flow. Use the standalone
[checkpoint evaluator](checkpoint_evaluation.md) to campaign checkpoints and
publishes their task metrics. Add it after any checkpoint-producing node, such
as `materialize` or `global_kd`. Use the standalone
[checkpoint evaluation](checkpoint_evaluation.md) command when campaign
lineage, filtering, and reports are not needed.

Materialization writes a reloadable Hugging Face checkpoint directory. Its
configuration records the realized per-layer block sizes, and its safetensors
contain the physically sliced weights. Evaluation passes that saved directory
unchanged to a fresh `lmms-eval` process backed by vLLM. It does not convert the
AnyModel instance back into an AutoModel instance. Modality-specific profiles
may prepare pinned task adapters and offline dataset snapshots first, but they
delegate checkpoint execution and completion validation to the same evaluator.

Global KD publishes a consolidated Hugging Face checkpoint and preserves the
realized pruning configuration and required tokenizer or processor assets. A
downstream evaluation node after KD therefore uses the same checkpoint contract
as one after materialization. The Qwen 3.5 text and VLM smoke flows evaluate the
selected checkpoint both before and after their short KD stage.

The post-MIP graph does not treat the teacher as a candidate revision. To compare
a distilled student with its teacher, evaluate the teacher separately with the
same task, evaluator version, dataset revision, prompt settings, and sample
limit, then compare those metrics with the post-KD node's metrics.

## Filters

`top_k` accepts one integer or separate homogeneous/heterogeneous quotas.
Expand Down
33 changes: 22 additions & 11 deletions examples/puzzletron/docs/qwen3p5_0p8b_smoke.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,15 @@
The checked-in `full_smoke` recipe runs a small end-to-end test of text-only
pruning for Qwen 3.5 0.8B. It searches the FFN intermediate sizes
`[3072, 2048]`, evaluates the candidates, saves the two strongest candidates as
physical checkpoints, and measures their serving performance with AIPerf. It
then distills the candidate with higher measured output-token throughput for
two steps, evaluates it again, and selects the final checkpoint. The recipe
pins the public checkpoint revision so repeated runs use the same starting
model.
physical checkpoints, and reloads each saved directory through vLLM for two
IFEval samples. It then measures both checkpoints with AIPerf, distills the
candidate with higher measured output-token throughput for two steps, evaluates
the resulting checkpoint with another two IFEval samples and the internal
two-sample LM-loss check, and selects the final checkpoint. The recipe pins the
public checkpoint revision so repeated runs use the same starting model. See
[evaluate saved checkpoints](post_mip_pipeline.md#evaluate-saved-checkpoints)
for how both Hugging Face directories are loaded without an AnyModel-to-AutoModel
conversion.

These small budgets check that the complete workflow runs and resumes
correctly. They do not establish model quality or production throughput.
Expand Down Expand Up @@ -61,9 +65,16 @@ python examples/puzzletron/orchestrate.py \
--stage full
```

The checked-in flow deliberately uses two evaluation samples per candidate,
four AIPerf requests per serving candidate, and two distillation steps. These
budgets validate lifecycle correctness, comparative serving selection, and
resumability; they are not quality or throughput claims. Final acceptance must
reload the selected checkpoint, verify the cumulative report, and confirm that
the resume submits no work for completed stages.
The checked-in flow deliberately uses two candidate-evaluation samples, two
IFEval samples, four AIPerf requests per serving candidate, and two
distillation steps. These budgets validate lifecycle correctness, comparative
serving selection, and resumability; they are not quality or throughput claims.
The worker environment must provide the [pinned evaluator
installation](checkpoint_evaluation.md#quick-start). IFEval task data must be
fetchable from each worker or already present in its Hugging Face cache.

After completion, inspect the `checkpoint_eval` and `post_kd_checkpoint_eval`
nodes under `artifacts/post_mip/nodes`. Their summaries must name the corresponding
pre-KD and post-KD checkpoints, report two effective IFEval samples, and contain
finite metrics. Also verify the cumulative report and confirm that resuming
submits no work for completed stages.
54 changes: 42 additions & 12 deletions examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,19 @@ The checked-in `full_vlm_smoke` recipe runs a small end-to-end test of
vision-language pruning for the public `Qwen/Qwen3.5-0.8B` checkpoint. It uses
real image-conversation examples to search the FFN intermediate sizes
`[3072, 2048]`, evaluate the candidates, and save the two strongest candidates
as physical checkpoints. It measures both checkpoints with 1-, 6-, and
12-image AIPerf requests, distills the candidate with the highest measured
12-image throughput for two steps, and runs a final image-and-text evaluation.
The immutable revision in `model.yaml` ensures that repeated runs use the same
starting model.
as physical checkpoints. It reloads each saved directory through vLLM for two
RealWorldQA image samples using the pinned Qwen 3.5 VLM evaluation profile. The
profile verifies the evaluator revision and immutable offline dataset snapshot,
strips inherited Hub credentials, and records preflight provenance before
delegating execution to the shared checkpoint evaluator. The workflow then
measures both checkpoints with 1-, 6-, and 12-image AIPerf requests, distills
the candidate with the highest measured 12-image throughput for two steps, and
runs the pinned RealWorldQA benchmark and internal image-and-text evaluation on
the resulting checkpoint. The immutable revision in `model.yaml` ensures that
repeated runs use the same starting model. See
[evaluate saved checkpoints](post_mip_pipeline.md#evaluate-saved-checkpoints)
for the direct pre-KD and post-KD reload paths; no AnyModel-to-AutoModel
conversion occurs.

These small budgets check that the complete workflow runs and resumes
correctly. They do not establish model quality or production throughput.
Expand All @@ -18,7 +26,11 @@ correctly. They do not establish model quality or production throughput.
Prepare the setup and worker environments described in
[environment setup](environment_setup.md). The worker environment must provide
ModelOpt, NeMo AutoModel's Qwen 3.5 VLM support, the Puzzletron requirements,
and the reviewed AIPerf/vLLM runtime selected by your runner contract.
the pinned `lmms-eval` dependency, and the reviewed AIPerf/vLLM runtime selected
by your runner contract. The worker-visible Hugging Face cache must already
contain the pinned RealWorldQA snapshot. Populate it as described in
[cache benchmark data](vlm_checkpoint_evaluation.md#cache-benchmark-data); the
profile verifies the local snapshot and evaluates it offline.

The worker-visible Hugging Face cache must contain, or be allowed to fetch,
`Qwen/Qwen3.5-0.8B` at the pinned revision in
Expand Down Expand Up @@ -128,11 +140,11 @@ python examples/puzzletron/orchestrate.py \

Confirm that all enabled model stages use one GPU, image-backed stages resolve
`data.modality=multimodal`, and no text tokenization stage is present. Confirm
that two quality candidates reach `post.params-90.vlm_serving`, which declares
a `chat` workload with 1, 6, and 12 1280x720 images per request rather than a
text-only serving proxy. The `fastest_vlm` filter selects one candidate from
the 12-image throughput metric before KD. Then launch by omitting only
`--dry-run`:
that two quality candidates reach `post.params-90.checkpoint_eval` for the
bounded RealWorldQA run and then `post.params-90.vlm_serving`, which declares a
`chat` workload with 1, 6, and 12 1280x720 images per request rather than a
text-only serving proxy. The `fastest_vlm` filter selects one candidate from the
12-image throughput metric before KD. Then launch by omitting only `--dry-run`:

```bash
python examples/puzzletron/orchestrate.py \
Expand All @@ -157,7 +169,14 @@ canonical stage summaries:
- width scoring and VLM KD processed real image tensors and report a nonzero
vision-forward count;
- sorting and physical slicing equivalence passed at the configured tolerance;
- the selected checkpoint reloads after materialization and after VLM KD;
- each materialized pre-KD checkpoint has a successful `checkpoint_eval`
summary whose `checkpoint` field names that saved artifact, whose
RealWorldQA sample count equals two, and whose metrics are finite; the
evaluation root's `profile.json` records the pinned dataset revision and
offline preflight;
- the selected post-KD checkpoint has a successful `post_kd_checkpoint_eval`
summary for two RealWorldQA samples with finite metrics, and the internal
final image evaluation reloads the same checkpoint;
- the two-step KD summary contains finite main CE/KD and MTP CE/KD metrics plus
nonzero trainable-group gradient evidence;
- AIPerf completes every 1-, 6-, and 12-image chat workload for both retained
Expand Down Expand Up @@ -195,3 +214,14 @@ vision path, and comparative selection while keeping the example bounded; it
does not isolate vision-encoder latency or establish production performance.
Increase request count and concurrency in a separate reviewed performance run
before making throughput claims.

The real-checkpoint lifecycle test is opt-in and is not part of default pytest
or routine CI smoke execution. Point it at the populated cache:

```bash
PUZZLETRON_VLM_BENCHMARK_HF_HOME=/path/to/hf-home \
pytest --run-manual tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_vlm_smoke.py
```

The checkpoint contract is documented in
[evaluate saved checkpoints](post_mip_pipeline.md#evaluate-saved-checkpoints).
18 changes: 18 additions & 0 deletions examples/puzzletron/evaluation/vlm/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Public interfaces for Puzzletron VLM evaluation."""

from .evaluator import *
Loading
Loading