[rl] Register customized config parser to vllm + less vllm config dependency by wwwjn · Pull Request #3242 · pytorch/torchtitan

wwwjn · 2026-05-06T19:07:58Z

Stack from ghstack (oldest at bottom):

vllm has this customized config parser registry support so we can plug in TorchTitan's config parser. Why we need this:

get rid of dependency on a HF format checkpoint folder when initializing. Don't implicitly depend on config.json as config source of truth

Another changes in this PR:

remove the round-trip translation from torchtitan config -> vllm config -> torchtitan config. Using closure to bypass.

[ghstack-poisoned]

wwwjn · 2026-05-06T19:17:16Z

+    # Dynamic config parser class capturing ModelSpec (and any registration-
+    # time custom fields) in the closure.
+    @register_config_parser(TORCHTITAN_CONFIG_FORMAT)
+    class TorchTitanConfigParserForSpec(ConfigParserBase):


Nested in the registry() function because we need to access model_spec when registering config_parser via closure

wwwjn · 2026-05-06T19:18:50Z

Not related to this PR. should be cleaned.

wwwjn · 2026-05-06T19:32:55Z

            during model construction.
    """
    from torchtitan.experiments.rl.models.vllm_wrapper import TorchTitanVLLMModelWrapper
+    from transformers import PretrainedConfig


We / RL side have to depend on PretrainedConfig definition from transformers as it's required to as return-ed type of ConfigParser. I think ConfigParser is a clean abstraction, but we would need to depend on transformers unfortunately

vllm has this customized config parser registry support so we can plug in TorchTitan's config parser. It serves as 2 purpose: 1. get rid of dependency on a HF format checkpoint folder when initializing 2. Passing customized args to VLLMModelWrapper, eg CompileConfig, skip_init_load_weights [ghstack-poisoned]

tianyu-l · 2026-05-06T20:11:01Z

-                compile_config=compile_config,
            )

-    # Set the class name so vLLM can identify it


why removing these comments

vllm has this customized config parser registry support so we can plug in TorchTitan's config parser. It serves as 2 purpose: 1. get rid of dependency on a HF format checkpoint folder when initializing 2. Passing customized args to VLLMModelWrapper, eg CompileConfig, skip_init_load_weights [ghstack-poisoned]

… config dependency" vllm has this customized config parser registry support so we can plug in TorchTitan's config parser. Why we need this: - get rid of dependency on a HF format checkpoint folder when initializing. Don't implicitly depend on `config.json` as config source of truth Another changes in this PR: - remove the round-trip translation from torchtitan config -> vllm config -> torchtitan config. Using closure to bypass. [ghstack-poisoned]

wwwjn · 2026-05-08T22:59:27Z

    device_id: torch.device | None = None
    if comm_config.mode == "torchcomms":
        try:
-            import torchcomms  # noqa: F401  # pyrefly: ignore [missing-import]


Will be reverted, not related to this PR

wwwjn · 2026-05-08T23:01:25Z

 from torchtitan.components.optimizer import OptimizersContainer
-from torchtitan.config import CommConfig, Configurable, TORCH_DTYPE_MAP
-from torchtitan.config.configs import (
+from torchtitan.config import (


This change is just consolidate the import path

tianyu-l · 2026-05-08T23:03:29Z

+            if p.enable_sequence_parallel:
+                logger.warning(
+                    "Generator enable_sequence_parallel=True hurts inference "
+                    "throughput; prefer SP=False."
+                )


this won't be supported by spmd_types erasure mode I think, so I don't mind we ban it
cc @pianpwk

hmm are we not supporting SP in spmd_types

we support SP in spmd_types when sequence length is evenly divisible; we don't support SP when sequence length is not evenly divisible.

spmd_types doesn't handle padding & unpadding like DTensor.

uneven usually show up only in inference

for inference we always use non-SP (vanilla TP)

tianyu-l · 2026-05-08T23:03:56Z

@@ -199,14 +214,17 @@ def __init__(
        engine_kwargs = dict(
            model=model_path,


what is this path for?

tianyu-l · 2026-05-08T23:05:25Z


        assert vllm_config is not None, "vllm_config is required"

+        # PP and CP are not supported on this inference path


this "raise ValueError" may better happen at grpo trainer post_init, to be consistent

here we only need assert

tianyu-l · 2026-05-08T23:21:55Z

+            wrapper's parallelize step.
    """
    from torchtitan.experiments.rl.models.vllm_wrapper import TorchTitanVLLMModelWrapper
+    from transformers import PretrainedConfig


oh, can we not depend on this? I know it's already implicit via vllm dependency, but trying to see if we can avoid explicit dependency (someday we can remove)

tianyu-l · 2026-05-08T23:25:08Z

+    # parser only produces HF-shaped fields; torchtitan-specific config is
+    # delivered through the model-class closure above.
+    @register_config_parser(TORCHTITAN_CONFIG_FORMAT)
+    class TorchTitanConfigParserForSpec(ConfigParserBase):


Suggested change

class TorchTitanConfigParserForSpec(ConfigParserBase):

class TorchTitanConfigParser(ConfigParserBase):

tianyu-l · 2026-05-08T23:26:58Z


-    # Create dynamic model class capturing ModelSpec in the closure
+    # Dynamic model class capturing torchtitan config in the closure.
    class TorchTitanVLLMModelFromSpec(TorchTitanVLLMModelWrapper):


maybe simplify, from a titan-centric view

Suggested change

class TorchTitanVLLMModelFromSpec(TorchTitanVLLMModelWrapper):

class VLLMModelFromSpec(VLLMModelWrapper):

tianyu-l · 2026-05-08T23:28:35Z

+            **kwargs,
+        ):
+            config_dict = model_spec_to_hf_config_dict(model_spec)
+            return config_dict, PretrainedConfig.from_dict(config_dict)


It's actually very weird that the contract is both a config_dict and a cls(config_dict), sounds redundant to me

tianyu-l · 2026-05-08T23:36:37Z

+    if moe is not None:
+        hf[
+            "num_experts"
+        ] = moe.experts.num_experts  # presence required: >0 toggles MoE/EP branches


nit: put # presence required above this field

tianyu-l · 2026-05-08T23:37:21Z

+    }

-def register_model_to_vllm_model_registry(
+    ffn = getattr(layer0, "feed_forward", None)


using layer0 is not robust? what if a transformer has 1st layer MoE, 2nd layer FFN?

tianyu-l · 2026-05-08T23:37:28Z

+        # Unused: only v1/metrics/perf.py reads it (off by default). SwiGLU hidden == w1.out_features.
+        hf["intermediate_size"] = ffn.w1.out_features
+
+    moe = getattr(layer0, "moe", None)


config parser

45c13fb

[ghstack-poisoned]

wwwjn requested review from fegin, tianyu-l and wconstab as code owners May 6, 2026 19:07

pytorch-bot Bot added the ciflow/8gpu label May 6, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 6, 2026

This was referenced May 6, 2026

[rl] Enable TP2EP for MoE inference in vLLM wrapper #3142

Open

[WIP] Enable DP+EP for MoE inference in vLLM wrapper #3236

Open

wwwjn changed the title ~~config parser~~ [rl] Register customized config parser to vllm May 6, 2026

wwwjn commented May 6, 2026

View reviewed changes

Comment thread torchtitan/experiments/rl/actors/generator.py Outdated

tianyu-l reviewed May 6, 2026

View reviewed changes

wwwjn added 2 commits May 6, 2026 13:33

pytorch-bot Bot added the ciflow/rl label May 7, 2026

wwwjn changed the title ~~[rl] Register customized config parser to vllm~~ [rl] Register customized config parser to vllm + less vllm config dependency May 7, 2026

wwwjn commented May 8, 2026

View reviewed changes

tianyu-l reviewed May 8, 2026

View reviewed changes

		@@ -199,14 +214,17 @@ def __init__(
		engine_kwargs = dict(
		model=model_path,


		assert vllm_config is not None, "vllm_config is required"

		# PP and CP are not supported on this inference path

	class TorchTitanConfigParserForSpec(ConfigParserBase):
	class TorchTitanConfigParser(ConfigParserBase):

	class TorchTitanVLLMModelFromSpec(TorchTitanVLLMModelWrapper):
	class VLLMModelFromSpec(VLLMModelWrapper):

Conversation

wwwjn commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pianpwk May 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l May 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wwwjn commented May 6, 2026 •

edited

Loading

pianpwk May 9, 2026 •

edited

Loading

tianyu-l May 9, 2026 •

edited

Loading