Fallback for unsupported Hybrid Engine policies

[LiRunGuo]

Summary

• register Hybrid Engine auxiliary Linear/Embedding/LayerNorm policies only when the model contains a supported complete transformer policy
• allow unsupported architectures such as Qwen2.5 to retain the native generate() fallback
• add CPU-only unit coverage for supported and unsupported policy registration

Fixes #8263.

Why

The generic wrappers are auxiliary pieces of a complete transformer injection policy. Registering them when no transformer layer matches creates a partial inference path. With a ZeRO-3-partitioned Qwen model, container construction can then reach _mark_uc_metadata() with a one-dimensional normalization weight and fail on weight.shape[1].

The change first builds the complete-policy map and checks it against the model's module classes. If no complete policy matches, the map remains empty, no partial containers are created, and Hybrid Engine leaves the model's native generation method intact. Supported model behavior is unchanged.

Validation

• pytest -q tests/unit/hybrid_engine/test_he_policy.py (2 passed)
• pre-commit run --files deepspeed/runtime/hybrid_engine.py tests/unit/hybrid_engine/test_he_policy.py
• Qwen2.5-0.5B / Qwen2.5-Math-7B OPSD completed a full prompt epoch and a separate 200-step run on 8 x MI250 with the fallback

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[delock]

Can you add a logger.warning here for not matching inference_policies? This would give user proper expectations.

[LiRunGuo]

Hi @delock, I have addressed the requested warning in the latest update. When no compatible transformer inference policy matches, Hybrid Engine now emits a logger.warning explaining that inference acceleration is unavailable and rollout will use the model native generate() path. I also rebased the PR onto the latest master (including #8264). Verification on the rebased branch: tests/unit/hybrid_engine/test_he_policy.py passes (2/2), git diff --check passes, and all applicable pre-commit hooks pass. The PR is mergeable; could you please re-review it and merge/add it to the merge queue if it looks good? Thanks!

[delock]

Hi @LiRunGuo there is a UT failure in CI, can you take a look? Thanks!

[LiRunGuo]

Hi @delock, thanks for pointing this out. I investigated the failed CPU UT. The warning itself was emitted correctly (it appeared under Captured stdout), but the test asserted against caplog.text, which was empty because DeepSpeed custom logger output was not propagated to pytest logging capture in that CI configuration. I pushed commit cc2aebf10, which directly mocks and verifies the logger.warning call, including the warning template and model type argument, instead of depending on logger-handler configuration. Validation: the targeted tests pass (2/2), the full tests/unit/hybrid_engine suite passes (2 passed, 4 hardware-dependent skipped), all applicable pre-commit checks pass, and git diff --check passes. Could you please rerun CI and add the updated PR to the merge queue again? Thanks!

[LiRunGuo]

Hi @delock, the latest Modal CI failed with exit code 137 after progressing through roughly 55% of the tests, without any assertion failure. This appears to be a transient sandbox resource/OOM termination. Could you please rerun the failed CI check and add the PR back to the merge queue if it passes? Thanks!