[training_utils] fix: no allocator set when using TMA for kernels

[HollowMan6]

What does this PR do?

Refer to triton-lang/triton@43625fc

Triton 3.5.0 stores allocators in ContextVar; values do not propagate to new threads by default. Some execution paths in verl use thread pools (e.g., concurrent.futures), so we set the allocator as ContextVar default to avoid falling back to NullAllocator in worker threads.

Related traceback:

   File "/usr/lib/python3.12/concurrent/futures/_base.py", line 456, in result
     return self.__get_result()
            ^^^^^^^^^^^^^^^^^^^
   File "/usr/lib/python3.12/concurrent/futures/_base.py", line 401, in __get_result
     raise self._exception
            ^^^^^^^^^^^^^^^^^^^^^
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/single_controller/ray/base.py", line 841, in func
     return getattr(self.worker_dict[key], name)(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/single_controller/base/decorator.py", line 456, in inner
     return func(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/transferqueue_utils.py", line 314, in dummy_inner
     output = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 105, in f
     return self.log(decorated_function, *args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 118, in log
     output = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/profile.py", line 256, in wrapper
     return func(self_instance, *args, **kwargs_inner)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/workers/megatron_workers.py", line 827, in compute_ref_log_prob
     return self.compute_log_prob(data)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/single_controller/base/decorator.py", line 456, in inner
     return func(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/transferqueue_utils.py", line 314, in dummy_inner
     output = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 105, in f
     return self.log(decorated_function, *args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 118, in log
     output = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/profile.py", line 256, in wrapper
     return func(self_instance, *args, **kwargs_inner)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/workers/megatron_workers.py", line 870, in compute_log_prob
     output, entropys, layers_topk_idx = self.actor.compute_log_prob(data=data, calculate_entropy=not is_lora)
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 105, in f
     return self.log(decorated_function, *args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/profiler/performance.py", line 118, in log
     output = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/workers/actor/megatron_actor.py", line 235, in compute_log_prob
     output = self.forward_backward_batch(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/workers/actor/megatron_actor.py", line 693, in forward_backward_batch
     losses_reduced = forward_backward_func(
                      ^^^^^^^^^^^^^^^^^^^^^^
   File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 632, in forward_backward_no_pipelining
     output_tensor, num_tokens = forward_step(
                                 ^^^^^^^^^^^^^
   File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 417, in forward_step
     output_tensor, loss_func = forward_step_func(data_iterator, model)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/workers/actor/megatron_actor.py", line 608, in forward_step
     output = forward_fn(
              ^^^^^^^^^^^
   File "/root/verl/verl/models/mcore/model_forward_fused.py", line 117, in fused_forward_model
     output_orig: CausalLMOutputForPPO = model(**input_args)
                                         ^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/distributed/data_parallel_base.py", line 22, in forward
     return self.module(*inputs, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/transformer/module.py", line 456, in forward
     outputs = self.module(*inputs, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
     return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
     return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/models/mcore/model_forward_fused.py", line 209, in _fused_GPTModel_forward
     logprobs, entropy = linear_cross_entropy(
                         ^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 581, in apply
     return super().apply(*args, **kwargs)  # type: ignore[misc]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/kernel/linear_cross_entropy.py", line 77, in forward
     logprobs, entropy, _maximum, _accumulate, _entropy_b = kernels.efficient_entropy_forward(
                                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/root/verl/verl/utils/kernel/kernels.py", line 626, in efficient_entropy_forward
     efficient_entropy_kernel_general_mainloop[mainloop_grid](
   File "/usr/local/lib/python3.12/dist-packages/triton/runtime/jit.py", line 419, in <lambda>
     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/triton/runtime/autotuner.py", line 250, in run
     ret = self.fn.run(
           ^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/triton/runtime/jit.py", line 757, in run
     kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata,
   File "/usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/driver.py", line 709, in __call__
     global_scratch = allocate_scratch(self.global_scratch_size, self.global_scratch_align, _allocation._allocator)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/driver.py", line 706, in allocate_scratch
     return alloc_fn(alloc_size, align, stream)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/usr/local/lib/python3.12/dist-packages/triton/runtime/_allocation.py", line 20, in __call__
     raise RuntimeError("Kernel requires a runtime memory allocation, but no allocator was set. " +
 RuntimeError: Kernel requires a runtime memory allocation, but no allocator was set. Use triton.set_allocator to specify an allocator.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

_{✨ Presented to you with Mind Lab - A Lab for Experiential Intelligence.}

[gemini-code-assist]

Code Review

This pull request addresses a critical runtime error when using Triton's Tensor Memory Accelerator (TMA) with thread pools by correctly setting the default allocator in a ContextVar. The approach of patching Triton's internal _allocator is a good workaround for the issue. My main feedback is to improve the robustness of the error handling in the patch to avoid silently ignoring unexpected errors, which could make future debugging difficult.