[dist] refactor: inject the async-offload host buffer pool

docs/requirements.txt

@@ -1,8 +1,8 @@
-# Sphinx build
-sphinx==9.1.0
 
 # Markdown support
 myst-parser==5.1.0
+# Sphinx build
+sphinx==9.1.0
 
 # Default book theme
 sphinx-book-theme==1.4.0

docs/usage/arguments.md

@@ -557,13 +557,17 @@ distinct from the first emission.
 | activation_gpu_limit | `float` | `0.0` | GB of activations allowed to remain on GPU. |
 | enable_async_activation | `bool` | `False` | Enable async activation offload via stream-based D2H/H2D. Mutually exclusive with `enable_activation`. When `activation_offload_modules` is empty, targets are discovered from `model._no_split_modules`; missing or unmatched model metadata fails closed. |
 | activation_offload_modules | `List[str]` | `[]` | Optional module name patterns for async offload, overriding `_no_split_modules` auto-discovery. Supports segment-aware glob (`model.layers.*` matches direct children only) and `{*}` for sequential groups (`model.layers.{*}`). |
-| activation_offload_host_cache_limit_gb | `float` | `4.0` | Maximum GB of free host buffers retained between steps. In-flight offloads may temporarily exceed this value. Set to `0` to disable reuse. |
+| activation_offload_host_cache_limit_gb | `float` | `4.0` | Idle-cache cap of **one** host-buffer pool, in GB. The trainer applies offload once with this limit, so it is the cap for that call. Each extra `apply_async_activation_offload` given only this limit gets its own pool (caps add); pass the same `host_buffer_pool` to share one cap. Bounds the idle cache only — in-flight offloads may temporarily exceed it. Set to `0` to disable reuse. |
 
 Async activation offload is enabled for CUDA/NPU tensors only; CPU tensors pass
 through unchanged. Only private, dense, contiguous activations are swapped so
-shared-storage views are never resized. Host buffers are pooled per model,
-keyed by shape, stride, and dtype, and evicted by least-recently-used layout to
-enforce `activation_offload_host_cache_limit_gb`. The manager is reset at every training-step
+shared-storage views are never resized. Host buffers are pooled, keyed by shape,
+stride, and dtype, and evicted by least-recently-used layout to enforce the
+pool's `max_cached_bytes`. Passing `host_cache_limit_bytes` (the trainer path)
+builds one pool of that size for that `apply_async_activation_offload` call.
+A caller that applies more than once may pass the same `host_buffer_pool` so
+several schedules share the cap, or omit it so each call owns a pool and the
+caps add. The manager is reset at every training-step
 boundary, including before a step after a failed forward/backward, so stale
 autograd keys cannot affect the next step. The path wraps selected module instances
 and is not intended to be captured by `torch.compile`.

tests/distributed/test_async_offload_unit.py

@@ -264,7 +264,7 @@ def test_wait_d2h_finished_preserves_offset_zero_set_alias():
     assert alias._base is None
     assert alias.storage_offset() == 0
 
-    manager = OffloadManager(host_cache_limit_bytes=1 << 20)
+    manager = OffloadManager(PinnedBufferPool(max_cached_bytes=1 << 20))
     swap = SwapTensor(original, "0_0", manager.host_buffer_pool)
     swap.launch_d2h(manager.swap_stream)
     swap.wait_d2h_finished()
@@ -278,6 +278,52 @@ def test_async_offload_rejects_cpu_tensors():
     assert not base_check_fn(torch.ones(4))
 
 
+def test_shared_host_buffer_pool_survives_repeated_application():
+    """The per-module application path must not multiply the configured limit.
+
+    A composed model applies offload one module at a time, and each call needs
+    its own manager because ``layer_idx`` restarts at 0. The host limit bounds
+    pinned memory for one pool. Sharing that pool across calls keeps one budget;
+    omitting it (passing only a limit) would give each call its own pool.
+    """
+    thinker = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4))
+    talker = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4))
+    pool = PinnedBufferPool(max_cached_bytes=1 << 20)
+
+    apply_async_activation_offload(thinker, ["0", "1"], host_buffer_pool=pool)
+    apply_async_activation_offload(talker, ["0", "1"], host_buffer_pool=pool)
+
+    offloaded = [thinker[0], thinker[1], talker[0], talker[1]]
+    assert len({id(module._veomni_offload_manager) for module in offloaded}) == 2
+    assert all(module._veomni_offload_manager.host_buffer_pool is pool for module in offloaded)
+
+
+def test_apply_async_activation_offload_builds_a_pool_when_given_a_limit():
+    model = nn.Sequential(nn.Linear(4, 4))
+
+    apply_async_activation_offload(model, ["0"], host_cache_limit_bytes=1 << 20)
+
+    assert model[0]._veomni_offload_manager.host_buffer_pool.max_cached_bytes == 1 << 20
+
+
+def test_offload_manager_rejects_a_non_pool_argument():
+    """Fail at wiring time, not from inside the pack hook on the first forward."""
+    with pytest.raises(TypeError, match="PinnedBufferPool"):
+        OffloadManager(1 << 20)
+
+
+def test_apply_async_activation_offload_rejects_a_limit_and_a_pool_together():
+    model = nn.Sequential(nn.Linear(4, 4))
+
+    with pytest.raises(ValueError, match="not both"):
+        apply_async_activation_offload(
+            model,
+            ["0"],
+            host_cache_limit_bytes=1 << 20,
+            host_buffer_pool=PinnedBufferPool(),
+        )
+
+
 def test_pinned_buffer_pool_reuses_matching_layout():
     pool = PinnedBufferPool()
     source = torch.empty((2, 3), dtype=torch.float32)
@@ -350,7 +396,7 @@ def test_async_offload_manager_resets_at_step_boundary():
 def test_unpack_swap_tensor_survives_repeated_access():
     """PyTorch may unpack the same saved tensor twice (``retain_graph=True``)."""
     device = torch.device(get_device_type())
-    manager = OffloadManager(host_cache_limit_bytes=1 << 20)
+    manager = OffloadManager(PinnedBufferPool(max_cached_bytes=1 << 20))
     original = torch.randn(8, 8, device=device)
     expected = original.detach().cpu().clone()
     swap = SwapTensor(original, "0_0", manager.host_buffer_pool)

veomni/arguments/arguments_types.py

@@ -516,8 +516,11 @@ class OffloadConfig:
         default=4.0,
         metadata={
             "help": (
-                "Maximum GB of free host buffers retained by async activation offload between steps. "
-                "In-flight offloads may temporarily use more host memory. Set to 0 to disable reuse."
+                "Maximum GB of free host buffers retained between steps by one host-buffer "
+                "pool. Each `apply_async_activation_offload` call given only this limit owns "
+                "a pool of this size; pass the same `host_buffer_pool` to share one budget "
+                "across calls. Bounds the idle cache only; in-flight offloads may temporarily "
+                "use more. Set to 0 to disable reuse."
             )
         },
     )

veomni/distributed/async_offload.py

@@ -23,7 +23,7 @@
 import fnmatch
 from collections import OrderedDict, deque
 from threading import Lock
-from typing import List
+from typing import List, Optional
 from weakref import WeakValueDictionary
 
 import torch
@@ -43,6 +43,8 @@
 
 logger = logging.get_logger(__name__)
 
+_DEFAULT_HOST_CACHE_LIMIT_BYTES = 4 * 1024**3
+
 
 def _module_name_match(pattern: str, name: str) -> bool:
     """Match module paths without allowing wildcards to cross ``.`` separators."""
@@ -82,7 +84,7 @@ def base_check_fn(tensor) -> bool:
 class PinnedBufferPool:
     """Reuse host buffers with a bounded, layout-level LRU cache."""
 
-    def __init__(self, max_cached_bytes: int = 4 * 1024**3):
+    def __init__(self, max_cached_bytes: int = _DEFAULT_HOST_CACHE_LIMIT_BYTES):
         if max_cached_bytes < 0:
             raise ValueError(f"max_cached_bytes must be non-negative, got {max_cached_bytes}.")
 
@@ -301,14 +303,37 @@ def launch_h2d(self, h2d_stream):
 
 
 class OffloadManager:
-    def __init__(self, host_cache_limit_bytes: int = 4 * 1024**3):
+    """Bookkeeping for one offload schedule, drawing host memory from a shared pool.
+
+    One manager serves exactly one matched set. Keys stay unique across sets --
+    ``GetCnt`` increments rather than resets on a repeated ``block_idx`` -- but
+    ``block_idx`` still means "position within this schedule", and everything
+    that reasons about ordering keys off it: ``wait_d2h_finished`` waits on the
+    previous block by key prefix, ``get_prefetch_keys`` returns the next-lowest
+    block, and ``block_idx == depth - 1`` skips the final block. Two sets both
+    numbering from 0 would prefetch and wait on each other's tensors, and the
+    ``layer_idx == 0`` reset in ``_patch_instance_call`` would fire again at the
+    second set's first layer, mid-step.
+
+    The host pool is the opposite: its buffers are keyed by layout alone and are
+    fungible across schedules. A manager never creates the pool; the caller
+    injects it. One pool can be shared by several managers (one idle-cache
+    budget) or each ``apply`` can supply its own (budgets add).
+    """
+
+    def __init__(self, host_buffer_pool: PinnedBufferPool):
+        if not isinstance(host_buffer_pool, PinnedBufferPool):
+            # Wiring mistakes otherwise surface as an AttributeError raised from
+            # inside the pack hook, several steps away from the bad call.
+            raise TypeError(f"host_buffer_pool must be a PinnedBufferPool, got {type(host_buffer_pool).__name__}.")
+
         # Autograd owns packed SwapTensor objects until backward. Weak references
         # let abandoned graphs release stale entries after a failed step.
         self.items = WeakValueDictionary()
         self.getcnt = GetCnt()
         self.swap_stream = create_stream()
 
-        self.host_buffer_pool = PinnedBufferPool(max_cached_bytes=host_cache_limit_bytes)
+        self.host_buffer_pool = host_buffer_pool
 
     def get_cnt(self, block_idx):
         return self.getcnt.get_cnt(block_idx)
@@ -401,8 +426,11 @@ def reset_async_activation_offload(model):
     The saved-tensor graph normally releases ``SwapTensor`` objects after
     backward.  A failed loss/backward step can leave the graph alive, though,
     so the next training step must explicitly discard its keys and counters.
-    The host pool remains attached to the model and can still reuse buffers
-    once the old graph is collected.
+    Only per-schedule bookkeeping is discarded; the host pool outlives the reset
+    and can still reuse buffers once the old graph is collected.
+
+    Managers are found by walking ``model``, so a caller that offloaded several
+    roots must reset each of them.
     """
     managers = {}
     for module in model.modules():
@@ -520,7 +548,7 @@ def _get_no_split_offload_modules(model):
     return matched_submodules
 
 
-def async_offload_modules(modules, prefetch=True, hidden_states_idx=0, host_cache_limit_bytes: int = 4 * 1024**3):
+def async_offload_modules(modules, *, host_buffer_pool: PinnedBufferPool, prefetch=True, hidden_states_idx=0):
     """Apply async activation offload via selected-instance ``__call__`` patching.
 
     For each module, per-instance attributes (``_veomni_offload_layer_idx``,
@@ -556,8 +584,12 @@ def async_offload_modules(modules, prefetch=True, hidden_states_idx=0, host_cach
     on the second unpack.  The VeOmni fix increments existing counts instead
     of resetting, so the second pass produces ``"0_1"``, ``"1_1"``, … rather
     than repeating ``"0_0"``, ``"1_0"``, ….
+
+    ``host_buffer_pool`` is keyword-only and supplied by the caller, who may
+    share one pool across matched sets or give each set its own; see
+    ``OffloadManager``.
     """
-    manager = OffloadManager(host_cache_limit_bytes=host_cache_limit_bytes)
+    manager = OffloadManager(host_buffer_pool)
     for name, module, layer_idx, depth in modules:
         logger.info_rank0(
             f"Applying activation offload to module: {name}, offload idx: {layer_idx}, offload_layers_num: {depth}"
@@ -638,7 +670,8 @@ def check_fn(x):
 def apply_async_activation_offload(
     model,
     activation_offload_modules: List[str],
-    host_cache_limit_bytes: int = 4 * 1024**3,
+    host_cache_limit_bytes: Optional[int] = None,
+    host_buffer_pool: Optional[PinnedBufferPool] = None,
 ):
     """Apply async activation offload to matched submodules.
 
@@ -651,7 +684,33 @@ def apply_async_activation_offload(
         GC handles intermediate activations via recomputation.
       - Without GC: async offload intercepts hidden_states directly,
         intermediate activations stay on GPU.
+
+    Several patterns in one call share a manager and a pool already, so
+    ``host_cache_limit_bytes`` is all a single apply needs: it builds a pool of
+    that size for this call alone. ``host_buffer_pool`` shares one budget
+    across several apply calls -- a composed model configured per module, or a
+    second model such as a DPO reference. Passing only the limit to each call
+    gives each its own pool and the limits add. The two arguments are mutually
+    exclusive.
+
+    Sharing a pool is not free: its LRU is global and keyed by layout, so
+    schedules with unlike activation shapes can evict each other and fall back
+    to fresh pinned allocations. Scale a shared limit with the number of
+    schedules rather than reusing the single-apply value.
+
+    ``reset_async_activation_offload`` walks one root, so a caller that offloads
+    several roots owes it a call per root.
     """
+    if host_buffer_pool is not None and host_cache_limit_bytes is not None:
+        raise ValueError(
+            "Pass either host_cache_limit_bytes or host_buffer_pool, not both: "
+            "a supplied pool already carries its own limit."
+        )
+    if host_buffer_pool is None:
+        if host_cache_limit_bytes is None:
+            host_cache_limit_bytes = _DEFAULT_HOST_CACHE_LIMIT_BYTES
+        host_buffer_pool = PinnedBufferPool(max_cached_bytes=host_cache_limit_bytes)
+
     if activation_offload_modules:
         matched_modules = get_offload_modules(model, activation_offload_modules)
         if not matched_modules:
@@ -660,4 +719,4 @@ def apply_async_activation_offload(
             )
     else:
         matched_modules = _get_no_split_offload_modules(model)
-    async_offload_modules(matched_modules, host_cache_limit_bytes=host_cache_limit_bytes)
+    async_offload_modules(matched_modules, host_buffer_pool=host_buffer_pool)