[AMD][CI] CI: fix AMD 2-GPU multimodal-gen partition-count abort

[michaelzhang-ai]

Motivation

Every multimodal-gen-test-2-gpu-amd shard on both pr-test-amd and pr-test-amd-rocm720 fails before running a single test (e.g. run 31327057626):

Error: total_partitions (3) must be >= standalone files (7)

run_suite.py reserved one whole partition per standalone file whenever CI did not pass a precomputed partition plan, and bailed out when a suite had more standalone files than --total-partitions. The AMD lanes hardcode --total-partitions 3, so the guard tripped as soon as STANDALONE_FILES["2-gpu"] grew past three entries (#33725 and #33775 took it to seven).

Raising the hardcoded count (as the previously closed #33879 did) fixes it only until the next standalone file is added, and each extra AMD partition is a serial 2-GPU job with its own container and dependency install. This removes the coupling instead: the shard count no longer has to track the standalone file count.

Worth noting the same guard had already silently cost the AMD lane its parametrized coverage — at exactly three standalone files, parametrized_partitions was 0 and all three shards ran only standalone files. Balanced packing restores those cases.

Modifications

• partitioning.py: add assign_partition(), returning one shard's slice of the deterministic LPT pass.
• run_suite.py: when no --partition-plan-json is given, build the shard assignment with build_local_partition_assignment() — the same LPT pass over parametrized cases and standalone files that compute_diffusion_partitions.py uses for the precomputed plans — instead of the index-based scheme, and drop the guard. Both plan sources now feed one executor (_run_partition_assignment), so the two code paths cannot drift; auto_partition() and the already-dead _get_standalone_file() go away.
• gen_diffusion_ci_outputs.py: import the partition helpers from partitioning directly, since run_suite no longer re-exports partition_items_by_lpt.
• New unit/test_suite_partitioning.py: every shard count schedules each case and standalone file exactly once, including counts below the standalone file count.
• Refresh the now-stale part: [0, 1, 2] comments in the two AMD workflows.

No workflow partition counts change. For the 2-GPU suite the three AMD shards come out at 2642s / 2597s / 2652s estimated, versus ~2506s per shard for the two parametrized shards before the breakage, so per-shard runtime is unchanged. Four of the seven standalone files (test_ipc_a2a_2_gpu, test_pynccl_a2a_capture_2_gpu, test_dp_serving_2_gpu, test_usp_replicated_parity_2_gpu) self-skip on non-CUDA platforms, so the ROCm shards only newly execute test_ar_models and test_diffusion_bcg_tp2_zimage_turbo.

Separate pre-existing bug found while investigating

scripts/ci/utils/diffusion/diffusion_case_parser.py AST-parses STANDALONE_FILES out of run_suite.py, but that dict moved to server/gpu_cases.py in #24630. The parser therefore reports zero standalone files, so the CUDA partition plans have not scheduled any standalone test since June, and verify_diffusion_coverage.py does not notice because it derives its expectations from the same parser. Fixing it is deliberately not part of this PR: it would newly schedule nine never-run standalone files on the CUDA lanes and would immediately fail the coverage job, which exits non-zero for standalone files missing a STANDALONE_FILE_EST_TIMES entry (test_generate_zimage_turbo_cli.py has none).

Which lanes change behavior

Only the lanes that run without a precomputed plan, which is AMD plus two small CUDA jobs:

• AMD 2-GPU (pr-test-amd, pr-test-amd-rocm720): aborted before any test; now 3 balanced shards of cases + standalone files.
• AMD 1-GPU (both lanes): composition changes too — 3 parametrized shards plus 1 dedicated standalone shard become 4 mixed shards. Better balanced, and test_generate_zimage_turbo_cli.py (the only standalone file left on HIP, since is_hip drops test_update_weights_from_disk) now shares a shard. It keeps running even when cases in that shard fail; see the fast-fail note below.
• CUDA multimodal-gen-test-1-5090: no standalone entries for the suite, one partition, so the same LPT call over the same items — identical case list and order.
• CUDA bcg-diffusion: empty PARAMETRIZED_CASE_GROUPS, one standalone file, one partition — the same file runs; only the report's is_standalone flag differs, and nothing consumes that job's report.
• Everything else is untouched: the plan-driven CUDA and NPU shards execute byte-identical code (NPU's STANDALONE_FILES is empty anyway), 1-gpu-b200 and unit take the file-suite branch, component-accuracy has its own branch, and MUSA has a separate run_suite.py.

A shard no longer stops before its standalone files when a parametrized case fails without --continue-on-error. Standalone files are separate pytest runs that only share a shard because the shard count is fixed; keeping the early return would have made the AMD 1-GPU standalone file collateral damage of unrelated case failures, silently, since the AMD lanes run no coverage check.

Validation

• AMD, targeted dispatch of the two affected stages on this branch: run 31445683539
  • multimodal-gen-test-2-gpu-amd shard 0: no longer aborts — it received an assignment (2-gpu 1/3, 7 parametrized cases) and ran for 22 min, 5 passed / 2 errors. Both errors are a broken LTX-2.3 checkpoint in the runner's model cache (transformer/ has no weight file), unrelated to partitioning. Shards 1 and 2 are still queued behind the serial 2-GPU matrix.
  • multimodal-gen-unit-test-amd: 1272 passed / 5 skipped, including the 22 new test_suite_partitioning.py cases.
• CUDA, full base suite with the plan-driven path plus the two no-plan jobs (multimodal-gen-test-1-5090, bcg-diffusion): run 31443692138 (attempt 2).

Accuracy Tests

N/A — CI orchestration only.

Speed Tests and Profiling

N/A.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #31636186876
Latest PR Test (Extra): ❌ Run #31636186592

[kangwangamd]

Offline-validated this fix against the exact 08/09 failure — it greens the 2-GPU multimodal-gen shards. 👍

Repro of the pre-fix abort (run 31327057626, multimodal-gen-test-2-gpu-amd-rocm720 parts 0/1/2): the 2-gpu suite has 7 standalone files but the matrix declares 3 partitions, so the old logic computes parametrized_partitions = 3 - 7 = -4 < 0 and hits:

Error: total_partitions (3) must be >= standalone files (7)   # exit 1, before any test runs

With this PR's partitioning.py / assign_partition() (head 6f77ef9d3), LPT-balancing all items (parametrized cases + standalone files) across the 3 partitions instead of reserving one partition per standalone file:

part 0: 4 items (~805s)   part 1: 4 items (~755s)   part 2: 4 items (~755s)

• ✅ every item scheduled exactly once (no missing / no dupes)
• ✅ no partition empty → no abort
• ✅ all 7 standalone files placed
• ✅ coverage-once also holds for total_partitions = 1, 2, 4, 8

So all 3 red 2-GPU multimodal-gen-test shards from that run will go green with this change. (The remaining 1-GPU flux_image_t2i failure in the same run is a separate ROCm PTX-inline-asm codegen bug, tracked in #34351 / #34352 — not addressed here.)

Tested on MI300X (gfx942), image rocm/sgl-dev:v0.5.17-rocm720-mi30x-20260809; pure-Python harness logic, no GPU required. LGTM.

[yctseng0211]

/rerun-failed-ci

[michaelzhang-ai]

/rerun-failed-ci

[michaelzhang-ai]

@amd-bot ci-status

[amd-bot]

@michaelzhang-ai

CI Status for PR #34328

Merge verdict: The fix works — the exact bug this PR targets (total_partitions (3) must be >= standalone files (7) abort) is gone, and the multimodal-gen shards now actually distribute and run tests. However, PR CI is still incomplete (AMD run 31443692177 is in-progress; a manual retest 31445683539 is still running) and the AMD/NPU multimodal-gen lanes remain RED for a different, pre-existing reason (model/attention-backend failures), not the partitioning change. Do not merge on "green," and do not block on those red X's either — they are unrelated to this diff.

Note

This PR's changed code is exercised and verified: the new unit test test/unit/test_suite_partitioning.py (covering assign_partition / build_local_partition_assignment) passed in multimodal-gen-unit-test-amd ✅, and the e2e AMD/NPU shards prove the guard-removal path runs (they no longer abort; they execute and deselect the balanced case set). The residual failures live in model-serving code this PR does not touch.

Changed files: .github/workflows/pr-test-amd.yml (+1/-1), .github/workflows/pr-test-amd-rocm720.yml (+1/-1), partitioning.py (+14/-0), run_suite.py (+144/-237), scripts/gen_diffusion_ci_outputs.py (+4/-2), test/unit/test_suite_partitioning.py (+68/-0, new)

Executed CI failure attribution: AMD: multiple multimodal-gen shard failures (0 related to this PR) · NPU: 1 shard failure (0 related) · Others: gate/finish/extra jobs are fast-fail cascades of the above. AMD run still pending (do not count as passed).

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
multimodal-gen-test-2-gpu-amd (shard 0, retest)	test/server/test_server_2_gpu.py	server startup ltx_2.3_one_stage_ti2v	ValueError: Attention backend 'aiter' is not supported by this attention layer; supported backends: ['torch_sdpa']	🟢	Model/attention-backend gap for LTX-2.3 on AMD; PR only changes shard assignment, not backend selection
multimodal-gen-test-2-gpu-amd (shard 0)	test/server/test_server_2_gpu.py	ltx_2_3_two_stage_ti2v_2gpus, ltx_2.3_one_stage_ti2v	same aiter backend error (5 passed, 2 errors)	🟢	Tests now run (guard fixed); LTX aiter failure is independent of partitioning
multimodal-gen-test-1-gpu-amd (shard 0)	test/server/test_server_1_gpu.py	test_diffusion_generation[flux_image_t2i], [joyai_image_edit_ti2i]	ValueError: ... mm_token_type_ids is missing / APITimeoutError (2 failed, 34 deselected)	🟢	M-RoPE / processor issue in model code; unrelated to shard packing

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
multimodal-gen-test-2-npu-a3	test/server/ascend/test_server_2_npu.py	[qwen_image_t2i_2npu], [flux_2_image_t2i_2npu], [wan2_2_t2v_14b_w8a8_2npu]	InternalServerError / Server exited early (code 1) (1 failed, 2 passed, 2 errors)	🟢	Tests run (partition path works on NPU too); failures are model-serving errors, not partitioning
pr-gate / *-finish / PR Test Extra (AMD)	N/A	N/A	fast-fail/aggregation cascade	🟢	These roll up the child failures above; not independent failures

Details / what to do before merge

• The partition fix is verified. Compare to the baseline symptom in the PR body: shards previously aborted with total_partitions (3) must be >= standalone files (7) before running anything. That message no longer appears; shards now report real pytest results (5 passed … 2 errors, 2 failed, 34 deselected, etc.). The new build_local_partition_assignment LPT packing is doing its job.
• The red lanes are pre-existing model-support gaps, not this PR. The diff touches only partitioning/run-suite/workflow-comment code — no model, attention-backend, or processor code. The two residual clusters are:
  • Attention backend 'aiter' is not supported … supported: ['torch_sdpa'] for LTX-2.3 (2-gpu). Related open work exists (e.g. #32926 AITER diffusion backend), but no exact in-flight fix for LTX-2.3 aiter support was found.
  • mm_token_type_ids is missing (flux/joyai, 1-gpu) and NPU InternalServerError — separate model/processor issues.
    These were masked before because the guard aborted the whole shard; this PR correctly surfaces them. They should be tracked/fixed separately and should not gate this PR.
• Before merge: let AMD run 31443692177 and the manual retest 31445683539 finish so the full shard matrix (1-gpu shard 3, 2-gpu shards 1/2) is observed — confirm none of them regress to a partitioning error (they should only ever show the model-level failures above). If any shard shows a partition/scheduling error, that would be PR-related and worth a second look.

---

Generated by amd-bot using Claude Code CLI

[michaelzhang-ai]

Thanks @bingxche — all three points addressed.

Fast-fail skipping standalone files (0ad51a5): fixed, and you're right that 1-gpu is the real regression — on HIP the only standalone file left is test_generate_zimage_turbo_cli.py, and with flux_image_t2i / joyai_image_edit_ti2i currently red it would be skipped every run with no coverage check to notice. I dropped the early return entirely rather than guarding it on not assignment.standalone_files: when a shard has no standalone files the end-of-function path writes the identical report and returns the same code, so behavior is unchanged there and there's one less copy of the report call. --continue-on-error keeps its meaning between standalone files. Added test_failing_cases_do_not_skip_the_shards_standalone_files to lock it down.

Unread estimated_time / missing_standalone_estimates (3eb5538): removed, along with the used_fallback_estimate plumbing that only fed them.

Description scope: added a "Which lanes change behavior" section covering the 1-gpu recomposition (3 parametrized + 1 standalone shard → 4 mixed shards) and why the other no-plan lanes are unaffected — matches your list.

Not actioned, both pre-existing and happy to take in a follow-up: the now-dead elif report["is_standalone"] in verify_diffusion_coverage.py, and _resolve_suite_files(strict=True) exiting without a report. Same family, also left alone: if not suite_files: return 0 skips standalone files too, but only when a parametrized test file is missing from the checkout.

[michaelzhang-ai]

@HaiShaw @bingxche conflict resolved. ready to merge. Thanks!