diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/all_axis_base.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/all_axis_base.yaml new file mode 100644 index 00000000000..8e2d68bd38d --- /dev/null +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/all_axis_base.yaml @@ -0,0 +1,94 @@ +# @package _global_ + +defaults: + - runs/vlm_pruning_base + - _self_ + +# Bounded all-axis search shared by the maintained 4B lifecycle smoke and +# campaign. Each supported structural axis receives an isolated +# dynamic-versus-physical diagnostic before it can enter MIP selection. +puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT,puzzle_runs/qwen3p5_4b_vlm_smoke} +display_name: Qwen3p5 4B VLM all-axis search + +model: + automodel_backend: {attn: sdpa} + +embedding_pruning: + widths: [2560, 2400] + alignment: 32 + +pruning: + intermediate_size_list: [8704] + # [query heads, KV groups]: teacher, query-only, KV-only, and joint targets. + attn_heads_list: [[16, 4], [12, 4], [8, 2], [6, 2]] + attention_scored_axes: [kv_groups, q_heads_per_group] + gdn_scored_axes: + - gdn_key_groups + - gdn_value_heads_per_group + - gdn_key_head_dim + - gdn_value_head_dim + +sort: + deferred_axes: [] + +width_sanity: + hidden_width_diagnostic: true + axes: + - hidden_width + - kv_groups + - q_heads_per_group + - ffn_intermediate + - gdn_key_groups + - gdn_value_heads_per_group + - gdn_key_head_dim + - gdn_value_head_dim + one_case_per_axis: true + target_count_per_axis: 1 + layer_count: 1 + eval_samples: 2 + target_values: + hidden_width: 2400 + kv_groups: 2 + q_heads_per_group: 3 + ffn_intermediate: 8704 + gdn_key_groups: 14 + gdn_value_heads_per_group: 1 + gdn_key_head_dim: 112 + gdn_value_head_dim: 112 + +depth_importance: + enabled: true + max_removals: 1 + max_subblocks_to_remove: 1 + eval_samples: 2 + +search_space: + axes: + hidden_width: {enabled: true, teacher_value: 2560, values: [2400]} + kv_groups: {enabled: true, teacher_value: 4, values: [2]} + q_heads_per_group: {enabled: true, teacher_value: 4, values: [3]} + ffn_intermediate: {enabled: true, teacher_value: 9216, values: [8704]} + gdn_key_groups: {enabled: true, teacher_value: 16, values: [14]} + gdn_value_heads_per_group: {enabled: true, teacher_value: 2, values: [1]} + gdn_key_head_dim: {enabled: true, teacher_value: 128, values: [112]} + gdn_value_head_dim: {enabled: true, teacher_value: 128, values: [112]} + +mip: + depth_scenario_count: 2 + runs: + params-80: + search_space: &all_axis_search + depth: [0, 1] + embedding: [2560, 2400] + axes_default: all + axes: {ffn.intermediate_size: all} + solver: &all_axis_solver + backend: auto + num_solutions: 1 + min_hamming_distance: 1 + max_seconds_per_solution: 30 + homogeneous: &no_homogeneous {enabled: false} + memory-85: + search_space: *all_axis_search + solver: *all_axis_solver + homogeneous: *no_homogeneous diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/all_axis_kd_search.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/all_axis_kd_search.yaml new file mode 100644 index 00000000000..e76b2380fb2 --- /dev/null +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/all_axis_kd_search.yaml @@ -0,0 +1,341 @@ +# @package _global_ + +defaults: + - /families/qwen3_5/qwen3p5_4b/all_axis_base@_global_ + - /families/qwen3_5/qwen3p5_4b/vlm_quality_evaluation@_global_ + - _self_ + +# Illustrative Qwen 3.5 4B VLM all-axis campaign. The finite pre-KD image-loss +# screen admits exactly four distinct architectures. Every admitted candidate +# then receives the same frozen evaluation, serving smoke, and resumable +# 128-step KD exposure before final quality and serving comparison. +puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT} +display_name: Qwen3p5 4B VLM exact-four all-axis KD128 campaign + +data: + require_content_digest: true + calibration: + # One immutable 512-row order backs every candidate trajectory. Fail when + # the prepared dataset is shorter instead of cycling a mechanics packet. + num_samples: 512 + +embedding_pruning: + widths: [2560, 2400, 2240] + +pruning: + experiment_id: qwen35_4b_vlm_all_axis_campaign + intermediate_size_list: [8704, 8192] + eval_samples: 64 + seed: 42 + +sort_sanity: + eval_samples: 8 + +width_sanity: + target_count_per_axis: 2 + layer_count: 4 + eval_samples: 8 + +depth_importance: + max_removals: 2 + max_subblocks_to_remove: 2 + +replacement_scoring: + eval_samples: 32 + +search_space: + axes: + hidden_width: {enabled: true, teacher_value: 2560, values: [2400, 2240]} + kv_groups: {enabled: true, teacher_value: 4, values: [2]} + q_heads_per_group: {enabled: true, teacher_value: 4, values: [3]} + ffn_intermediate: {enabled: true, teacher_value: 9216, values: [8704, 8192]} + gdn_key_groups: {enabled: true, teacher_value: 16, values: [14]} + gdn_value_heads_per_group: {enabled: true, teacher_value: 2, values: [1]} + gdn_key_head_dim: {enabled: true, teacher_value: 128, values: [112]} + gdn_value_head_dim: {enabled: true, teacher_value: 128, values: [112]} + +global_distillation: + enabled: false + domain: vlm + automodel: + activation_checkpointing: true + force_hf: false + teacher_force_hf: false + student_force_hf: false + freeze_policy: vision_frozen + validation_enabled: false + objective: &memory_safe_kd_objective + main_ce: {weight: 1.0} + main_kd: {weight: 1.0, chunk_size: 64} + mtp_ce: {weight: 1.0} + mtp_kd: {weight: 1.0, chunk_size: 64} + +mip: + depth_scenario_count: 3 + runs: + params-80: false + memory-85: false + search-candidates: + objectives: + - metric: metrics.cosine_embedding_loss_hidden_states + direction: minimize + search_space: + depth: [0, 1, 2] + embedding: [2560, 2400, 2240] + axes_default: all + axes: {ffn.intermediate_size: all} + variants: + params-82: + constraints: + params: {max: 82%} + memory-85: + constraints: + memory: + at: + serving-default: {max: 85%} + solver: + backend: auto + num_solutions: 1 + min_hamming_distance: 1 + max_seconds_per_solution: 300 + homogeneous: {enabled: false} + teacher-reference: + objectives: + - metric: metrics.cosine_embedding_loss_hidden_states + direction: minimize + search_space: + depth: [0] + embedding: [2560] + axes_default: teacher + axes: {ffn.intermediate_size: teacher} + variants: + teacher: + constraints: + params: {min: 100%, max: 100%} + solver: + backend: auto + num_solutions: 1 + min_hamming_distance: 1 + max_seconds_per_solution: 30 + homogeneous: {enabled: false} + +post_mip: + flows: + candidate-evaluation: + source: + run: search-candidates + variants: all + objectives: all + nodes: + online_eval: + type: evaluation + failure_policy: strict + config: {eval_samples: 32, block_size: 512} + + # MIP architectures are deduplicated before this deterministic screen. + # Missing or non-finite evidence excludes a candidate; fewer than four + # eligible outputs stops the campaign before materialization and KD. + selected: + type: filter + input: online_eval + mode: top_k + metric: online_eval.lm_loss + direction: minimize + top_k: 4 + require_exact_count: true + + materialized: + type: materialize + input: selected + failure_policy: strict + + pre_kd_quality: + type: downstream_evaluation + input: materialized + failure_policy: strict + config: ${vlm_4b_campaign_evaluation} + + serving_smoke: + type: aiperf + input: pre_kd_quality + failure_policy: strict + config: + endpoint_type: chat + input_tokens: 32 + output_tokens: 8 + image_batch_sizes: [1] + image_width_mean: 224 + image_height_mean: 224 + concurrency: [1] + request_count: 1 + extra_inputs: {min_tokens: 8} + use_server_token_count: true + allow_aiperf_v011_online_tokenizer_resolution: true + seed: 42 + readiness_timeout: 600 + benchmark_timeout: 600 + 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 + server_context_overhead_tokens: 2048 + extra_vllm_args: + - -cc.cudagraph_mode=NONE + - --no-enable-flashinfer-autotune + - --gdn-prefill-backend + - triton + - --gpu-memory-utilization + - "0.7" + - --reasoning-parser + - qwen3 + - --default-chat-template-kwargs + - '{"enable_thinking": false}' + + kd_128: + type: global_kd + input: serving_smoke + model_source: materialized + trajectory: all_axis_recovery + failure_policy: strict + config: + resume: true + seed: 1111 + validation_seed: 445 + shuffle_training_data: true + freeze_policy: vision_frozen + activation_checkpointing: true + automodel: + parallel: + tp: 2 + cp: 1 + pp: 1 + ep: 1 + dp_shard: 1 + dp_replicate: 1 + sequence_parallel: false + pipeline_schedule: 1f1b + objective: *memory_safe_kd_objective + max_steps: 128 + global_batch_size: 4 + local_batch_size: 1 + checkpoint_every_steps: 128 + exposure: + cumulative_steps: 128 + global_batch_size: 4 + cumulative_examples: 512 + max_sample_length: 512 + estimated_cumulative_gpu_hours: 1.0 + + final_eval: + type: evaluation + input: kd_128 + failure_policy: strict + config: {eval_samples: 64, block_size: 512} + + quality_benchmarks: + type: downstream_evaluation + input: final_eval + failure_policy: strict + config: ${vlm_4b_campaign_evaluation} + + comparison_ready: + type: result_manifest + input: quality_benchmarks + config: + pre_kd_source: materialized + pre_kd_evaluation: pre_kd_quality + profile: ${vlm_4b_campaign_evaluation.profile} + row_manifest: profile:core-3_344-examples_r1-vllm + row_manifest_sha256: ${vlm_4b_campaign_evaluation.row_manifest_sha256} + reference_checkpoint: ${vlm_4b_campaign_evaluation.reference_checkpoint} + reference_cache_id: ${vlm_4b_campaign_evaluation.reference_cache_id} + milestones: [{steps: 128, kd: kd_128, evaluation: quality_benchmarks}] + + # Repetitions share one server lifetime per checkpoint. Quality ranks + # dominate the final aggregate; serving resolves equal quality ranks. + student_performance: + type: aiperf + input: comparison_ready + failure_policy: strict + config: &performance_config + endpoint_type: chat + input_tokens: 100 + output_tokens: 80 + image_batch_sizes: [1, 6, 12] + concurrency: [1, 4] + request_count: 64 + warmup_request_count: 32 + repetitions: 3 + seed: 42 + warmup_seed: 314 + image_width_mean: 1280 + image_height_mean: 720 + extra_inputs: {min_tokens: 80} + trust_remote_code: false + allow_aiperf_v011_online_tokenizer_resolution: true + use_server_token_count: true + collect_peak_gpu_memory: true + readiness_timeout: 600 + benchmark_timeout: 1800 + 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 + server_context_overhead_tokens: 32768 + extra_vllm_args: + - --gdn-prefill-backend + - triton + - --reasoning-parser + - qwen3 + - --default-chat-template-kwargs + - '{"enable_thinking": false}' + + best: + type: filter + input: student_performance + mode: aggregate_rank + metrics: + - metric: quality_benchmarks.modelopt_vlm_benchmark_realworldqa.exact_match_flexible-extract + direction: maximize + weight: 100 + - metric: quality_benchmarks.modelopt_vlm_benchmark_mmmu_val.mmmu_acc_none + direction: maximize + weight: 100 + - metric: student_performance.images_12.concurrency_4.image_throughput + direction: maximize + weight: 1 + top_k: 1 + + teacher-performance: + source: + run: teacher-reference + variants: teacher + objectives: all + nodes: + teacher_materialized: + type: materialize + failure_policy: strict + teacher_performance: + type: aiperf + input: teacher_materialized + failure_policy: strict + config: *performance_config + +zero_shot_evaluation: + enabled: false +aiperf: + enabled: false +global_distillation_sanity: + enabled: false +post_distillation_evaluation: + enabled: false diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml deleted file mode 100644 index edc2e04a717..00000000000 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml +++ /dev/null @@ -1,255 +0,0 @@ -# @package _global_ - -defaults: - - vlm_pruning_base - - _self_ - -# Qwen 3.5 4B VLM FFN-width 10%-to-20% KD search. Every MIP -# candidate receives the same serving, screening-KD, and quality contracts -# before aggregate ranking selects one student for longer KD and comparison -# with the teacher. -puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT} -display_name: Qwen3p5 4B VLM FFN-width 10%-to-20% KD search - -pruning: - experiment_id: qwen35_4b_vlm_ffn_campaign - intermediate_size_list: [7168, 6144, 5120] - eval_samples: 64 - seed: 42 - -sort_sanity: - eval_samples: 8 - -width_sanity: - target_count_per_axis: 2 - layer_count: 4 - eval_samples: 8 - -replacement_scoring: - eval_samples: 32 - -search_space: - axes: - ffn_intermediate: - enabled: true - teacher_value: 9216 - values: [7168, 6144, 5120] - -mip: - runs: - params-80: false - memory-85: false - ffn-candidates: - search_space: - depth: [0] - embedding: [2560] - axes_default: teacher - axes: {ffn.intermediate_size: all} - variants: - width-7168: - constraints: - params: {max: 92%} - search_space: - axes: {ffn.intermediate_size: [7168]} - width-6144: - constraints: - params: {max: 87%} - search_space: - axes: {ffn.intermediate_size: [6144]} - width-5120: - constraints: - params: {max: 82%} - search_space: - axes: {ffn.intermediate_size: [5120]} - solver: - backend: auto - num_solutions: 1 - min_hamming_distance: 2 - max_seconds_per_solution: 300 - homogeneous: {enabled: true, keep: 1, rank_by: objective} - -global_distillation: - enabled: false - domain: vlm - automodel: - activation_checkpointing: true - force_hf: false - teacher_force_hf: false - student_force_hf: false - # The maintained search prunes language FFNs only. Keep the unchanged vision - # tower fixed so KD memory is spent on the student parameters being adapted. - freeze_policy: vision_frozen - validation_enabled: false - objective: &memory_safe_kd_objective - main_ce: {weight: 1.0} - main_kd: {weight: 1.0, chunk_size: 64} - mtp_ce: {weight: 1.0} - mtp_kd: {weight: 1.0, chunk_size: 64} - -post_mip: - flows: - candidate-evaluation: - source: - run: ffn-candidates - variants: all - objectives: all - nodes: - online_eval: - type: evaluation - config: {eval_samples: 32, block_size: 512} - - materialized: - type: materialize - input: online_eval - - serving: - type: aiperf - input: materialized - config: - endpoint_type: chat - input_tokens: 100 - output_tokens: 80 - image_batch_sizes: [1, 4, 8] - image_width_mean: 1280 - image_height_mean: 720 - concurrency: [1, 4] - request_count: 8 - extra_inputs: {min_tokens: 80} - use_server_token_count: true - allow_aiperf_v011_online_tokenizer_resolution: true - seed: 42 - readiness_timeout: 600 - benchmark_timeout: 1800 - 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 - server_context_overhead_tokens: 16384 - extra_vllm_args: - - -cc.cudagraph_mode=NONE - - --no-enable-flashinfer-autotune - - --gpu-memory-utilization - - "0.7" - - --reasoning-parser - - qwen3 - - --default-chat-template-kwargs - - '{"enable_thinking": false}' - - screening_kd: - type: global_kd - input: serving - config: &screening_kd_config - seed: 1111 - validation_seed: 445 - shuffle_training_data: true - freeze_policy: vision_frozen - activation_checkpointing: true - # One complete teacher/student pair exceeds an 80 GB GPU on late - # multimodal samples. TP2 preserves the batch and sample contract - # while sharding supported language FFN and full-attention state; - # vision and GDN modules remain replicated. - automodel: - parallel: - tp: 2 - cp: 1 - pp: 1 - ep: 1 - dp_shard: 1 - dp_replicate: 1 - sequence_parallel: false - pipeline_schedule: 1f1b - objective: *memory_safe_kd_objective - max_steps: 64 - global_batch_size: 1 - local_batch_size: 1 - # Avoid a mid-run consolidated export, which can leave too little - # device memory for the next high-visual-token sample. - checkpoint_every_steps: 64 - - screening_eval: - type: evaluation - input: screening_kd - config: {eval_samples: 32, block_size: 512} - - quality_screen: - type: downstream_evaluation - input: screening_eval - failure_policy: strict - config: &bounded_vlm_evaluation - profile: qwen35_vlm_realworldqa_mmmu_prefix100_x2 - disable_thinking: true - batch_size: 1 - timeout_seconds: 14400 - dtype: bfloat16 - gpu_memory_utilization: 0.7 - max_model_len: 16384 - limit_mm_per_prompt: {image: 12} - 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 - - selected: - type: filter - input: quality_screen - mode: aggregate_rank - metrics: - - metric: screening_eval.lm_loss - direction: minimize - - metric: quality_screen.modelopt_vlm_benchmark_realworldqa.exact_match_flexible-extract - direction: maximize - - metric: quality_screen.modelopt_vlm_benchmark_mmmu_val.mmmu_acc_none - direction: maximize - top_k: 1 - - global_kd: - type: global_kd - input: selected - model_source: materialized - config: - <<: *screening_kd_config - # Exploratory fresh 256-step run from the selected materialized student. - # Screening KD weights are used for ranking only, not initialization. - # This is not a convergence threshold or a campaign completion gate. - max_steps: 256 - checkpoint_every_steps: 256 - - final_eval: - type: evaluation - input: global_kd - config: {eval_samples: 64, block_size: 512} - - quality_benchmarks: - type: downstream_evaluation - input: final_eval - failure_policy: strict - config: - <<: *bounded_vlm_evaluation - reference_checkpoint: ${teacher_dir} - - best: - type: filter - input: quality_benchmarks - mode: top_k - metric: final_eval.lm_loss - direction: minimize - top_k: 1 - -zero_shot_evaluation: - enabled: false -aiperf: - enabled: false -global_distillation_sanity: - enabled: false -post_distillation_evaluation: - enabled: false diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml index 83da9c63edd..a8d14b5c1fb 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml @@ -1,12 +1,16 @@ # @package _global_ defaults: - - vlm_pruning_base + - /families/qwen3_5/qwen3p5_4b/all_axis_base@_global_ + - /families/qwen3_5/qwen3p5_4b/vlm_quality_evaluation@_global_ - _self_ -# Bounded lifecycle check of FFN pruning through physical materialization, -# checkpoint reload, bounded RealWorldQA evaluation, and two-step VLM KD. +# Complete bounded all-axis lifecycle through physical materialization, +# checkpoint reload, frozen image-text evaluation, multimodal serving, and +# two-step VLM KD. puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT} +display_name: Qwen3p5 4B VLM all-axis lifecycle smoke + data: # Tested public snapshot; callers may override with another immutable SHA. revision: ${oc.env:PUZZLETRON_DATASET_REVISION,51f4f4d219315c3283950994d4eb3d7fc30aa87b} @@ -26,6 +30,19 @@ global_distillation: mtp_ce: {weight: 1.0} mtp_kd: {weight: 1.0, chunk_size: 64} +mip: + runs: + params-80: + constraints: + params: {max: 92%} + solver: + backend: pulp + num_solutions: 1 + min_hamming_distance: 1 + max_seconds_per_solution: 30 + homogeneous: {enabled: false} + memory-85: false + post_mip: flows: params-80: @@ -47,18 +64,31 @@ post_mip: materialized: type: materialize input: best_vlm_loss + failure_policy: strict checkpoint_eval: type: downstream_evaluation input: materialized failure_policy: strict - config: &bounded_realworldqa - profile: qwen35_vlm_realworldqa - batch_size: 1 - timeout_seconds: 900 - dtype: bfloat16 - gpu_memory_utilization: 0.8 - max_model_len: 4096 - limit_mm_per_prompt: {image: 1} + config: ${vlm_4b_smoke_evaluation} + serving_smoke: + type: aiperf + input: checkpoint_eval + failure_policy: strict + config: &serving_smoke_config + endpoint_type: chat + input_tokens: 32 + output_tokens: 8 + image_batch_sizes: [1] + image_width_mean: 224 + image_height_mean: 224 + concurrency: [1] + request_count: 1 + extra_inputs: {min_tokens: 8} + use_server_token_count: true + allow_aiperf_v011_online_tokenizer_resolution: true + seed: 42 + readiness_timeout: 600 + benchmark_timeout: 600 topology: tensor_parallel_size: 1 pipeline_parallel_size: 1 @@ -68,15 +98,25 @@ post_mip: enable_expert_parallel: false distributed_executor_backend: mp gpu_group_size: 1 + server_context_overhead_tokens: 2048 + extra_vllm_args: + - -cc.cudagraph_mode=NONE + - --no-enable-flashinfer-autotune + - --gdn-prefill-backend + - triton + - --gpu-memory-utilization + - "0.7" + - --reasoning-parser + - qwen3 + - --default-chat-template-kwargs + - '{"enable_thinking": false}' short_vlm_kd: type: global_kd - input: checkpoint_eval + input: serving_smoke + failure_policy: strict config: freeze_policy: vision_frozen activation_checkpointing: true - # A complete 4B teacher/student pair exceeds one 80 GB GPU on - # multimodal samples. This matches the campaign's validated TP2 - # memory-safe KD recipe. automodel: parallel: tp: 2 @@ -92,14 +132,38 @@ post_mip: global_batch_size: 1 local_batch_size: 1 checkpoint_every_steps: 2 + exposure: + cumulative_steps: 2 + global_batch_size: 1 + cumulative_examples: 2 + max_sample_length: 512 + estimated_cumulative_gpu_hours: 0.2 post_kd_checkpoint_eval: type: downstream_evaluation input: short_vlm_kd failure_policy: strict - config: *bounded_realworldqa + config: ${vlm_4b_smoke_evaluation} + result: + type: result_manifest + input: post_kd_checkpoint_eval + config: + pre_kd_source: materialized + pre_kd_evaluation: checkpoint_eval + profile: ${vlm_4b_smoke_evaluation.profile} + row_manifest: profile:core-3_24-examples_r1-vllm + row_manifest_sha256: ${vlm_4b_smoke_evaluation.row_manifest_sha256} + reference_checkpoint: ${vlm_4b_smoke_evaluation.reference_checkpoint} + reference_cache_id: ${vlm_4b_smoke_evaluation.reference_cache_id} + milestones: + - {steps: 2, kd: short_vlm_kd, evaluation: post_kd_checkpoint_eval} + final_serving_smoke: + type: aiperf + input: result + failure_policy: strict + config: *serving_smoke_config final_image_eval: type: evaluation - input: post_kd_checkpoint_eval + input: final_serving_smoke config: {eval_samples: 2, block_size: 512} best: type: filter diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/vlm_quality_evaluation.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/vlm_quality_evaluation.yaml new file mode 100644 index 00000000000..421927a27eb --- /dev/null +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/vlm_quality_evaluation.yaml @@ -0,0 +1,59 @@ +# @package _global_ + +# Shared 4B VLM evaluator identities. The resolver replaces the evaluator +# revision with the immutable worker-source identity before sealing a run. +vlm_4b_quality_evaluation: + profile: qwen35_vlm_realworldqa64_mmmu120_mvbench160_frozen_rows_v3 + evaluator_revision: ${oc.env:PUZZLETRON_SOURCE_REVISION,unpublished} + row_manifest_sha256: be7ed6f695fca19b11a90ecd4f400a95162ede7d421573a08d81b6b70be03536 + reference_checkpoint: ${teacher_dir} + reference_once: true + reference_cache_id: qwen3p5-4b-core-3_344-examples-r1-vllm-teacher + batch_size: 1 + timeout_seconds: 14400 + dtype: bfloat16 + gdn_prefill_backend: triton + gpu_memory_utilization: 0.7 + max_model_len: 32768 + limit_mm_per_prompt: {image: 32} + topology: &single_gpu_vllm + 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_4b_campaign_evaluation: + profile: ${vlm_4b_quality_evaluation.profile} + evaluator_revision: ${vlm_4b_quality_evaluation.evaluator_revision} + row_manifest_sha256: ${vlm_4b_quality_evaluation.row_manifest_sha256} + reference_checkpoint: ${vlm_4b_quality_evaluation.reference_checkpoint} + reference_once: ${vlm_4b_quality_evaluation.reference_once} + reference_cache_id: ${vlm_4b_quality_evaluation.reference_cache_id} + batch_size: ${vlm_4b_quality_evaluation.batch_size} + timeout_seconds: ${vlm_4b_quality_evaluation.timeout_seconds} + dtype: ${vlm_4b_quality_evaluation.dtype} + gdn_prefill_backend: ${vlm_4b_quality_evaluation.gdn_prefill_backend} + gpu_memory_utilization: ${vlm_4b_quality_evaluation.gpu_memory_utilization} + max_model_len: ${vlm_4b_quality_evaluation.max_model_len} + limit_mm_per_prompt: ${vlm_4b_quality_evaluation.limit_mm_per_prompt} + topology: ${vlm_4b_quality_evaluation.topology} + +vlm_4b_smoke_evaluation: + profile: qwen35_vlm_core3_24row_smoke_v2 + evaluator_revision: ${oc.env:PUZZLETRON_SOURCE_REVISION,unpublished} + row_manifest_sha256: 08eb41371b2f40fdde31b45573a7276fa370ec96be7c1d95554c916d872d6cde + reference_checkpoint: ${teacher_dir} + reference_once: true + reference_cache_id: qwen3p5-4b-core-3_24-examples-r1-vllm-teacher + batch_size: 1 + timeout_seconds: 3000 + dtype: bfloat16 + gdn_prefill_backend: triton + gpu_memory_utilization: 0.7 + max_model_len: 16384 + limit_mm_per_prompt: {image: 32} + topology: *single_gpu_vllm diff --git a/examples/puzzletron/docs/maintained_recipes.md b/examples/puzzletron/docs/maintained_recipes.md index 7ecd7a5a58a..2b2cbffa0c0 100644 --- a/examples/puzzletron/docs/maintained_recipes.md +++ b/examples/puzzletron/docs/maintained_recipes.md @@ -9,8 +9,8 @@ model and modality before running a campaign. | `qwen3p5_0p8b_text_smoke.yaml` | Text pruning lifecycle with bounded evaluation, serving, and distillation | `single-gpu` | | `qwen3p5_0p8b_vlm_smoke.yaml` | Image-text pruning lifecycle with bounded evaluation, serving, and distillation | `smoke` | | `qwen3p5_0p8b_vlm_campaign.yaml` | Larger multi-axis image-text integration example | `campaign` | -| `qwen3p5_4b_vlm_smoke.yaml` | 4B FFN-pruning lifecycle with bounded evaluation and distillation | `smoke` | -| `qwen3p5_4b_vlm_campaign.yaml` | 4B FFN-width search and screening example | `campaign` | +| `qwen3p5_4b_vlm_smoke.yaml` | 4B all-axis lifecycle with bounded evaluation, serving, and distillation | `smoke` | +| `qwen3p5_4b_vlm_campaign.yaml` | 4B exact-four all-axis search with matched KD128 | `campaign` | Smoke workloads validate that the configured lifecycle runs, produces usable checkpoints, and resumes. Their scores and throughput are not model-quality or @@ -76,19 +76,23 @@ python examples/puzzletron/materialize_dataset.py nemotron_vlm_v2 \ --max-shards-per-subset 1 ``` -Set `data.path` to `/shared/data/qwen3p5-vlm`. The campaign needs at least 64 +Set `data.path` to `/shared/data/qwen3p5-vlm`. The campaign needs at least 512 samples, so prepare a separate directory with the same command and -`--num-samples 64`, then point the campaign recipe to it. - -The smoke searches aligned FFN widths and checks materialization, checkpoint -reload, bounded RealWorldQA evaluation, two TP2 distillation steps, final -image-text evaluation, reporting, and resume on two colocated GPUs. - -The campaign compares roughly 10%, 15%, and 20% FFN-pruning bands using matched -image-text loss, serving, screening distillation, and evaluation. Its ranking -policy and training budgets are examples. Compare candidates only when their -model, data, evaluator, sampling, generation, teacher, and runtime identities -match. +`--num-samples 512`, then point the campaign recipe to it. + +The smoke exercises hidden width, FFN width, grouped attention, GDN geometry, +and depth before checking materialization, checkpoint reload, a frozen 24-row +image-text evaluation, serving, two TP2 distillation steps, structured results, +reporting, and resume on two colocated GPUs. + +The campaign generates candidates from parameter and serving-memory budgets +over every supported structural axis. It ranks complete finite pre-KD +image-text loss, fails before fanout unless exactly four distinct architectures +remain, and gives all four the same frozen pre/post evaluation and resumable +128-step, 512-example TP2 distillation exposure. Final quality ranks dominate +the serving tie-breaker. These search, ranking, and training budgets are example +settings. Compare candidates only when their model, data, evaluator, sampling, +generation, teacher, and runtime identities match. ## Inspect results diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index 01019d668c7..b633ae0d496 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -225,6 +225,8 @@ sample counts and parser outcomes remain attached as result evidence. ## Filters `top_k` accepts one integer or separate homogeneous/heterogeneous quotas. +With a scalar quota, set `require_exact_count: true` to stop before downstream +fanout when missing or non-finite evidence leaves fewer than `top_k` candidates. `threshold` accepts `min`, `max`, or both. `pareto` retains the nondominated set. `aggregate_rank` computes a weighted mean rank and retains `top_k`: diff --git a/modelopt/torch/puzzletron/orchestration/_route_catalog.py b/modelopt/torch/puzzletron/orchestration/_route_catalog.py index 4de6e3dbf3c..e7a175b0603 100644 --- a/modelopt/torch/puzzletron/orchestration/_route_catalog.py +++ b/modelopt/torch/puzzletron/orchestration/_route_catalog.py @@ -64,13 +64,16 @@ def route_id(self) -> str: } _FOUR_B_CAMPAIGN_STAGES: dict[str, dict[str, Any]] = { - "replacement_scoring": {"strategy": "single"}, - "post.candidate-evaluation.online_eval": {"instances": 4}, + "depth_importance": {"strategy": "single"}, + "replacement_scoring": {"strategy": "persistent_pool", "instances": 8}, + "post.candidate-evaluation.online_eval": {"instances": 8}, "post.candidate-evaluation.materialized": {"instances": 4}, - "post.candidate-evaluation.serving": {"instances": 4}, - "post.candidate-evaluation.screening_kd": {"instances": 4}, - "post.candidate-evaluation.screening_eval": {"instances": 4}, - "post.candidate-evaluation.quality_screen": {"instances": 4}, + "post.candidate-evaluation.pre_kd_quality": {"instances": 4}, + "post.candidate-evaluation.serving_smoke": {"instances": 4}, + "post.candidate-evaluation.kd_128": {"instances": 4}, + "post.candidate-evaluation.final_eval": {"instances": 4}, + "post.candidate-evaluation.quality_benchmarks": {"instances": 4}, + "post.candidate-evaluation.student_performance": {"instances": 4}, } ROUTES = ( @@ -109,21 +112,25 @@ def route_id(self) -> str: model="qwen3.5-4b", workflow="vlm-pruning", mode="smoke", - search="bounded-ffn", + search="bounded-all-axis", evaluation="smoke", distillation="smoke", experiment_template="families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml", - execution_stages={"replacement_scoring": {"strategy": "single"}}, + execution_stages={ + "depth_importance": {"strategy": "single"}, + "replacement_scoring": {"strategy": "single"}, + "post.params-80.image_eval": {"instances": 2}, + }, requires_data=True, ), RouteProfile( model="qwen3.5-4b", workflow="vlm-pruning", mode="campaign", - search="ffn-10-to-20pct", + search="multi-axis-exact-four", evaluation="quality", - distillation="screening", - experiment_template=("families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml"), + distillation="matched-kd128", + experiment_template="families/qwen3_5/qwen3p5_4b/runs/all_axis_kd_search.yaml", execution_stages=_FOUR_B_CAMPAIGN_STAGES, requires_data=True, ), diff --git a/modelopt/torch/puzzletron/post_mip/filters.py b/modelopt/torch/puzzletron/post_mip/filters.py index 8c78422b8de..7ad28ab0608 100644 --- a/modelopt/torch/puzzletron/post_mip/filters.py +++ b/modelopt/torch/puzzletron/post_mip/filters.py @@ -57,7 +57,8 @@ def validate_filter_config(config: Mapping[str, Any]) -> None: mode = str(config.get("mode") or "") common = {"type", "input", "model_source", "failure_policy", "config", "mode"} allowed = { - "top_k": common | {"metric", "direction", "top_k", "best_selection_mode"}, + "top_k": common + | {"metric", "direction", "top_k", "best_selection_mode", "require_exact_count"}, "threshold": common | {"metric", "min", "max"}, "pareto": common | {"metrics"}, "aggregate_rank": common | {"metrics", "top_k"}, @@ -70,6 +71,9 @@ def validate_filter_config(config: Mapping[str, Any]) -> None: if mode in {"top_k", "threshold"} and not config.get("metric"): raise ValueError(f"{mode} filter requires metric") if mode == "top_k": + require_exact_count = config.get("require_exact_count", False) + if not isinstance(require_exact_count, bool): + raise TypeError("top_k.require_exact_count must be a boolean") direction = str(config.get("direction", "minimize")) if direction not in {"minimize", "maximize"}: raise ValueError("top_k direction must be minimize or maximize") @@ -89,6 +93,10 @@ def validate_filter_config(config: Mapping[str, Any]) -> None: top_k = config.get("top_k") if best_selection_mode is not None and isinstance(top_k, Mapping): raise ValueError("best_selection_mode requires an integer top_k") + if require_exact_count and (best_selection_mode is not None or isinstance(top_k, Mapping)): + raise ValueError( + "require_exact_count requires a scalar top_k without best_selection_mode" + ) if isinstance(top_k, Mapping): if set(top_k) - {"homogeneous", "heterogeneous"}: raise ValueError("top_k quotas accept homogeneous and heterogeneous only") @@ -128,6 +136,14 @@ def _origin_kind(ledger: CandidateLedger, revision_id: str) -> str: return str(revision.artifact.get("kind", "heterogeneous")) +def _require_exact_top_k(selected: Sequence[str], config: Mapping[str, Any]) -> None: + if config.get("require_exact_count") and len(selected) != int(config["top_k"]): + raise RuntimeError( + f"top_k requires exactly {int(config['top_k'])} selected revisions; " + f"found {len(selected)} with complete finite metric evidence" + ) + + def _ordered_metric_rows( rows: Sequence[tuple[float, str]], *, @@ -176,6 +192,7 @@ def _apply_sweep_top_k( selected = tuple(revision_id for _value, revision_id in rows[:top_k]) for _value, revision_id in rows[top_k:]: excluded[revision_id] = "outside top_k" + _require_exact_top_k(selected, config) return selected, excluded, scores selected_ids = set() @@ -238,6 +255,7 @@ def apply_filter( for _value, revision_id in rows: if revision_id not in selected: excluded[revision_id] = "outside top_k" + _require_exact_top_k(selected, config) return selected, excluded, scores if mode == "threshold": diff --git a/tests/unit/torch/puzzletron/test_post_mip_filters.py b/tests/unit/torch/puzzletron/test_post_mip_filters.py index b9f61cee27c..f86b63654af 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_filters.py +++ b/tests/unit/torch/puzzletron/test_post_mip_filters.py @@ -224,6 +224,43 @@ def test_top_k_rejects_unknown_best_selection_mode(best_selection_mode): ) +def test_top_k_exact_count_fails_before_fanout_when_finite_candidates_are_missing(tmp_path): + ledger = _ledger( + tmp_path, + { + "revision-a": {"loss": 1.0}, + "revision-b": {"loss": math.nan}, + "revision-c": {"loss": 3.0}, + }, + ) + + with pytest.raises(RuntimeError, match="requires exactly 3 selected revisions; found 2"): + apply_filter( + ledger, + tuple(ledger.revisions), + { + "mode": "top_k", + "metric": "serving.loss", + "direction": "minimize", + "top_k": 3, + "require_exact_count": True, + }, + ) + + +@pytest.mark.parametrize("value", [1, "true", None]) +def test_top_k_exact_count_requires_a_boolean(value): + with pytest.raises(TypeError, match="require_exact_count"): + validate_filter_config( + { + "mode": "top_k", + "metric": "serving.loss", + "top_k": 3, + "require_exact_count": value, + } + ) + + def test_best_selection_mode_is_rejected_on_non_top_k_filter(): with pytest.raises(ValueError, match="best_selection_mode"): validate_filter_config( diff --git a/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py b/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py index d4000d47707..eec6c648451 100644 --- a/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py +++ b/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py @@ -20,6 +20,7 @@ import yaml +from modelopt.torch.puzzletron.mip.profiles import normalize_mip_profiles from puzzletron_orchestrator.compiler import ( compile_campaign_plan, load_execution_config, @@ -41,11 +42,23 @@ REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml" ) - -def _compile_plan( - tmp_path: Path, - recipe_source: Path, -): +ALL_AXIS_DOMAINS = { + "hidden_width": {"enabled": True, "teacher_value": 2560, "values": [2400]}, + "kv_groups": {"enabled": True, "teacher_value": 4, "values": [2]}, + "q_heads_per_group": {"enabled": True, "teacher_value": 4, "values": [3]}, + "ffn_intermediate": {"enabled": True, "teacher_value": 9216, "values": [8704]}, + "gdn_key_groups": {"enabled": True, "teacher_value": 16, "values": [14]}, + "gdn_value_heads_per_group": { + "enabled": True, + "teacher_value": 2, + "values": [1], + }, + "gdn_key_head_dim": {"enabled": True, "teacher_value": 128, "values": [112]}, + "gdn_value_head_dim": {"enabled": True, "teacher_value": 128, "values": [112]}, +} + + +def _compile_plan(tmp_path: Path, recipe_source: Path): run_root = tmp_path / recipe_source.stem dataset = tmp_path / "dataset" recipe = yaml.safe_load(recipe_source.read_text()) @@ -88,80 +101,137 @@ def test_qwen3p5_4b_model_pins_the_bounded_ffn_grid() -> None: assert model["pruning"] == {"intermediate_size_list": widths} -def test_qwen3p5_4b_smoke_materializes_reloads_and_bounds_kd_and_evaluation( - tmp_path: Path, -) -> None: - plan = _compile_plan( - tmp_path, - SMOKE_RECIPE_PATH, - ) +def test_qwen3p5_4b_smoke_covers_all_axes_and_emits_comparable_results(tmp_path: Path) -> None: + plan = _compile_plan(tmp_path, SMOKE_RECIPE_PATH) config = plan.experiment_config - post_stages = tuple(stage for stage in plan.stages if stage.stage_id.startswith("post.")) nodes = config["post_mip"]["flows"]["params-80"]["nodes"] + post_stages = tuple(stage for stage in plan.stages if stage.stage_id.startswith("post.")) + assert config["search_space"]["axes"] == ALL_AXIS_DOMAINS + assert config["width_sanity"]["axes"] == list(ALL_AXIS_DOMAINS) + assert config["depth_importance"]["max_removals"] == 1 + profiles = normalize_mip_profiles( + config["mip"], available_depths=[0, 1], available_embeddings=[2560, 2400] + ) + assert len(profiles) == 1 + assert ( + sum( + len(profile.embedding_widths) + * len(profile.depth_selections) + * profile.solver.num_solutions + for profile in profiles + ) + == 4 + ) assert tuple(stage.stage_id for stage in post_stages) == ( "post.params-80.image_eval", "post.params-80.best_vlm_loss", "post.params-80.materialized", "post.params-80.checkpoint_eval", + "post.params-80.serving_smoke", "post.params-80.short_vlm_kd", "post.params-80.post_kd_checkpoint_eval", + "post.params-80.result", + "post.params-80.final_serving_smoke", "post.params-80.final_image_eval", "post.params-80.best", ) - assert post_stages[0].parents == ("mip",) - for parent, stage in pairwise(post_stages): + for parent, stage in pairwise(post_stages[:7]): assert stage.parents == (parent.stage_id,) - assert nodes["materialized"]["input"] == "best_vlm_loss" - assert nodes["checkpoint_eval"]["input"] == "materialized" - assert nodes["checkpoint_eval"]["failure_policy"] == "strict" - kd = nodes["short_vlm_kd"] - assert kd["input"] == "checkpoint_eval" - assert kd["config"]["automodel"]["parallel"]["tp"] == 2 - assert ( - next(stage for stage in plan.stages if stage.stage_id.endswith("short_vlm_kd")).total_gpus - == 2 - ) + result_stage = next(stage for stage in post_stages if stage.stage_id.endswith(".result")) + assert set(result_stage.parents) == { + "post.params-80.materialized", + "post.params-80.checkpoint_eval", + "post.params-80.short_vlm_kd", + "post.params-80.post_kd_checkpoint_eval", + } + assert nodes["checkpoint_eval"]["config"]["profile"] == "qwen35_vlm_core3_24row_smoke_v2" + assert nodes["checkpoint_eval"]["config"] == nodes["post_kd_checkpoint_eval"]["config"] + assert nodes["serving_smoke"]["config"] == nodes["final_serving_smoke"]["config"] + assert nodes["short_vlm_kd"]["config"]["max_steps"] == 2 + assert nodes["short_vlm_kd"]["config"]["automodel"]["parallel"]["tp"] == 2 + assert nodes["result"]["config"]["milestones"] == [ + {"steps": 2, "kd": "short_vlm_kd", "evaluation": "post_kd_checkpoint_eval"} + ] + stages = {stage.stage_id: stage for stage in plan.stages} + assert stages["post.params-80.image_eval"].instances == 2 + assert stages["post.params-80.short_vlm_kd"].total_gpus == 2 + assert plan.final_report_partition == plan.runner.slurm.partition_cpu -def test_qwen3p5_4b_campaign_compares_pruning_bands_and_teacher(tmp_path) -> None: - plan = _compile_plan( - tmp_path, - CAMPAIGN_RECIPE_PATH, - ) +def test_qwen3p5_4b_campaign_runs_exactly_four_candidates_through_matched_kd128( + tmp_path: Path, +) -> None: + plan = _compile_plan(tmp_path, CAMPAIGN_RECIPE_PATH) config = plan.experiment_config - candidates = config["mip"]["runs"]["ffn-candidates"] nodes = config["post_mip"]["flows"]["candidate-evaluation"]["nodes"] - assert tuple(config["mip"]["runs"]) == ("params-80", "memory-85", "ffn-candidates") - assert config["mip"]["runs"]["params-80"] is False - assert config["mip"]["runs"]["memory-85"] is False - assert set(candidates["variants"]) == {"width-7168", "width-6144", "width-5120"} - assert nodes["quality_benchmarks"]["config"]["reference_checkpoint"] == config["teacher_dir"] - assert nodes["global_kd"]["model_source"] == "materialized" - assert tuple(stage.stage_id for stage in plan.stages)[-11:-1] == ( - "post.candidate-evaluation.online_eval", - "post.candidate-evaluation.materialized", - "post.candidate-evaluation.serving", - "post.candidate-evaluation.screening_kd", - "post.candidate-evaluation.screening_eval", - "post.candidate-evaluation.quality_screen", - "post.candidate-evaluation.selected", - "post.candidate-evaluation.global_kd", - "post.candidate-evaluation.final_eval", - "post.candidate-evaluation.quality_benchmarks", + assert config["embedding_pruning"]["widths"] == [2560, 2400, 2240] + assert set(config["search_space"]["axes"]) == set(ALL_AXIS_DOMAINS) + assert config["depth_importance"]["max_removals"] == 2 + profiles = normalize_mip_profiles( + config["mip"], available_depths=[0, 1, 2], available_embeddings=[2560, 2400, 2240] ) - assert tuple(stage.stage_id for stage in plan.stages)[-1] == "post.candidate-evaluation.best" + search_profiles = [profile for profile in profiles if profile.run_id == "search-candidates"] + assert {profile.variant_id for profile in search_profiles} == {"params-82", "memory-85"} + assert ( + sum( + len(profile.embedding_widths) + * len(profile.depth_selections) + * profile.solver.num_solutions + for profile in search_profiles + ) + == 18 + ) + assert nodes["selected"] == { + "type": "filter", + "input": "online_eval", + "mode": "top_k", + "metric": "online_eval.lm_loss", + "direction": "minimize", + "top_k": 4, + "require_exact_count": True, + } + kd = nodes["kd_128"] + assert kd["input"] == "serving_smoke" + assert kd["model_source"] == "materialized" + assert kd["config"]["resume"] is True + assert kd["config"]["max_steps"] == 128 + assert kd["config"]["global_batch_size"] == 4 + assert kd["config"]["checkpoint_every_steps"] == 128 + assert kd["exposure"]["cumulative_examples"] == 512 + assert nodes["pre_kd_quality"]["config"] == nodes["quality_benchmarks"]["config"] + assert nodes["pre_kd_quality"]["config"]["profile"] == ( + "qwen35_vlm_realworldqa64_mmmu120_mvbench160_frozen_rows_v3" + ) + assert nodes["comparison_ready"]["config"]["milestones"] == [ + {"steps": 128, "kd": "kd_128", "evaluation": "quality_benchmarks"} + ] + assert [entry["weight"] for entry in nodes["best"]["metrics"]] == [100, 100, 1] stages = {stage.stage_id: stage for stage in plan.stages} - candidate_stages = { - "post.candidate-evaluation.online_eval", + assert stages["replacement_scoring"].instances == 8 + assert stages["post.candidate-evaluation.online_eval"].instances == 8 + four_candidate_stages = { "post.candidate-evaluation.materialized", - "post.candidate-evaluation.serving", - "post.candidate-evaluation.screening_kd", - "post.candidate-evaluation.screening_eval", - "post.candidate-evaluation.quality_screen", + "post.candidate-evaluation.pre_kd_quality", + "post.candidate-evaluation.serving_smoke", + "post.candidate-evaluation.kd_128", + "post.candidate-evaluation.final_eval", + "post.candidate-evaluation.quality_benchmarks", + "post.candidate-evaluation.student_performance", } - assert all(stages[stage_id].instances == 4 for stage_id in candidate_stages) - assert stages["post.candidate-evaluation.screening_kd"].total_gpus == 8 - assert all(stages[stage_id].gpus_per_node == 8 for stage_id in candidate_stages) - assert stages["post.candidate-evaluation.global_kd"].total_gpus == 2 + assert all(stages[stage_id].instances == 4 for stage_id in four_candidate_stages) + assert stages["post.candidate-evaluation.kd_128"].total_gpus == 8 + assert all( + stages[stage_id].total_gpus == 4 + for stage_id in four_candidate_stages - {"post.candidate-evaluation.kd_128"} + ) + assert plan.runner.slurm.max_nodes == 1 + assert all(stages[stage_id].nodes <= 1 for stage_id in four_candidate_stages) + performance = nodes["student_performance"]["config"] + assert performance["repetitions"] == 3 + assert performance["image_batch_sizes"] == [1, 6, 12] + assert performance["concurrency"] == [1, 4] + teacher_nodes = config["post_mip"]["flows"]["teacher-performance"]["nodes"] + assert teacher_nodes["teacher_performance"]["config"] == performance + assert plan.final_report_partition == plan.runner.slurm.partition_cpu