From dec513854b6017b6699448f6729bb07e930c4e18 Mon Sep 17 00:00:00 2001 From: guoqiong song Date: Thu, 9 Jul 2026 14:12:11 -0700 Subject: [PATCH] [rl] Add XPU support to GRPO training pipeline Enable the full GRPO RL pipeline (trainer + vLLM generator) to run on Intel XPU alongside CUDA with minimal, config-driven changes. Key changes: 1. Device detection via ZE_AFFINITY_MASK (always present on Intel GPU systems) -- no runtime torch.xpu.is_available() checks needed. 2. expandable_segments guard: PYTORCH_CUDA_ALLOC_CONF is only set when ZE_AFFINITY_MASK is absent. On XPU the CUDA allocator config corrupts oneDNN USM pointers, causing oneCCL collective failures. 3. Bootstrap callable in PerHostProvisioner sets ZE_AFFINITY_MASK (XPU) or CUDA_VISIBLE_DEVICES (CUDA) inside the child process before Level Zero / CUDA runtime initialization. 4. Generator attention_config field using vLLM's AttentionConfig type. Default (None) preserves existing behavior; XPU config passes AttentionConfig() to let vLLM's platform auto-select the backend. 5. decoder.py: gate separate_full_blocks kwarg on inspect.signature for PyTorch version compatibility. 6. XPU config in config_registry.py: disable max_autotune (avoids OUT_OF_RESOURCES on backward), override FlexAttention class var. 7. Remove xpu_compat.py monkey-patching in favor of the above. Files changed: - torchtitan/models/common/decoder.py (version compat, no behavior change) - torchtitan/experiments/rl/__init__.py (ZE_AFFINITY_MASK guard, lazy imports) - torchtitan/experiments/rl/train.py (bootstrap callable, XPU guards) - torchtitan/experiments/rl/actors/generator.py (attention_config field) - torchtitan/experiments/rl/actors/trainer.py (remove apply_patches) - torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py (XPU config) Verified: 10-step GRPO on 4x Intel Max 1550, reward 0.194 -> 0.351. --- torchtitan/experiments/rl/__init__.py | 37 +++++--- torchtitan/experiments/rl/actors/generator.py | 34 ++++--- torchtitan/experiments/rl/actors/trainer.py | 3 +- .../examples/alphabet_sort/config_registry.py | 91 +++++++++++++++++++ torchtitan/experiments/rl/train.py | 63 ++++++------- torchtitan/models/common/decoder.py | 27 +++--- 6 files changed, 182 insertions(+), 73 deletions(-) diff --git a/torchtitan/experiments/rl/__init__.py b/torchtitan/experiments/rl/__init__.py index 6643ca8dec..2965cdb6c7 100644 --- a/torchtitan/experiments/rl/__init__.py +++ b/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", ] diff --git a/torchtitan/experiments/rl/actors/generator.py b/torchtitan/experiments/rl/actors/generator.py index 6cdfb5f176..2ae8a825d7 100644 --- a/torchtitan/experiments/rl/actors/generator.py +++ b/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=).""" + 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: diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py index 75d9079751..7e76f762f4 100644 --- a/torchtitan/experiments/rl/actors/trainer.py +++ b/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 diff --git a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py index 9c0b2cea3f..5133422588 100644 --- a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py +++ b/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). diff --git a/torchtitan/experiments/rl/train.py b/torchtitan/experiments/rl/train.py index 2cc25eaa57..50b90e4141 100644 --- a/torchtitan/experiments/rl/train.py +++ b/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) diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index e170150410..7ab0a08cdc 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -5,10 +5,11 @@ # LICENSE file in the root directory of this source tree. import dataclasses +import inspect from dataclasses import dataclass import torch -from torch.nn.attention.flex_attention import and_masks +from torch.nn.attention.flex_attention import and_masks, create_block_mask from torchtitan.distributed.minimal_async_ep.api import ( maybe_update_minimal_async_ep_config, @@ -33,6 +34,10 @@ from torchtitan.protocols.module import Module, ModuleDict +_supports_separate_full_blocks = ( + "separate_full_blocks" in inspect.signature(create_block_mask).parameters +) + __all__ = ["Decoder", "TransformerBlock"] @@ -295,20 +300,16 @@ def _create_flex_attention_mask_for_document( ] B = positions.shape[0] seq_len = positions.shape[1] - return create_attention_mask( - and_masks(*mask_mods), - B, - None, - seq_len, - seq_len, + kwargs = dict( device=positions.device, BLOCK_SIZE=attn_config.inner_attention.block_size, - # when separate_full_blocks = True, kernel iterates through - # full blocks first (blocks where all elements are unmasked) - # but which blocks are "full" vs "partial" changes depending - # on the particular batch - # for batch invariance, we disable this optimization - separate_full_blocks=not is_in_batch_invariant_mode(), + ) + if _supports_separate_full_blocks: + # Iterating full blocks first is faster but which blocks are "full" + # vs "partial" changes per batch -- disable for batch invariance. + kwargs["separate_full_blocks"] = not is_in_batch_invariant_mode() + return create_attention_mask( + and_masks(*mask_mods), B, None, seq_len, seq_len, **kwargs ) def get_attention_masks(