[yuchenwang3]

What does this PR do?

Make worker_extension_cls_fqn the common selection path for policy and vLLM workers. Today a config-provided FQN is rejected whenever quant_cfg is set, even if it names the correct quantized worker, while the constructor override can replace that worker without the same check.

The shared resolver fills an omitted FQN with the backend's quantized worker and warns, accepts an exact match, and rejects a mismatch with the expected FQN before allocating workers. Existing quant-only recipes keep automatic worker selection.

The Policy constructor argument now emits a deprecation warning and feeds the config path. The research template selects its extension in YAML, and frozen teachers clear the copied student extension alongside quantization settings.

Issues

Closes #4032.

Usage and compatibility

Leave worker_extension_cls_fqn unset to keep automatic selection, or set it to the FQN printed by the warning. An explicit quant worker must match the policy backend or vLLM sync/async mode. The Automodel backend uses automodel_cfg. The CLI entrypoints' existing check_outdated_config() rejects the old dtensor_cfg block and the removed _v2 key for either value, including when an extension is provided.

Compatibility change: custom subclasses passed through Policy(worker_extension_cls_fqn=...) with quant_cfg are now rejected, just like config-provided mismatches. Non-quantized custom extensions remain supported. The deprecated argument still works when it satisfies the config validation rules; conflicting constructor/config values remain errors.

Validation (October 8)

Rebased onto main 2c2ecd2ea814d666edae2a8ce8b458403cb692b1 at 24c27600e4df9330b48e90fb78850ae2e0f7150e. The conflicts retain upstream vLLM FP32 patch imports and the Automodel class/config rename. The template YAML and tests use the new FQNs. Stale-config rejection stays in main's unchanged check_outdated_config() entrypoint path; eight cases cover quantization with an omitted or explicit worker FQN, the old/new backend block names, and both _v2 values.

• 144 focused CPU tests passed, with one expected eval-config skip: 73 worker/policy cases, 21 vLLM generation cases, and 50 outdated-config cases. Driver cases mock Ray workers, model loading, and GPU execution.
• The broader shipped-config suite ran 1,088 passed and 345 skipped. Eighteen cases initially failed because local HF_HOME was unset; the same failures reproduce on main, and all 18 pass after setting it. No test-code workaround or project dependency change was made.
• Ruff, import order, format and Python syntax passed for all 15 changed Python files. YAML/TOML parsing, the template extension FQN, recipe minimization, and git diff --check passed. Targeted Pyrefly on the shared helper and policy/vLLM resolvers reported zero shown errors (9 ignored under repository configuration). The full-repository Pyrefly hook was skipped.
• Tests used Python 3.13.14 on macOS with the existing .venv-4032, --no-sync, and --noconftest. Full-repository type checking, distributed tests, and GPU model execution were not run. These tests validate worker selection; upstream's Automodel quantized worker is still a placeholder. NVIDIA-runner CI requires maintainer authorization for this SHA.

Reproduce the focused selections in a configured test environment:

PYTHONPATH=. TORCH_CUDA_ARCH_LIST=9.0 uv run --group test pytest --noconftest -o addopts='' -q \
  tests/unit/models/test_worker_config.py \
  tests/unit/models/policy/test_policy_validation.py \
  tests/unit/models/generation/test_vllm_worker_selection.py \
  tests/unit/models/policy/test_teacher_worker_group.py \
  -k 'worker_extension or quant_worker or quantization or unquantized_worker or clears_student_settings or fp32_lm_head'

PYTHONPATH=. TORCH_CUDA_ARCH_LIST=9.0 uv run --group test pytest --noconftest -o addopts='' -q \
  tests/unit/models/generation/test_vllm_generation.py -k worker_extension

PYTHONPATH=. TORCH_CUDA_ARCH_LIST=9.0 uv run --group test pytest --noconftest -o addopts='' -q \
  tests/unit/utils/test_outdated_config_checks.py

--noconftest skips the repository's autouse Ray/GPU-monitor fixtures for these CPU cases. Locally, these commands used UV_PROJECT_ENVIRONMENT=.venv-4032 uv run --no-sync.

Pre-checks

• Read the contributor guidelines.
• Added regression tests and updated existing expectations.
• Full unit and GPU functional suites run locally (focused CPU results above).
• Updated quantization docs and the research template documentation.

Assisted-by: OpenAI Codex

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuchenwang3]

Rebased onto main (fd7112c3) at fe6d468e6f372be4665e352f63f2cdfbe62c543c. The conflict resolution preserves the upstream FP32 LM-head validation and the new teacher_index argument.

The focused CPU selection passes 84 tests, including the FP32 LM-head cases. Changed-file formatting, lint, TOML, and recipe checks pass. Targeted Pyrefly reports zero errors under the repository configuration. The full-repository Pyrefly hook could not start because the local Automodel submodule is not populated; GPU/distributed validation still requires the full environment and maintainer-authorized CI.

[yuchenwang3]

Rebased onto main (592149ef) at ea794c201418c35d848d5a245d14adeed6e4d4d2 and resolved the conflict while retaining the upstream DTensor V1 guard and V2 default. Updated the worker-selection expectations and added three cases showing that implicit/config/constructor quantized workers cannot bypass the V1 rejection before allocation.

The focused CPU selection passes 88 tests (39 deselected), including the FP32 LM-head cases. Changed-file formatting, lint, TOML, and recipe checks pass; targeted Pyrefly reports zero shown errors under the repository configuration. Full-repository Pyrefly remains blocked locally by the unpopulated Automodel submodule. GPU/distributed validation was not run and still requires maintainer-authorized CI for this SHA. The PR description now contains the current results and reproduction command.

[yuchenwang3]

@yuki-97 Rebased onto main at a297e5d6f18184ba4b39102714869341e2b49ed3. This keeps the new _v2-key rejection intact and adds coverage for both values with implicit, config, and constructor quantized-worker selection before allocation.

The focused CPU selections pass 79 + 21 tests; changed-file lint/format and targeted type checks also pass. GPU/distributed validation has not run. Could you review the worker-selection path for #4032 and help authorize CI for this SHA?

[yuchenwang3]

@yuki-97 Rebased on main and fixed the conflicts. The template and worker tests now use the Automodel names. All 144 focused CPU tests pass; GPU tests still need CI. Could you review the worker-selection change and authorize CI for 24c27600e?

[yuki-97]

hi @yuchenwang3 , sorry for late reply. thanks for helping refactor this, much more clean than before!

@RayenTian could you help review it?