[rl] Add XPU support to upstream GRPO training pipeline

torchtitan/experiments/rl/__init__.py

@@ -36,22 +36,33 @@
 import sys
 import warnings
 
-# Avoid memory fragmentation and peak reserved memory increasing over time
-# To overwrite, set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
-if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
-    if "torch" in sys.modules:
-        warnings.warn(
-            "The 'torch' module has already been imported. "
-            "Setting PYTORCH_CUDA_ALLOC_CONF may not have an effect."
-            "For best results, set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True before importing 'torch'."
-        )
-    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+# Avoid memory fragmentation and peak reserved memory increasing over time.
+# expandable_segments corrupts XPU's oneCCL USM pointers (PyTorch bug: XPU
+# allocator reads PYTORCH_CUDA_ALLOC_CONF). Detect XPU via ZE_AFFINITY_MASK.
+if "ZE_AFFINITY_MASK" not in os.environ:
+    if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
+        if "torch" in sys.modules:
+            warnings.warn(
+                "The 'torch' module has already been imported. "
+                "Setting PYTORCH_CUDA_ALLOC_CONF may not have an effect."
+                "For best results, set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True before importing 'torch'."
+            )
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
-from torchtitan.experiments.rl.models.vllm_registry import register_to_vllm
-from torchtitan.experiments.rl.models.vllm_wrapper import VLLMModelWrapper
+
+def __getattr__(name: str):
+    if name == "register_to_vllm":
+        from torchtitan.experiments.rl.models.vllm_registry import register_to_vllm
+
+        return register_to_vllm
+    if name == "VLLMModelWrapper":
+        from torchtitan.experiments.rl.models.vllm_wrapper import VLLMModelWrapper
+
+        return VLLMModelWrapper
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
 __all__ = [
     "VLLMModelWrapper",
-    "register_to_vllm",  # Export register function for manual use
+    "register_to_vllm",
 ]

torchtitan/experiments/rl/actors/generator.py

@@ -143,9 +143,9 @@ def _prepare_generation_request_metrics(
                 inputs.last_token_ts - inputs.first_token_ts
             ) * 1000
             metric_values[f"{prefix}/decode_time_ms"] = first_to_last_token_ms
-            metric_values[
-                f"{prefix}/inter_token_latency_ms"
-            ] = first_to_last_token_ms / (inputs.num_generation_tokens - 1)
+            metric_values[f"{prefix}/inter_token_latency_ms"] = (
+                first_to_last_token_ms / (inputs.num_generation_tokens - 1)
+            )
 
     # Emit each value with both Mean and Max aggregators.
     return [
@@ -722,6 +722,13 @@ class Config(Configurable.Config):
         the new weights. No effect under strict-drain (engine idle at pull time); async hot-swap only.
         Default True to avoid reusing stale-weight KV."""
 
+        attention_config: AttentionConfig | None = None
+        """vLLM attention config passed to the engine. When None (default), the
+        backend is derived from the model's inner_attention config
+        (FlexAttention -> FLEX_ATTENTION, VarlenAttention -> CUSTOM).
+        Platforms that don't support those backends (e.g. XPU) should set this
+        to AttentionConfig(backend=<supported_backend>)."""
+
         def __post_init__(self):
             # The generator runs vLLM full expert parallelism: vLLM forms the EP
             # group from all DP*TP ranks, so expert_parallel_degree must equal
@@ -818,6 +825,7 @@ def __init__(
 
         # Build vLLM engine
         enable_ep = config.parallelism.expert_parallel_degree > 1
+
         engine_kwargs = dict(
             # ``model`` is the path to the HF checkpoint directory. The
             # config is sourced from torchtitan's ModelSpec via
@@ -844,7 +852,7 @@ def __init__(
             distributed_executor_backend="external_launcher",
             gpu_memory_utilization=config.gpu_memory_limit,
             enforce_eager=not config.cudagraph.enable,
-            attention_config=AttentionConfig(
+            attention_config=config.attention_config or AttentionConfig(
                 backend=(
                     AttentionBackendEnum.FLEX_ATTENTION
                     if isinstance(inner_attn, FlexAttention.Config)
@@ -1128,11 +1136,13 @@ async def _decide_next_action(self) -> LoopDecision:
         # the predicate. In-flight requests keep the predicate true, so they need no notify.
         async with self._engine_loop_condition:
             await self._engine_loop_condition.wait_for(
-                lambda: self._close_request is not None
-                or self._model_state_dict_pull_request is not None
-                or self._queued_generation_requests
-                # In-flight requests (on any DP rank) keep rank 0 issuing STEP.
-                or self._request_dispatcher.rank0_has_pending_futures()
+                lambda: (
+                    self._close_request is not None
+                    or self._model_state_dict_pull_request is not None
+                    or self._queued_generation_requests
+                    # In-flight requests (on any DP rank) keep rank 0 issuing STEP.
+                    or self._request_dispatcher.rank0_has_pending_futures()
+                )
             )
 
             if self._close_request is not None:
@@ -1206,9 +1216,9 @@ async def pull_model_state_dict(self, version: int) -> None:
         self._rank0_check_engine_loop_running("pull_model_state_dict")
 
         # A placeholder future for the engine loop to resolve once the pull has been applied.
-        pull_model_state_dict_future: asyncio.Future[
-            int
-        ] = asyncio.get_running_loop().create_future()
+        pull_model_state_dict_future: asyncio.Future[int] = (
+            asyncio.get_running_loop().create_future()
+        )
 
         # `_engine_loop_condition` wakes the engine loop, if asleep, when a pull is queued.
         async with self._engine_loop_condition:

torchtitan/experiments/rl/actors/trainer.py

@@ -371,8 +371,7 @@ async def forward_backward(
             dict[str, float]: Globally-reduced metrics.
         """
         logger.debug(
-            f"{os.getpid()=} PolicyTrainer forward_backward "
-            f"step {self.policy_version}"
+            f"{os.getpid()=} PolicyTrainer forward_backward step {self.policy_version}"
         )
 
         # RL does not support pipeline parallelism yet, so the trainer

torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py

@@ -13,8 +13,11 @@
 
 import dataclasses
 
+import torch
+
 from torchtitan.components.checkpoint import CheckpointManager
 from torchtitan.components.loss import ChunkedLossWrapper
+from torchtitan.components.lora import LoRAConverter
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 from torchtitan.components.optimizer import default_adamw
 from torchtitan.config import (
@@ -24,6 +27,8 @@
     ParallelismConfig,
     TrainingConfig,
 )
+from vllm.config import AttentionConfig
+
 from torchtitan.experiments.rl.actors.generator import (
     SamplingConfig,
     VLLMCudagraphConfig,
@@ -191,6 +196,92 @@ def rl_grpo_qwen3_0_6b_flex() -> Controller.Config:
     )
 
 
+def rl_grpo_lora_qwen3_0_6b() -> Controller.Config:
+    """GRPO + LoRA config for Qwen3-0.6B with flex attention (4 GPUs: 2 gen + 2 train).
+
+    Uses FSDP (dp_shard=2) instead of TP for training, LoRA adapters (rank=8).
+
+    On XPU, set ZE_AFFINITY_MASK for device isolation and
+    TORCHINDUCTOR_MAX_AUTOTUNE=0 to avoid backward kernel resource exhaustion.
+    """
+    group_size = 8
+    model_spec = _qwen3_rl_model_registry(
+        "0.6B",
+        attn_backend="flex",
+        converters=[
+            LoRAConverter.Config(rank=8, alpha=16.0, target_modules=["wqkv", "wo"]),
+        ],
+    )
+    # Disable max_autotune for XPU: backward autotuning tries kernel configs
+    # that exceed XPU register limits (OUT_OF_RESOURCES). Harmless on CUDA
+    # (just uses the default heuristic config instead of autotuning).
+    # Must redefine _compiled_flex_attn since torch.compile captures options at
+    # definition time.
+    from torch.nn.attention.flex_attention import flex_attention
+    from torchtitan.models.common.attention import FlexAttention
+
+    FlexAttention.inductor_configs = {
+        **FlexAttention.inductor_configs,
+        "max_autotune": False,
+        "coordinate_descent_tuning": False,
+    }
+    FlexAttention._compiled_flex_attn = torch.compile(
+        flex_attention, options=FlexAttention.inductor_configs
+    )
+    return Controller.Config(
+        model_spec=model_spec,
+        hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B",
+        async_loop=AsyncLoopConfig(
+            num_training_steps=10,
+            num_groups_per_train_step=8,
+            group_size=group_size,
+            validation=ValidationConfig(num_samples=20),
+            batcher=Batcher.Config(
+                batch=BatchConfig(local_batch_size=2, seq_len=2048),
+            ),
+        ),
+        compile=CompileConfig(enable=True, backend="aot_eager"),
+        rollouter=AlphabetSortRollouter.Config(),
+        renderer=RendererConfig(name="qwen3", enable_thinking=False),
+        metrics=MetricsProcessor.Config(enable_wandb=False),
+        trainer=PolicyTrainer.Config(
+            optimizer=default_adamw(lr=2e-6),
+            lr_scheduler=LRSchedulersContainer.Config(
+                warmup_steps=2,
+                decay_type="linear",
+            ),
+            training=TrainingConfig(dtype="bfloat16"),
+            parallelism=ParallelismConfig(
+                data_parallel_shard_degree=2,
+                tensor_parallel_degree=1,
+            ),
+            checkpoint=CheckpointManager.Config(
+                enable=True,
+                initial_load_in_hf=True,
+                interval=10,
+                last_save_model_only=False,
+            ),
+            loss=GRPOLoss.Config(),
+        ),
+        generator=VLLMGenerator.Config(
+            model_dtype="bfloat16",
+            gpu_memory_limit=0.85,
+            cudagraph=VLLMCudagraphConfig(enable=False),
+            attention_config=AttentionConfig(),
+            parallelism=InferenceParallelismConfig(
+                data_parallel_degree=2,
+                tensor_parallel_degree=1,
+            ),
+            checkpoint=CheckpointManager.Config(enable=False),
+            sampling=SamplingConfig(
+                temperature=0.8,
+                top_p=0.95,
+                max_tokens=100,
+            ),
+        ),
+    )
+
+
 def rl_grpo_qwen3_0_6b_flex_batch_invariant() -> Controller.Config:
     """GRPO training config for Qwen3-0.6B with flex attention and batch invariance
     for bitwise-identical numerics between trainer and generator (4 GPUs: 2 gen + 2 train).

torchtitan/experiments/rl/train.py

@@ -25,15 +25,15 @@
 import asyncio
 import logging
 import os
+from collections.abc import Callable
 from dataclasses import dataclass
 
 # must run before torch import. Set it as early as possible to avoid other
 # imports transitively importing torch.
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+# expandable_segments corrupts XPU's oneCCL USM pointers. Detect XPU via ZE_AFFINITY_MASK.
+if "ZE_AFFINITY_MASK" not in os.environ:
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-# TODO: Remove `_src` after monarch cuts a new release. This is already a public
-# API in monarch nightly. https://github.com/meta-pytorch/monarch/pull/4327
-from monarch._src.actor.host_mesh import default_bootstrap_cmd
 from monarch.actor import HostMesh, ProcMesh, this_host
 
 from torchtitan.config import ConfigManager, ParallelismConfig
@@ -52,7 +52,7 @@ def breakable_cudagraph_env(generator_cfg) -> dict[str, str]:
     makes prefill attention a cudagraph break (run eager at replay). The import happens before the
     generator actor's ``__init__`` and in the vLLM EngineCore worker subprocesses -- which do NOT
     inherit a runtime-set ``os.environ`` -- so setting it at runtime is too late. It must go in the
-    proc's LAUNCH env (the spawn ``bootstrap_command``, or the MAST role.env). Without it the decorator no-ops
+    proc's LAUNCH env (the spawn bootstrap, or the MAST role.env). Without it the decorator no-ops
     and prefill attention is captured as ``output.fill_(0)`` (zeroed) -> the model never reads the
     prompt -> coherent-but-unrelated output. FULL_DECODE_ONLY never captures prefill so it needs
     nothing. Shared so the OSS spawn path and the fbcode MAST launcher use one source of truth.
@@ -67,21 +67,14 @@ def breakable_cudagraph_env(generator_cfg) -> dict[str, str]:
     return {}
 
 
-def _preimport_torch() -> None:
-    """``bootstrap`` setup callable: pre-import torch on the spawned proc.
-    """
-    # TODO: Remove once Monarch/PyTorch fixes concurrent import during unpickling.
-    import torch  # noqa: F401
-
-
 class PerHostProvisioner:
     """Allocates non-overlapping GPU ranges within a single host.
 
     On the same host, the trainer and generator run on separate GPU
     meshes (e.g. GPUs 0-3 for training, GPUs 4-7 for generation). Each
-    call to `allocate(n)` reserves the next *n* GPUs and returns the launch
-    env (`CUDA_VISIBLE_DEVICES` plus any `extra_env`) for those GPUs. The
-    returned env vars are applied to the spawned process via ``bootstrap_command``.
+    call to `allocate(n)` reserves the next *n* GPUs and returns a
+    bootstrap callable that sets device visibility env vars before the
+    accelerator runtime initializes in the spawned process.
     """
 
     def __init__(self, total_gpus: int = 8):
@@ -94,7 +87,7 @@ def available(self) -> int:
 
     def allocate(
         self, num_gpus: int, *, extra_env: dict[str, str] | None = None
-    ) -> dict[str, str]:
+    ) -> Callable[[], None]:
         if num_gpus > self.available:
             raise RuntimeError(
                 f"Requested {num_gpus} GPUs but only {self.available} "
@@ -103,10 +96,22 @@ def allocate(
         gpu_ids = list(range(self.next_gpu, self.next_gpu + num_gpus))
         self.next_gpu += num_gpus
 
-        env = {"CUDA_VISIBLE_DEVICES": ",".join(str(g) for g in gpu_ids)}
-        if extra_env:
-            env.update(extra_env)
-        return env
+        def _bootstrap():
+            gpu_str = ",".join(str(g) for g in gpu_ids)
+            # Device visibility must be set before torch import initializes
+            # the accelerator runtime. Detect XPU by checking inherited env
+            # from the launcher (ZE_AFFINITY_MASK is standard for XPU systems).
+            if "ZE_AFFINITY_MASK" in os.environ:
+                os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+                os.environ["ZE_AFFINITY_MASK"] = gpu_str
+            else:
+                os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str
+            if extra_env:
+                os.environ.update(extra_env)
+            # TODO: Remove once Monarch/PyTorch fixes concurrent import during unpickling.
+            import torch  # noqa: F401
+
+        return _bootstrap
 
 
 @dataclass
@@ -151,11 +156,9 @@ def _spawn_proc_mesh(
     )
     role_gpus_per_node = role_world_size // nodes
     provisioner = PerHostProvisioner(total_gpus=gpus_per_node)
-    env = provisioner.allocate(role_gpus_per_node, extra_env=extra_env)
     return host_mesh.spawn_procs(
         per_host={"gpus": role_gpus_per_node},
-        bootstrap=_preimport_torch,
-        bootstrap_command=default_bootstrap_cmd().with_env(env),
+        bootstrap=provisioner.allocate(role_gpus_per_node, extra_env=extra_env),
     )
 
 
@@ -213,24 +216,18 @@ def spawn_proc_mesh(
         ]
     else:
         # Single-node mode: partition GPUs on this_host() via
-        # CUDA_VISIBLE_DEVICES
+        # device visibility env vars (CUDA_VISIBLE_DEVICES or ZE_AFFINITY_MASK)
         host_mesh = this_host()
         provisioner = PerHostProvisioner(total_gpus=total_gpus)
         trainer_mesh = host_mesh.spawn_procs(
             per_host={"gpus": trainer_world_size},
-            bootstrap=_preimport_torch,
-            bootstrap_command=default_bootstrap_cmd().with_env(
-                provisioner.allocate(trainer_world_size)
-            ),
+            bootstrap=provisioner.allocate(trainer_world_size),
         )
         generator_meshes = [
             host_mesh.spawn_procs(
                 per_host={"gpus": per_generator_world_size},
-                bootstrap=_preimport_torch,
-                bootstrap_command=default_bootstrap_cmd().with_env(
-                    provisioner.allocate(
-                        per_generator_world_size, extra_env=generator_env
-                    )
+                bootstrap=provisioner.allocate(
+                    per_generator_world_size, extra_env=generator_env
                 ),
             )
             for _ in range(num_generators)