fix(dtensor): keep v1 model loading on Transformers classes

[jepio]

What does this PR do ?

Keep DTensor V1 model loading on plain Transformers classes.

DTensor V1 loads the full checkpoint on rank 0, then distributes it through FSDP. NeMo AutoModel wrappers may enter distributed collectives during construction, while the other ranks wait for DTensor V1 to finish its rank-0 load. A separate Transformers model map keeps this loading path under DTensor V1's control. AutoModel and DTensor V2 retain the NeMo-aware default.

Issues

N/A

Usage

No configuration change is required.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

Validation:

• All changed-file pre-commit hooks pass in Lima.
• The six focused model-class resolver tests pass in Lima.
• No documentation change is needed; the default behavior remains unchanged.

[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.

[jepio]

/ok to test 0a3f073

[yuki-97]

@RayenTian could you help review?

[jepio]

PR looks good — correctly fixes a real hang hazard (DTensor V1 loads weights on rank 0 only, then broadcasts; NeMo AutoModel wrappers can enter distributed collectives during construction, which would only run on rank 0 and hang the other ranks). resolve_model_class now lets callers opt out of the NeMo-aware default via a keyword-only use_nemo_automodel flag, and DTensor V1 does so with a clear comment explaining why.

Verified:

• All call sites checked (automodel/setup.py:370 keeps the NeMo-aware default; dtensor_policy_worker_v2.py doesn't call resolve_model_class at all) — no unintended behavior change elsewhere.
• New test test_resolve_model_class_selects_requested_loader genuinely covers all 4 (use_nemo_automodel, NEMO_AUTOMODEL_AVAILABLE) branch combinations across 3 model types, correct monkeypatch targets, no coverage gaps.
• ruff and pyrefly-typecheck both pass on the changed files.
• No existing PR comments/reviews need a response.

LGTM.

Generated by Claude Code

[RayenTian]

One doc comment from a local review of this change. The fix itself looks right — pinning DTensor V1 to the plain Transformers classes makes V1's loading contract (rank-0-only from_pretrained plus an all-ranks meta from_config) independent of what NeMo AutoModel's _build_model does.

Generated by Claude Code

[RayenTian]

nemo_rl/models/policy/utils.py:59-70

Now that this block sits above HF_AUTOMODEL_FACTORY, one sentence in it is no longer true for the dict it annotates:

Models with a custom NeMo automodel impl (e.g. qwen3_5_moe) don't need an entry — the custom impl intercepts from_pretrained regardless of the parent AutoModel class.

That interception happens in _resolve_custom_model_cls_for_config, whose only production call sites are get_is_hf_model and __init_model — both inside NeMo's _build_model. HF_AUTOMODEL_FACTORY is exactly the dict the new use_nemo_automodel=False path reads, so nothing consults the custom impl there.

The comment's own example shows the gap. qwen3_5_moe is absent from both dicts, so V1 falls back to AutoModelForCausalLM, and on transformers 5.12.1 that resolves without error:

AutoModelForCausalLM        [Qwen3_5MoeConfig] -> Qwen3_5MoeForCausalLM
AutoModelForImageTextToText [Qwen3_5MoeConfig] -> Qwen3_5MoeForConditionalGeneration

i.e. a silent drop to the text-only tower, rather than the clean KeyError you'd get for qwen2_5_vl / mistral3 / internvl / smolvlm.

Suggested wording below scopes the carve-out to AUTOMODEL_FACTORY. Alternatively (or additionally) you may want to add "qwen3_5_moe": AutoModelForImageTextToText to the dict — I left that out of the suggestion since it's a behavior call rather than a doc fix.

Suggested change

	# Plain Hugging Face classes remain separate from the NeMo AutoModel wrappers so
	# callers that manage distribution can request them when NeMo AutoModel is installed.
	# Add an entry here when a model (1) uses HF's standard loading path
	# (no custom NeMo automodel impl) AND (2) its architecture isn't
	# loadable via AutoModelForCausalLM (e.g. VLMs using
	# ForConditionalGeneration / ForImageTextToText). Models with a
	# custom NeMo automodel impl (e.g. qwen3_5_moe) don't need an entry
	# — the custom impl intercepts from_pretrained regardless of the
	# parent AutoModel class. Check MODEL_ARCH_MAPPING in the NeMo
	# automodel registry to see which architectures have custom impls:
	# https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/registry.py#L32-L146
	HF_AUTOMODEL_FACTORY: Dict[str, Any] = {
	# Plain Hugging Face classes remain separate from the NeMo AutoModel wrappers so
	# callers that manage distribution can request them when NeMo AutoModel is installed.
	# Add an entry here whenever a model's architecture isn't loadable via
	# AutoModelForCausalLM (e.g. VLMs using ForConditionalGeneration /
	# ForImageTextToText). Unlike AUTOMODEL_FACTORY below, this dict is also read on
	# the ``use_nemo_automodel=False`` path (DTensor V1), where no NeMo custom impl
	# intercepts from_pretrained -- so a model that has a custom NeMo automodel impl
	# still needs an entry here when its parent AutoModel class is not
	# AutoModelForCausalLM. Check MODEL_ARCH_MAPPING in the NeMo automodel registry
	# to see which architectures have custom impls:
	# https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/registry.py#L32-L146
	HF_AUTOMODEL_FACTORY: Dict[str, Any] = {

(Nit, pre-existing: MODEL_ARCH_MAPPING's keys are architecture class names like Qwen3_5MoeForConditionalGeneration, not model_type strings.)

[RayenTian]

Thanks for adding this. Left an non-blocking comment. Approved.