initial effort

Author: acisseJZhong

torchtitan/experiments/rl/config_registry.py

@@ -104,6 +104,47 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
     )
 
 
+def rl_grpo_qwen3_30b_a3b() -> RLTrainer.Config:
+    """GRPO training config for Qwen3-30B-A3B MoE (6 GPUs: 4 gen + 2 train)."""
+    return RLTrainer.Config(
+        model_spec=model_registry("30B-A3B"),
+        hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-30B-A3B",
+        num_steps=10,
+        batch_invariant_mode=True,
+        trainer=PolicyTrainer.Config(
+            optimizer=OptimizersContainer.Config(lr=2e-6),
+            lr_scheduler=LRSchedulersContainer.Config(
+                warmup_steps=2,
+                decay_type="linear",
+            ),
+            training=TrainingConfig(),
+            parallelism=ParallelismConfig(
+                tensor_parallel_degree=2,
+                expert_parallel_degree=1,
+                expert_tensor_parallel_degree=1,
+            ),
+        ),
+        generator=VLLMGenerator.Config(
+            model_dtype="bfloat16",
+            compile=GeneratorCompileConfig(
+                backend="none",
+                cudagraph_mode="none",
+            ),
+            parallelism=ParallelismConfig(
+                tensor_parallel_degree=4,
+                data_parallel_replicate_degree=1,
+            ),
+            num_samples_per_prompt=8,
+            sampling=SamplingConfig(
+                temperature=0.8,
+                top_p=0.95,
+                max_tokens=100,
+            ),
+            attention_backend="CUSTOM",
+        ),
+    )
+
+
 def rl_grpo_qwen3_debug() -> RLTrainer.Config:
     """Debug config for quick iteration -- small model, few steps (2 GPUs: 1 gen + 1 train)."""
     return RLTrainer.Config(

torchtitan/experiments/rl/models/parallelize.py

@@ -27,6 +27,7 @@
 from torchtitan.config.configs import CompileConfig
 from torchtitan.distributed import ParallelDims
 from torchtitan.distributed.compile import apply_compile_dense_rl
+from torchtitan.models.llama4.parallelize import apply_moe_ep_tp
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +66,15 @@ def parallelize_qwen3(
             has_position_id=has_position_id,
         )
 
+    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
+        apply_moe_ep_tp(
+            model,
+            tp_mesh=parallel_dims.get_optional_mesh("tp"),
+            ep_mesh=parallel_dims.get_optional_mesh("ep"),
+            etp_mesh=parallel_dims.get_optional_mesh("etp"),
+            ep_etp_mesh=parallel_dims.get_optional_mesh(["ep", "etp"]),
+        )
+
     if (
         compile_config is not None
         and compile_config.enable
@@ -169,11 +179,6 @@ def apply_non_moe_tp(
                     "feed_forward.w3": ColwiseParallel(use_local_output=False),
                 }
             )
-        else:
-            raise ValueError(
-                "Running vLLM inference with torchtitan Qwen3 MoE model is not supported yet."
-            )
-
         parallelize_module(
             # pyrefly: ignore [bad-argument-type]
             module=transformer_block,

torchtitan/models/llama4/parallelize.py

@@ -214,7 +214,7 @@ def parallelize_llama(
             if parallel_dims.dp_replicate_enabled
             else ["efsdp"]
         )
-        edp_mesh = parallel_dims.get_optional_mesh(edp_mesh_names)
+        edp_mesh: DeviceMesh | None = parallel_dims.get_optional_mesh(edp_mesh_names)
 
         apply_fsdp(
             model,