ITADN

[rl] vLLM torch.compile (VLLM_COMPILE) on the generator fails: inductor "out variant should have at least one out arg"

#3669Openyichuan-w 创建于 2026-06-15
bug
Y
yichuan-wcommented
### Summary Enabling vLLM **torch.compile** for the RL generator (`CompilationConfig(mode=CompilationMode.VLLM_COMPILE, cudagraph_mode="none")`) crashes during inductor compilation, before any generation: ``` torch/_dynamo/aot_compile.py ... aot_compile_fullgraph -> torch/_inductor/standalone_compile.py ... standalone_compile -> torch/_inductor/compile_fx.py ... _compile_fx_main ... raise AssertionError("out variant should have at least one out arg") AssertionError: out variant should have at least one out arg ``` The failing rank aborts and the run dies (the rest is Monarch/NCCL shutdown noise). ### Context `torchtitan/experiments/rl` generator (`VLLMGenerator`). The model is a torchtitan model loaded into vLLM via the custom config/registry, with the **CUSTOM attention backend**. The assertion is in inductor's "out variant" handling, which suggests an op (most likely the custom attention op registered for the CUSTOM backend) has a schema inductor can't lower — an `out`-variant op with no `out` argument. ### Repro Set the generator's `CompilationConfig` to compile (`mode=VLLM_COMPILE`, any `cudagraph_mode`) instead of the current `mode=NONE` + cudagraph, then start the engine. It fails at engine init / first compile. ### Status of the alternatives - **cudagraph** (full capture at large batch): corrupts generation + NaN logprobs — tracked in #3668. - **torch.compile** (this issue): inductor `out variant` assertion crash. - **eager**: works; this is what the Search-R1 example currently ships (~2x slower than cudagraph). ### Ask - Is the custom attention op's schema (out-variant without an `out` arg) the trigger? If so, fixing the op schema would unblock torch.compile for the generator and recover generation speed without cudagraph.
8 条评论