diff --git a/docs/requirements.txt b/docs/requirements.txt index 65ac6f6029..c18a55bf19 100644 --- a/docs/requirements.txt +++ b/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 diff --git a/docs/usage/arguments.md b/docs/usage/arguments.md index 8e09751211..6b7cbf95b2 100644 --- a/docs/usage/arguments.md +++ b/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`. diff --git a/tests/distributed/test_async_offload_unit.py b/tests/distributed/test_async_offload_unit.py index db57af2311..5949ed172d 100644 --- a/tests/distributed/test_async_offload_unit.py +++ b/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) diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index 7637c098b1..974ac1397a 100644 --- a/veomni/arguments/arguments_types.py +++ b/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." ) }, ) diff --git a/veomni/distributed/async_offload.py b/veomni/distributed/async_offload.py index 38ada68e8d..069bd77afc 100644 --- a/veomni/distributed/async_offload.py +++ b/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)