[veomni] feat: add async_activation_offload config

docs/ascend_tutorial/zh/feature_support/ascend_backend_features.md

@@ -305,7 +305,9 @@ VeOmni 是一个统一的强化学习训练后端，专为大规模模型的高
 | `actor_rollout_ref.actor.veomni.enable_reentrant` | 是否使用可重入的梯度检查点，默认值为false|
 | `actor_rollout_ref.actor.veomni.ckpt_manager` | 检查点管理器，默认值为dcp|
 | `actor_rollout_ref.actor.veomni.init_device` | 模型权重初始化设备，支持cpu、cuda、meta、npu，默认值为meta|
-| `actor_rollout_ref.actor.veomni.activation_gpu_limit` | 激活卸载时GPU上允许保留的激活显存限制（GB），默认值为0.0|
+| `actor_rollout_ref.actor.veomni.enable_async_activation_offload` | 是否启用基于独立流的异步激活卸载，默认值为false；VeOmni后端仅支持该模式，设置`actor_rollout_ref.model.enable_activation_offload`会直接报错|
+| `actor_rollout_ref.actor.veomni.activation_offload_modules` | 异步激活卸载的模块名匹配模式，默认值为[]（从模型的`_no_split_modules`自动发现）|
+| `actor_rollout_ref.actor.veomni.activation_offload_host_cache_limit_gb` | 异步激活卸载在步与步之间保留的空闲锁页主机缓冲上限（GB），默认值为4.0|
 | `actor_rollout_ref.rollout.moe_load_balance_metrics_interval` | Rollout侧MoE专家负载指标上报间隔，默认值为0（禁用）；需要同时开启`actor_rollout_ref.rollout.enable_rollout_routing_replay`以记录路由决策|
 
 #### Router Replay 支持

docs/perf/perf_tuning.rst

@@ -123,6 +123,16 @@ Therefore, users may need to tune the ``*micro_batch_size_per_gpu`` to accelerat
    Set ``actor_rollout_ref.model.enable_activation_offload=True`` and ``critic.model.enable_activation_offload=True``.
    This often works together with gradient checkpointing to get larger micro-batch sizes and it's only available in FSDP backend now.
 
+   The VeOmni backend does not support this synchronous mode and rejects the flag; use
+   ``actor_rollout_ref.actor.veomni.enable_async_activation_offload=True`` instead, which overlaps
+   the device-to-host and host-to-device copies with compute rather than blocking on them.
+   By default the offloaded modules are discovered from the model's
+   ``_no_split_modules``; override that with
+   ``actor_rollout_ref.actor.veomni.activation_offload_modules=[model.layers.{*}]`` when the model
+   does not declare them or when only part of the stack should be offloaded. Free pinned-host
+   buffers are capped by ``actor_rollout_ref.actor.veomni.activation_offload_host_cache_limit_gb``
+   (default 4.0 GB).
+
 Tuning for Dynamic Batch Size
 -----------------------------
 

examples/grpo_trainer/run_deepseek_v4_veomni.sh

@@ -70,6 +70,7 @@ ENGINE_CONFIG=(
     actor_rollout_ref.actor.veomni.enable_fsdp_offload=True
     # actor_rollout_ref.actor.veomni.param_offload=True
     # actor_rollout_ref.actor.veomni.optimizer_offload=True
+    actor_rollout_ref.actor.veomni.enable_async_activation_offload=True
     actor_rollout_ref.actor.veomni.enable_full_shard=True
     actor_rollout_ref.actor.veomni.ulysses_parallel_size=$usp_size
     actor_rollout_ref.actor.veomni.expert_parallel_size=$expert_size
@@ -80,12 +81,14 @@ ENGINE_CONFIG=(
     actor_rollout_ref.actor.veomni.mhc_implementation=tilelang
     actor_rollout_ref.actor.veomni.qat_implementation=fp8_blockwise
     actor_rollout_ref.actor.veomni.router_replay.mode=R3
+    actor_rollout_ref.actor.veomni.cross_entropy_loss_implementation=chunk_loss
 )
 
 # Actor model config
 ACTOR_CONFIG=(
     actor_rollout_ref.model.path=$model_path
     actor_rollout_ref.model.use_remove_padding=$use_remove_padding
+    actor_rollout_ref.model.use_fused_kernels=True
     actor_rollout_ref.actor.optim.lr=$actor_lr
     actor_rollout_ref.actor.use_kl_loss=$use_kl_loss
     actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef

tests/workers/rollout/rollout_vllm/test_vllm_abort.py

@@ -23,30 +23,8 @@
 import asyncio
 import os
 import time
-from types import SimpleNamespace
 from uuid import uuid4
 
-import pytest
-
-
-def test_abort_all_requests_propagates_pause_failure():
-    from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMHttpServer
-
-    async def fail_pause_generation(**_kwargs):
-        raise RuntimeError("pause failed")
-
-    server = vLLMHttpServer.__new__(vLLMHttpServer)
-    server.engine = SimpleNamespace(
-        output_processor=SimpleNamespace(request_states={}),
-        pause_generation=fail_pause_generation,
-    )
-    server._submission_paused = False
-    server._admitting = 0
-    server._resume_event = asyncio.Event()
-    server._resume_event.set()
-    with pytest.raises(RuntimeError, match="pause failed"):
-        asyncio.run(server.abort_all_requests())
-
 
 def test_vllm_abort():
     # ==================== Configuration ====================

verl/trainer/config/_generated_ppo_veomni_trainer.yaml

@@ -59,7 +59,9 @@ actor_rollout_ref:
       mhc_implementation: eager
       qat_implementation: none
       force_use_huggingface: false
-      activation_gpu_limit: 0.0
+      enable_async_activation_offload: false
+      activation_offload_modules: []
+      activation_offload_host_cache_limit_gb: 4.0
       basic_modules: []
       pad_to_length: ${oc.select:actor_rollout_ref.actor.pad_to_length,False}
       pad_to_length_bucket: 1024
@@ -254,7 +256,9 @@ actor_rollout_ref:
       mhc_implementation: eager
       qat_implementation: none
       force_use_huggingface: false
-      activation_gpu_limit: 0.0
+      enable_async_activation_offload: false
+      activation_offload_modules: []
+      activation_offload_host_cache_limit_gb: 4.0
       basic_modules: []
       pad_to_length: ${oc.select:actor_rollout_ref.ref.pad_to_length,False}
       pad_to_length_bucket: ${oc.select:actor_rollout_ref.actor.veomni.pad_to_length_bucket,1024}
@@ -561,7 +565,9 @@ critic:
     mhc_implementation: eager
     qat_implementation: none
     force_use_huggingface: false
-    activation_gpu_limit: 0.0
+    enable_async_activation_offload: false
+    activation_offload_modules: []
+    activation_offload_host_cache_limit_gb: 4.0
     basic_modules: []
     pad_to_length: false
     pad_to_length_bucket: 1024

verl/trainer/config/engine/veomni.yaml

@@ -80,7 +80,19 @@ qat_implementation: none
 
 force_use_huggingface: false
 
-activation_gpu_limit: 0.0
+# Offload activations to CPU with VeOmni's stream-based D2H/H2D transfers. This is the only
+# activation offload mode this engine supports; model.enable_activation_offload (the synchronous
+# saved_tensors_hooks path) is rejected.
+enable_async_activation_offload: false
+
+# Module name patterns whose activations are offloaded when enable_async_activation_offload=true.
+# Supports segment-aware globs (`model.layers.*` matches direct children only) and `{*}` for
+# sequential groups (`model.layers.{*}`). Empty means auto-discovery from model._no_split_modules.
+activation_offload_modules: []
+
+# Upper bound in GB on the free pinned-host buffers async activation offload keeps for reuse
+# between steps. In-flight offloads may exceed it; 0 disables buffer reuse.
+activation_offload_host_cache_limit_gb: 4.0
 
 # List of basic modules to use
 basic_modules: []

verl/workers/config/engine.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import math
 import os
 import warnings
 from dataclasses import dataclass, field
@@ -356,8 +357,18 @@ class VeOmniEngineConfig(EngineConfig):
         load_balancing_loss_implementation (str): MoE load-balancing loss kernel.
             ``"eager"`` (default) or ``"triton"``.
         force_use_huggingface (bool): Force loading model from huggingface, default False
-        activation_gpu_limit (float): When enabling activation offload, `activation_gpu_limit` GB
-            activations are allowed to reserve on GPU, default 0.0
+        enable_async_activation_offload (bool): Offload activations to CPU with VeOmni's
+            stream-based D2H/H2D transfers. This is the only activation offload mode the VeOmni
+            engine supports; ``model.enable_activation_offload`` (the synchronous
+            ``saved_tensors_hooks`` path) is rejected. default False
+        activation_offload_modules (list[str]): Module name patterns whose activations are
+            offloaded when ``enable_async_activation_offload=True``. Supports segment-aware globs
+            (``model.layers.*`` matches direct children only) and ``{*}`` for sequential groups
+            (``model.layers.{*}``). Empty means auto-discovery from ``model._no_split_modules``,
+            which fails loudly when the model does not declare them. default []
+        activation_offload_host_cache_limit_gb (float): Upper bound in GB on the free pinned-host
+            buffers that async activation offload keeps around for reuse between steps. In-flight
+            offloads may exceed it; 0 disables buffer reuse. default 4.0
         basic_modules (list[str]): List of basic modules to use, default None
         forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False
         model_dtype (str): Model data type used to initialize the transformers model. default "fp32"
@@ -439,7 +450,9 @@ class VeOmniEngineConfig(EngineConfig):
     mhc_implementation: str = "eager"
     qat_implementation: str = "none"
     force_use_huggingface: bool = False
-    activation_gpu_limit: float = 0.0
+    enable_async_activation_offload: bool = False
+    activation_offload_modules: list[str] = field(default_factory=list)
+    activation_offload_host_cache_limit_gb: float = 4.0
     basic_modules: Optional[list[str]] = field(default_factory=list)
     pad_to_length: bool = False
     pad_to_length_bucket: int = 1024
@@ -448,6 +461,14 @@ def __post_init__(self):
         super().__post_init__()
         assert self.strategy in ["veomni"], f"strategy {self.strategy} not supported"
 
+        if not math.isfinite(self.activation_offload_host_cache_limit_gb) or (
+            self.activation_offload_host_cache_limit_gb < 0
+        ):
+            raise ValueError(
+                "activation_offload_host_cache_limit_gb must be a finite non-negative value, got "
+                f"{self.activation_offload_host_cache_limit_gb}."
+            )
+
         replacements = {
             "flash_attention_2": "veomni_flash_attention_2_with_sp",
             "flash_attention_3": "veomni_flash_attention_3_with_sp",

verl/workers/engine/veomni/transformer_impl.py

@@ -14,6 +14,7 @@
 
 
 import logging
+from contextlib import nullcontext
 from dataclasses import dataclass, field, fields
 from typing import Any, Callable, Optional, Sequence
 
@@ -23,7 +24,6 @@
 from torch.distributed.tensor import DTensor
 from veomni.arguments import MixedPrecisionConfig, OpsImplementationConfig
 from veomni.distributed import parallel_state
-from veomni.distributed.offloading import build_activation_offloading_context
 from veomni.distributed.torch_parallelize import build_parallelize_model
 from veomni.models.auto import build_foundation_model
 from veomni.models.checkpoint_tensor_loading import get_checkpoint_tensor_converter
@@ -306,6 +306,28 @@ def _get_model_config_path(self):
         """
         return self.model_config.local_hf_config_path
 
+    def _maybe_apply_async_activation_offload(self, module):
+        """Attach VeOmni's stream-based activation offload to the decoder layers.
+
+        Must run on the unparallelized module: the offload wraps each selected layer's
+        ``__call__`` so that ``saved_tensors_hooks`` sits *outside* the gradient-checkpoint
+        boundary and the FSDP2 sharding that ``build_parallelize_model`` installs. Applying it
+        afterwards would put the hooks inside the checkpoint region, where the packed
+        activations are recomputed rather than saved, so nothing would be offloaded.
+        """
+        if not self.engine_config.enable_async_activation_offload:
+            return
+
+        # Imported lazily so a VeOmni predating stream-based offload only fails when asked for it.
+        from veomni.distributed.async_offload import apply_async_activation_offload
+
+        apply_async_activation_offload(
+            module,
+            list(self.engine_config.activation_offload_modules),
+            host_cache_limit_bytes=int(self.engine_config.activation_offload_host_cache_limit_gb * 1024**3),
+        )
+        log_gpu_memory_usage("After apply async activation offload", logger=logger)
+
     def _build_model_optimizer(self):
         # build_foundation_model runs apply_ops_config(ops_implementation)
         # before constructing the model, so per-model device_patch files see
@@ -325,6 +347,8 @@ def _build_model_optimizer(self):
         )
         log_gpu_memory_usage("After load base model", logger=logger)
 
+        self._maybe_apply_async_activation_offload(module)
+
         # Applies parallel strategies to the model.
         log_gpu_memory_usage("Before parallelize model", logger=logger)
         module = build_parallelize_model(
@@ -357,11 +381,6 @@ def _build_model_optimizer(self):
         self.module = module
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
-        self.model_fwd_context, self.model_bwd_context = build_activation_offloading_context(
-            self.model_config.enable_activation_offload,
-            self.model_config.enable_gradient_checkpointing,
-            self.engine_config.activation_gpu_limit,
-        )
 
     def optimizer_step(self):
         """
@@ -420,6 +439,8 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw
             else:
                 self._router_replay.begin_replay()
 
+        ctx = torch.no_grad() if forward_only else nullcontext()
+
         # Wrap the per-step body in try/finally so the controller is always
         # reset to DISABLED even if forward / backward / postprocess raises.
         # Without this, an exception leaves _recorded / _targets pinned
@@ -430,13 +451,12 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw
             output_lst = []
 
             for micro_batch in micro_batches:
-                with self.model_fwd_context:
+                with ctx:
                     loss, meta_info = self.forward_step(
                         micro_batch, loss_function=loss_function, forward_only=forward_only
                     )
                 if not forward_only:
-                    with self.model_bwd_context:
-                        loss.backward()
+                    loss.backward()
 
                 output_lst.append(meta_info)