[rl] vLLM torch.compile (VLLM_COMPILE) on the generator fails: inductor "out variant should have at least one out arg"
bug
### 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 条评论