diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index d095f374a9..fc4234e730 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -15,6 +15,7 @@ veomni/ │ ├── parallel_state.py init_parallel_state(), ParallelState, device mesh setup │ ├── torch_parallelize.py build_parallelize_model(), parallelize_model_fsdp2() │ ├── parallel_plan.py ParallelPlan for ExtraParallel (EP, embedding shard) +│ ├── async_offload.py Async activation offload (SwapTensor, OffloadManager, async_save_on_cpu) │ ├── fsdp2/ FSDP2 (composable fully_shard), gradient clipping │ ├── moe/ MoE expert parallelism: token routing, all-to-all, EPGroupGemm │ └── sequence_parallel/ Ulysses SP: all-to-all head/seq exchange, async variants diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml index d70d6521c3..7e8f1fc1e3 100644 --- a/.github/workflows/gpu_unit_tests.yml +++ b/.github/workflows/gpu_unit_tests.yml @@ -173,6 +173,7 @@ jobs: uv run --frozen pytest -s -x tests/checkpoints/test_trainer_saveload.py - name: Run distributed tests run: | + uv run --frozen pytest -s -x tests/distributed/test_async_offload_unit.py uv run --frozen pytest -s -x tests/distributed/test_dummy_forward.py uv run --frozen pytest -s -x tests/distributed/test_torch_compile.py - name: Run data tests diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml index 046e38d068..e8d6e91baa 100644 --- a/.github/workflows/npu_unit_tests.yml +++ b/.github/workflows/npu_unit_tests.yml @@ -182,6 +182,7 @@ jobs: uv run --frozen pytest -v -s -x tests/checkpoints/test_trainer_saveload.py - name: Run distributed tests (dummy_forward) run: | + uv run --frozen pytest -v -s -x tests/distributed/test_async_offload_unit.py uv run --frozen pytest -v -s -x tests/distributed/test_dummy_forward.py - name: Run data tests run: | diff --git a/docs/usage/arguments.md b/docs/usage/arguments.md index b3d1e08ee1..dadd2e256d 100644 --- a/docs/usage/arguments.md +++ b/docs/usage/arguments.md @@ -439,8 +439,20 @@ distinct from the first emission. | Field | Type | Default | Description | | --- | --- | --- | --- | -| enable_activation | `bool` | `False` | Enable activation offload to CPU. | +| enable_activation | `bool` | `False` | Enable synchronous activation offload to CPU. | | 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. | + +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 +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`. ### CheckpointConfig diff --git a/tests/distributed/test_async_offload_unit.py b/tests/distributed/test_async_offload_unit.py new file mode 100644 index 0000000000..db57af2311 --- /dev/null +++ b/tests/distributed/test_async_offload_unit.py @@ -0,0 +1,534 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for MindSpeed-style async activation offload helpers.""" + +from functools import partial +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint +from transformers.modeling_layers import GradientCheckpointingLayer + +from veomni.arguments.arguments_types import OffloadConfig +from veomni.distributed.async_offload import ( + GetCnt, + OffloadManager, + PinnedBufferPool, + SwapTensor, + _has_private_dense_storage, + _unpack_swap_tensor, + apply_async_activation_offload, + base_check_fn, + get_offload_modules, + reset_async_activation_offload, +) +from veomni.utils.device import IS_CUDA_AVAILABLE, IS_NPU_AVAILABLE, get_device_type, get_torch_device + + +class _AsyncOffloadFSDPBlock(GradientCheckpointingLayer): + def __init__(self, device=None): + super().__init__() + self.proj = nn.Linear(16, 16, device=device) + + def forward(self, hidden_states): + return torch.nn.functional.gelu(self.proj(hidden_states)) + + +class _AsyncOffloadFSDPModel(nn.Module): + _no_split_modules = ["_AsyncOffloadFSDPBlock"] + + def __init__(self, device=None): + super().__init__() + self.layers = nn.ModuleList([_AsyncOffloadFSDPBlock(device=device) for _ in range(4)]) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + checkpoint_func = partial(checkpoint, **(gradient_checkpointing_kwargs or {})) + for layer in self.layers: + layer.gradient_checkpointing = True + layer._gradient_checkpointing_func = checkpoint_func + + def init_weights(self): + for module in self.modules(): + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +def test_offload_config_allows_no_modules_for_auto_discovery(): + cfg = OffloadConfig(enable_async_activation=True) + assert cfg.enable_async_activation is True + assert cfg.activation_offload_modules == [] + + +def test_offload_config_rejects_sync_and_async_activation_together(): + with pytest.raises(ValueError, match="mutually exclusive"): + OffloadConfig( + enable_activation=True, + enable_async_activation=True, + activation_offload_modules=["model.layers.{*}"], + ) + + +@pytest.mark.parametrize("cache_limit", [-1.0, float("inf"), float("nan"), float.fromhex("0x1.fffffffffffffp+1023")]) +def test_offload_config_rejects_invalid_host_cache_limit(cache_limit): + with pytest.raises(ValueError, match="activation_offload_host_cache_limit_gb"): + OffloadConfig(activation_offload_host_cache_limit_gb=cache_limit) + + +def test_get_offload_modules_glob_is_segment_aware(): + class Toy(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + + matched = get_offload_modules(Toy(), ["model.layers.*"]) + + assert [item[0] for item in matched] == ["model.layers.0", "model.layers.1"] + + +def test_get_cnt_unique_keys_across_second_pass(): + """Second forward over the same layer indices must not collide keys.""" + cnt = GetCnt() + first = [cnt.get_cnt(i)[0] for i in range(3)] + second = [cnt.get_cnt(i)[0] for i in range(3)] + assert first == ["0_0", "1_0", "2_0"] + assert second == ["0_1", "1_1", "2_1"] + assert set(first).isdisjoint(second) + + cnt = GetCnt() + assert [cnt.get_cnt(i)[1] for i in (0, 5, 2)] == [None, 0, 5] + + +def test_get_prefetch_keys_use_previous_offloaded_layer_and_its_tensor_count(): + cnt = GetCnt() + cnt.get_cnt(2) + cnt.get_cnt(2) + cnt.get_cnt(5) + + assert cnt.get_prefetch_keys(5) == ["2_0", "2_1"] + + +def test_get_offload_modules_brace_star_expands_sequential(): + class Toy(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]) + + model = Toy() + matched = get_offload_modules(model, ["model.layers.{*}"]) + names = [item[0] for item in matched] + assert names == ["model.layers.0", "model.layers.1", "model.layers.2"] + # depth field is rewritten to total offload layer count + assert all(item[-1] == 3 for item in matched) + assert [item[2] for item in matched] == [0, 1, 2] + + +def test_apply_async_activation_offload_rejects_unmatched_patterns(): + model = nn.Sequential(nn.Linear(4, 4)) + + with pytest.raises(ValueError, match="did not match any model modules"): + apply_async_activation_offload(model, ["model.layers.{*}"]) + + +def test_apply_async_activation_offload_auto_discovers_no_split_modules(): + class DecoderLayer(nn.Module): + pass + + class Toy(nn.Module): + _no_split_modules = ["DecoderLayer"] + + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([DecoderLayer(), DecoderLayer()]) + + model = Toy() + apply_async_activation_offload(model, []) + + assert [layer._veomni_offload_layer_idx for layer in model.layers] == [0, 1] + # FSDP2 and torch.compile discover blocks by exact class name after this + # hook is installed, so the private wrapper must preserve that identity. + assert [layer.__class__.__name__ for layer in model.layers] == ["DecoderLayer", "DecoderLayer"] + assert len([module for module in model.modules() if module.__class__.__name__ in model._no_split_modules]) == 2 + assert all(layer._veomni_offload_depth == 2 for layer in model.layers) + assert model.layers[0]._veomni_offload_manager is model.layers[1]._veomni_offload_manager + + +@pytest.mark.parametrize( + "no_split_modules, error_match", + [ + (None, "requires model._no_split_modules"), + (["MissingDecoderLayer"], "did not match any modules"), + ], +) +def test_apply_async_activation_offload_auto_discovery_fails_closed(no_split_modules, error_match): + model = nn.Sequential(nn.Linear(4, 4)) + if no_split_modules is not None: + model._no_split_modules = no_split_modules + + with pytest.raises(ValueError, match=error_match): + apply_async_activation_offload(model, []) + + +def test_async_offload_patch_is_confined_to_selected_instances_and_models(): + class DecoderLayer(nn.Module): + def forward(self, hidden_states): + return hidden_states * 2 + + class Toy(nn.Module): + def __init__(self): + super().__init__() + self.selected = DecoderLayer() + self.unselected = DecoderLayer() + + first_model = Toy() + second_model = Toy() + apply_async_activation_offload(first_model, ["selected"]) + apply_async_activation_offload(second_model, ["selected"]) + + assert type(first_model.selected) is not DecoderLayer + assert type(first_model.unselected) is DecoderLayer + assert type(second_model.unselected) is DecoderLayer + assert first_model.selected._veomni_offload_manager is not second_model.selected._veomni_offload_manager + torch.testing.assert_close(first_model.selected(torch.ones(2)), torch.full((2,), 2.0)) + + +def test_async_offload_manager_resets_after_forward_error(): + class FailingLayer(nn.Module): + def forward(self, hidden_states): + raise RuntimeError("expected failure") + + model = nn.Sequential(FailingLayer()) + apply_async_activation_offload(model, ["0"]) + manager = model[0]._veomni_offload_manager + manager.get_cnt(7) + + with pytest.raises(RuntimeError, match="expected failure"): + model(torch.ones(2)) + + assert not manager.items + assert manager.getcnt._block_idx == -1 + assert manager.getcnt._block_tensor_nums == {} + + +def test_async_offload_rejects_tensor_views_with_shared_storage(): + base = torch.arange(8.0) + view = base[2:6] + + assert not _has_private_dense_storage(view) + assert not base_check_fn(view) + torch.testing.assert_close(base, torch.arange(8.0)) + + +def test_private_dense_storage_allows_padded_contiguous_owner(): + tensor = torch.arange(4.0) + tensor.untyped_storage().resize_(tensor.untyped_storage().nbytes() * 2) + + assert _has_private_dense_storage(tensor) + alias = tensor.as_strided((2,), (1,), storage_offset=2) + assert not _has_private_dense_storage(alias) + + +@pytest.mark.skipif( + not (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE), + reason="CUDA or NPU is required for async D2H/H2D validation", +) +def test_wait_d2h_finished_preserves_offset_zero_set_alias(): + device = torch.device(get_device_type()) + original = torch.randn(8, 8, device=device) + original.untyped_storage().resize_(original.untyped_storage().nbytes() * 2) + expected = original.detach().cpu().clone() + alias = torch.empty_like(original) + alias.set_(original.untyped_storage(), 0, original.size(), original.stride()) + assert alias._base is None + assert alias.storage_offset() == 0 + + manager = OffloadManager(host_cache_limit_bytes=1 << 20) + swap = SwapTensor(original, "0_0", manager.host_buffer_pool) + swap.launch_d2h(manager.swap_stream) + swap.wait_d2h_finished() + + torch.testing.assert_close(alias.cpu(), expected) + restored = _unpack_swap_tensor(manager, swap, prefetch=False) + torch.testing.assert_close(restored.cpu(), expected) + + +def test_async_offload_rejects_cpu_tensors(): + assert not base_check_fn(torch.ones(4)) + + +def test_pinned_buffer_pool_reuses_matching_layout(): + pool = PinnedBufferPool() + source = torch.empty((2, 3), dtype=torch.float32) + first, key = pool.acquire(source) + pool.release(first, key) + second, second_key = pool.acquire(source) + + assert first.data_ptr() == second.data_ptr() + assert key == second_key + assert pool.allocations == 1 + assert pool.reuses == 1 + + +def test_pinned_buffer_pool_cache_is_bounded_across_dynamic_shapes(): + pool = PinnedBufferPool(max_cached_bytes=64) + + for width in (4, 8, 12): + source = torch.empty((width,), dtype=torch.float32) + buffer, key = pool.acquire(source) + pool.release(buffer, key) + + assert pool.cached_bytes <= 64 + assert pool.evictions > 0 + + +def test_pinned_buffer_pool_refreshes_layout_recency_on_reuse(): + pool = PinnedBufferPool(max_cached_bytes=48) + layouts = [ + torch.empty((4,), dtype=torch.float32), + torch.empty((2, 2), dtype=torch.float32), + torch.empty((1, 4), dtype=torch.float32), + torch.empty((4, 1), dtype=torch.float32), + ] + + first_key_buffers = [pool.acquire(layouts[0]) for _ in range(2)] + for buffer, key in first_key_buffers: + pool.release(buffer, key) + second_buffer, second_key = pool.acquire(layouts[1]) + pool.release(second_buffer, second_key) + + reused_buffer, reused_key = pool.acquire(layouts[0]) + pool.release(reused_buffer, reused_key) + for layout in layouts[2:]: + buffer, key = pool.acquire(layout) + pool.release(buffer, key) + + allocations_before = pool.allocations + pool.acquire(layouts[1]) + + assert pool.allocations == allocations_before + 1 + + +def test_async_offload_manager_resets_at_step_boundary(): + model = nn.Sequential(nn.Linear(2, 2)) + apply_async_activation_offload(model, ["0"]) + manager = model[0]._veomni_offload_manager + manager.get_cnt(9) + + reset_async_activation_offload(model) + + assert not manager.items + assert manager.getcnt._block_idx == -1 + assert manager.getcnt._block_tensor_nums == {} + + +@pytest.mark.skipif( + not (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE), + reason="CUDA or NPU is required for async D2H/H2D validation", +) +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) + original = torch.randn(8, 8, device=device) + expected = original.detach().cpu().clone() + swap = SwapTensor(original, "0_0", manager.host_buffer_pool) + swap.launch_d2h(manager.swap_stream) + swap.wait_d2h_finished() + manager.put("0_0", swap) + + first = _unpack_swap_tensor(manager, swap, prefetch=False) + assert not manager.exist("0_0") + torch.testing.assert_close(first.cpu(), expected) + + second = _unpack_swap_tensor(manager, swap, prefetch=False) + torch.testing.assert_close(second.cpu(), expected) + assert second.data_ptr() == first.data_ptr() + + +@pytest.mark.skipif( + not (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE), + reason="CUDA or NPU is required for async D2H/H2D validation", +) +def test_async_offload_accelerator_gradient_parity_and_peak_memory(): + device_api = get_torch_device() + device = torch.device(get_device_type()) + device_api.synchronize() + device_api.empty_cache() + + class Block(GradientCheckpointingLayer): + def __init__(self): + super().__init__() + # Same nesting as production: patched ``__call__`` (async pack) wraps + # ``GradientCheckpointingLayer.__call__`` (checkpoint). NPU Linear + # often saves a tensor whose ``data_ptr`` is not the module input, so + # a checkpoint-less Sequential never packs (allocations == 0). + self.gradient_checkpointing = True + self._gradient_checkpointing_func = partial(checkpoint, use_reentrant=False) + self.proj = nn.Linear(512, 512) + + def forward(self, hidden_states): + return torch.nn.functional.gelu(self.proj(hidden_states)) + + def make_model(): + model = nn.Sequential(*(Block() for _ in range(4))).to(device) + model.train() + return model + + torch.manual_seed(0) + baseline = make_model() + initial_state = {name: parameter.detach().cpu().clone() for name, parameter in baseline.state_dict().items()} + input_cpu = torch.randn(4, 128, 512) + + device_api.reset_peak_memory_stats() + baseline_before = device_api.memory_allocated() + baseline_input = input_cpu.to(device).requires_grad_(True) + baseline_output = baseline(baseline_input) + baseline_output.square().mean().backward() + device_api.synchronize() + baseline_peak = device_api.max_memory_allocated() - baseline_before + + baseline_output_cpu = baseline_output.detach().cpu() + baseline_input_grad_cpu = baseline_input.grad.detach().cpu() + baseline_parameter_grads = [parameter.grad.detach().cpu() for parameter in baseline.parameters()] + + del baseline_output, baseline_input, baseline + device_api.empty_cache() + + offloaded = make_model() + offloaded.load_state_dict(initial_state) + apply_async_activation_offload(offloaded, ["0", "1", "2", "3"]) + + device_api.reset_peak_memory_stats() + offloaded_before = device_api.memory_allocated() + offloaded_input = input_cpu.to(device).requires_grad_(True) + offloaded_output = offloaded(offloaded_input) + offloaded_output.square().mean().backward() + device_api.synchronize() + offloaded_peak = device_api.max_memory_allocated() - offloaded_before + + manager = offloaded[0]._veomni_offload_manager + assert not manager.items + assert manager.host_buffer_pool.allocations > 0 + + torch.testing.assert_close(offloaded_output.detach().cpu(), baseline_output_cpu) + torch.testing.assert_close(offloaded_input.grad.detach().cpu(), baseline_input_grad_cpu) + for offloaded_parameter, baseline_parameter_grad in zip(offloaded.parameters(), baseline_parameter_grads): + torch.testing.assert_close(offloaded_parameter.grad.detach().cpu(), baseline_parameter_grad) + + # This is a numerical and memory gate, not only an execution smoke test: + # offloading must preserve gradients and must not increase peak allocated + # device memory for the same model and input. + assert offloaded_peak <= baseline_peak, ( + f"async offload increased peak allocated memory: baseline={baseline_peak}, offloaded={offloaded_peak}" + ) + + del offloaded_output, offloaded_input + offloaded.zero_grad(set_to_none=True) + second_input = input_cpu.to(device).requires_grad_(True) + second_output = offloaded(second_input) + second_output.square().mean().backward() + device_api.synchronize() + assert manager.host_buffer_pool.reuses > 0 + + +def _run_async_offload_base_trainer_fsdp2_gc(): + import torch.distributed as dist + + from veomni.arguments import ( + FSDPConfig, + GradientCheckpointingConfig, + MixedPrecisionConfig, + OptimizerConfig, + TorchCompileConfig, + ) + from veomni.distributed.parallel_state import init_parallel_state, use_parallel_state + from veomni.trainer.base import BaseTrainer + + world_size = dist.get_world_size() + init_parallel_state( + dp_size=world_size, + dp_shard_size=world_size, + dp_mode="fsdp2", + device_type=get_device_type(), + name="base", + ) + + trainer = object.__new__(BaseTrainer) + trainer.model = _AsyncOffloadFSDPModel(device="meta") + trainer.args = SimpleNamespace( + model=SimpleNamespace( + lora_config=None, + fqn_to_index_mapping=None, + model_path=None, + basic_modules=[], + ), + train=SimpleNamespace( + accelerator=SimpleNamespace( + offload_config=OffloadConfig( + enable_async_activation=True, + activation_offload_host_cache_limit_gb=0.01, + ), + fsdp_config=FSDPConfig(mixed_precision=MixedPrecisionConfig(enable=False)), + ), + optimizer=OptimizerConfig(), + checkpoint=SimpleNamespace(load_path=None), + init_device="meta", + gradient_checkpointing=GradientCheckpointingConfig(enable=True), + broadcast_model_weights_from_rank0=False, + ep_sharded_stream_load=False, + torch_compile=TorchCompileConfig(enable=False), + ), + ) + + torch.manual_seed(0) + with use_parallel_state("base"): + trainer._build_parallelized_model() + + assert all(layer.gradient_checkpointing for layer in trainer.model.layers) + manager = trainer.model.layers[0]._veomni_offload_manager + hidden_states = torch.randn(2, 8, 16, device=get_device_type(), requires_grad=True) + with use_parallel_state("base"): + output = trainer.model(hidden_states) + with use_parallel_state("base"): + output.float().square().mean().backward() + + assert torch.isfinite(output).all() + assert manager.host_buffer_pool.allocations > 0 + assert not manager.items + for parameter in trainer.model.parameters(): + if parameter.grad is not None: + local_grad = parameter.grad.to_local() if hasattr(parameter.grad, "to_local") else parameter.grad + assert torch.isfinite(local_grad).all() + dist.barrier() + + +@pytest.mark.skipif( + not (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE), + reason="CUDA or NPU is required for FSDP2 async-offload integration", +) +def test_async_offload_base_trainer_fsdp2_gc_integration(): + from ..tools.launch_utils import torchrun + + torchrun(_run_async_offload_base_trainer_fsdp2_gc, world_size=2) diff --git a/tests/distributed/test_torch_compile.py b/tests/distributed/test_torch_compile.py index eeeb42d811..3064f8ad06 100644 --- a/tests/distributed/test_torch_compile.py +++ b/tests/distributed/test_torch_compile.py @@ -12,6 +12,7 @@ from veomni.arguments.arguments_types import ( DataArguments, ModelArguments, + OffloadConfig, OpsImplementationConfig, TrainingArguments, VeOmniArguments, @@ -421,6 +422,7 @@ def test_mark_compile_step_begin_skips_without_torch_compiler(monkeypatch): def test_vlm_train_step_marks_each_compile_micro_batch(monkeypatch): + from veomni.trainer.base import BaseTrainer from veomni.trainer.vlm_trainer import VLMTrainer marks = [] @@ -432,7 +434,12 @@ def test_vlm_train_step_marks_each_compile_micro_batch(monkeypatch): trainer = VLMTrainer.__new__(VLMTrainer) trainer.base = SimpleNamespace( - args=SimpleNamespace(train=SimpleNamespace(optimizer=SimpleNamespace(max_grad_norm=1.0))), + args=SimpleNamespace( + train=SimpleNamespace( + optimizer=SimpleNamespace(max_grad_norm=1.0), + accelerator=SimpleNamespace(offload_config=OffloadConfig()), + ) + ), state=SimpleNamespace(global_step=0), model=SimpleNamespace(_veomni_compile_uses_cuda_graphs=True), model_reshard=lambda *_: None, @@ -444,6 +451,9 @@ def test_vlm_train_step_marks_each_compile_micro_batch(monkeypatch): on_step_begin=lambda **_: None, on_step_end=lambda **_: None, ) + trainer.base._reset_async_activation_offload_if_enabled = ( + lambda: BaseTrainer._reset_async_activation_offload_if_enabled(trainer.base) + ) trainer.train_step(iter([[{}, {}]])) diff --git a/tests/trainer/test_aux_metrics.py b/tests/trainer/test_aux_metrics.py index 5d11b67168..e689ef7239 100644 --- a/tests/trainer/test_aux_metrics.py +++ b/tests/trainer/test_aux_metrics.py @@ -39,6 +39,7 @@ import veomni.trainer.callbacks.trace_callback as trace_callback_module import veomni.trainer.text_trainer as text_trainer_module import veomni.trainer.vlm_trainer as vlm_trainer_module +from veomni.arguments.arguments_types import OffloadConfig from veomni.trainer.base import BaseTrainer from veomni.trainer.callbacks.base import TrainerState from veomni.trainer.callbacks.trace_callback import RESERVED_TRAINING_METRIC_NAMES @@ -191,6 +192,7 @@ def _accumulating_trainer(outputs, recorded): accelerator=SimpleNamespace( dp_replicate_size=1, fsdp_config=SimpleNamespace(fsdp_mode="fsdp2", reshard_after_backward=True), + offload_config=OffloadConfig(), ), sync_each_train_step=False, ) diff --git a/tests/trainer/test_step_sync.py b/tests/trainer/test_step_sync.py index 2c1f43df37..27d7b1ec00 100644 --- a/tests/trainer/test_step_sync.py +++ b/tests/trainer/test_step_sync.py @@ -1,8 +1,9 @@ import inspect +from contextlib import nullcontext from types import SimpleNamespace import veomni.trainer.base as base_module -from veomni.trainer.base import BaseTrainer +from veomni.trainer.base import BaseTrainer, _resolve_offload_config from veomni.trainer.dit_trainer import DiTTrainer from veomni.trainer.text_dpo_trainer import TextDPOTrainer from veomni.trainer.text_trainer import TextTrainer @@ -33,6 +34,47 @@ def test_train_step_uses_sync_helper(): assert "synchronize()" not in source +def test_reset_async_activation_offload_skips_missing_config(monkeypatch): + calls = [] + monkeypatch.setattr(base_module, "reset_async_activation_offload", lambda model: calls.append(model)) + + trainer = BaseTrainer.__new__(BaseTrainer) + trainer.model = object() + trainer.args = SimpleNamespace(train=SimpleNamespace(accelerator=SimpleNamespace())) + + trainer._reset_async_activation_offload_if_enabled() + assert calls == [] + + trainer.args.train.accelerator.offload_config = SimpleNamespace(enable_async_activation=True) + trainer._reset_async_activation_offload_if_enabled() + assert calls == [trainer.model] + + +def test_resolve_offload_config_defaults_when_absent(): + cfg = _resolve_offload_config(SimpleNamespace(train=SimpleNamespace(accelerator=SimpleNamespace()))) + + assert cfg.enable_async_activation is False + assert cfg.enable_activation is False + + +def test_build_training_context_without_offload_config(monkeypatch): + contexts = (nullcontext(), nullcontext()) + monkeypatch.setattr(base_module, "build_activation_offloading_context", lambda *args, **kwargs: contexts) + + trainer = BaseTrainer.__new__(BaseTrainer) + trainer.args = SimpleNamespace( + train=SimpleNamespace( + accelerator=SimpleNamespace(), + gradient_checkpointing=SimpleNamespace(enable=False), + ) + ) + + trainer._build_training_context() + + assert trainer.model_fwd_context is contexts[0] + assert trainer.model_bwd_context is contexts[1] + + def test_configure_hsdp_allreduce_toggles_outer_micro_steps(): calls = [] trainer = BaseTrainer.__new__(BaseTrainer) @@ -52,6 +94,13 @@ def test_configure_hsdp_allreduce_toggles_outer_micro_steps(): assert calls == [False, True] +def test_train_step_uses_async_offload_reset_helper(): + for wrapper_cls in (BaseTrainer, TextTrainer, VLMTrainer, TextDPOTrainer, DiTTrainer): + source = inspect.getsource(wrapper_cls.train_step) + + assert "_reset_async_activation_offload_if_enabled()" in source + + def test_train_step_uses_hsdp_allreduce_helper(): for wrapper_cls in (BaseTrainer, TextTrainer, VLMTrainer, TextDPOTrainer, DiTTrainer): source = inspect.getsource(wrapper_cls.train_step) diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index 57ba0e3fe1..9daf7209f4 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -14,6 +14,7 @@ import math import os +import sys from dataclasses import MISSING, dataclass, field, fields from pathlib import Path from typing import Dict, List, Literal, Optional, Tuple @@ -24,6 +25,8 @@ logger = logging.get_logger(__name__) +_MAX_HOST_CACHE_LIMIT_GB = sys.maxsize // 1024**3 + def _resolve_hdfs_path(path: Optional[str]) -> Optional[str]: """Copy an ``hdfs://`` path to a local cache and return the local path. @@ -478,7 +481,7 @@ class OffloadConfig: enable_activation: bool = field( default=False, - metadata={"help": "Enable activation offload to CPU."}, + metadata={"help": "Enable synchronous activation offload to CPU."}, ) activation_gpu_limit: float = field( default=0.0, @@ -486,6 +489,55 @@ class OffloadConfig: "help": "When enabling activation offload, `activation_gpu_limit` GB activations are allowed to reserve on GPU." }, ) + enable_async_activation: bool = field( + default=False, + metadata={ + "help": ( + "Enable async activation offload to CPU via stream-based D2H/H2D transfers. " + "More efficient than synchronous offload. Mutually exclusive with " + "`enable_activation`. Uses `model._no_split_modules` when " + "`activation_offload_modules` is empty." + ) + }, + ) + activation_offload_modules: List[str] = field( + default_factory=list, + metadata={ + "help": ( + "Module name patterns for async activation offload. Supports glob patterns " + "(e.g. `model.layers.*`) and the special `{*}` wildcard for sequential " + "module groups (e.g. `model.layers.{*}` expands to all decoder layers)." + ) + }, + ) + activation_offload_host_cache_limit_gb: float = field( + 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." + ) + }, + ) + + def __post_init__(self): + if self.enable_activation and self.enable_async_activation: + raise ValueError( + "enable_activation and enable_async_activation are mutually exclusive; " + "select exactly one activation offload mode." + ) + 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, " + f"got {self.activation_offload_host_cache_limit_gb}." + ) + if self.activation_offload_host_cache_limit_gb > _MAX_HOST_CACHE_LIMIT_GB: + raise ValueError( + "activation_offload_host_cache_limit_gb is too large to convert to bytes: " + f"got {self.activation_offload_host_cache_limit_gb}, maximum {_MAX_HOST_CACHE_LIMIT_GB}." + ) @dataclass diff --git a/veomni/distributed/async_offload.py b/veomni/distributed/async_offload.py new file mode 100644 index 0000000000..c0a8a644ce --- /dev/null +++ b/veomni/distributed/async_offload.py @@ -0,0 +1,663 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# Adapted from MindSpeed-MM's async_offload.py. +# Key adaptations for VeOmni: +# - Uses VeOmni's device-agnostic stream/event helpers (veomni.utils.device) +# - Uses VeOmni's logger (veomni.utils.logging) instead of print_rank +# - Uses segment-aware module-name matching instead of mindspeed's str_match +# - Supports wildcard pattern ``{*}`` for sequential module groups (e.g. model.layers.{*}) + +import fnmatch +from collections import OrderedDict, deque +from threading import Lock +from typing import List +from weakref import WeakValueDictionary + +import torch +from torch.autograd.graph import saved_tensors_hooks + +from ..utils import logging +from ..utils.device import ( + IS_CUDA_AVAILABLE, + IS_NPU_AVAILABLE, + create_event, + create_stream, + get_current_stream, + get_device_type, + switch_to_specified_stream, +) + + +logger = logging.get_logger(__name__) + + +def _module_name_match(pattern: str, name: str) -> bool: + """Match module paths without allowing wildcards to cross ``.`` separators.""" + pattern_segments = pattern.split(".") + name_segments = name.split(".") + return len(pattern_segments) == len(name_segments) and all( + fnmatch.fnmatchcase(name_segment, pattern_segment) + for pattern_segment, name_segment in zip(pattern_segments, name_segments) + ) + + +def _is_active_accelerator_tensor(tensor) -> bool: + return (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE) and tensor.device.type == get_device_type() + + +def _has_private_dense_storage(tensor) -> bool: + """Return whether ``tensor`` can be swapped without mutating another live view. + + Expandable CUDA segments can make ``storage.nbytes()`` larger than the + logical tensor, and may also give a non-zero ``storage_offset``. ``SwapTensor`` + unbinds via ``set_`` rather than ``resize_(0)``, so padding and a non-zero + offset are safe when this tensor is not a view (``_base is None``). + """ + if isinstance(tensor, torch.nn.parameter.Parameter) or tensor._base is not None: + return False + if not tensor.is_contiguous(): + return False + storage_nbytes = tensor.untyped_storage().nbytes() + tensor_nbytes = tensor.numel() * tensor.element_size() + return storage_nbytes > 0 and storage_nbytes >= tensor_nbytes + tensor.storage_offset() * tensor.element_size() + + +def base_check_fn(tensor) -> bool: + return _is_active_accelerator_tensor(tensor) and _has_private_dense_storage(tensor) + + +class PinnedBufferPool: + """Reuse host buffers with a bounded, layout-level LRU cache.""" + + def __init__(self, max_cached_bytes: int = 4 * 1024**3): + if max_cached_bytes < 0: + raise ValueError(f"max_cached_bytes must be non-negative, got {max_cached_bytes}.") + + self.max_cached_bytes = int(max_cached_bytes) + self._free_buffers = OrderedDict() + self._cached_bytes = 0 + self._lock = Lock() + self.allocations = 0 + self.reuses = 0 + self.evictions = 0 + self.drops = 0 + + @property + def cached_bytes(self): + return self._cached_bytes + + @staticmethod + def _key(tensor): + # Strides are part of the key even though the current offload + # predicate accepts contiguous tensors only; this keeps the pool safe + # if another dense layout is supported later. + return (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) + + def acquire(self, tensor): + key = self._key(tensor) + with self._lock: + free_buffers = self._free_buffers.get(key) + if free_buffers: + buffer = free_buffers.pop() + self._cached_bytes -= buffer.numel() * buffer.element_size() + if free_buffers: + self._free_buffers.move_to_end(key) + else: + del self._free_buffers[key] + self.reuses += 1 + return buffer, key + + pin_memory = _is_active_accelerator_tensor(tensor) + try: + buffer = torch.empty_strided( + tensor.shape, + tensor.stride(), + dtype=tensor.dtype, + device="cpu", + pin_memory=pin_memory, + ) + except RuntimeError: + if not pin_memory: + raise + logger.warning_rank0( + "Pinned host allocation for async activation offload failed; falling back to pageable host memory." + ) + buffer = torch.empty_strided( + tensor.shape, + tensor.stride(), + dtype=tensor.dtype, + device="cpu", + pin_memory=False, + ) + with self._lock: + self.allocations += 1 + return buffer, key + + def release(self, buffer, key): + if buffer is None: + return + + buffer_nbytes = buffer.numel() * buffer.element_size() + with self._lock: + if self.max_cached_bytes == 0 or buffer_nbytes > self.max_cached_bytes: + self.drops += 1 + return + + while self._cached_bytes + buffer_nbytes > self.max_cached_bytes: + oldest_key, oldest_buffers = next(iter(self._free_buffers.items())) + evicted = oldest_buffers.popleft() + self._cached_bytes -= evicted.numel() * evicted.element_size() + self.evictions += 1 + if not oldest_buffers: + del self._free_buffers[oldest_key] + + free_buffers = self._free_buffers.setdefault(key, deque()) + free_buffers.append(buffer) + self._free_buffers.move_to_end(key) + self._cached_bytes += buffer_nbytes + + +class GetCnt: + def __init__(self): + self._block_idx = -1 + self._block_tensor_nums = {} + + def get_cnt(self, block_idx): + previous_block_idx = self._block_idx if self._block_idx >= 0 and block_idx != self._block_idx else None + if block_idx > self._block_idx: + if block_idx in self._block_tensor_nums: + self._block_tensor_nums[block_idx] += 1 + else: + self._block_tensor_nums[block_idx] = 1 + self._block_idx = block_idx + elif block_idx == self._block_idx: + self._block_tensor_nums[block_idx] += 1 + else: + if block_idx in self._block_tensor_nums: + self._block_tensor_nums[block_idx] += 1 + else: + self._block_tensor_nums[block_idx] = 1 + self._block_idx = block_idx + + offload_tensor_key = "{}_{}".format(self._block_idx, self._block_tensor_nums[self._block_idx] - 1) + return offload_tensor_key, previous_block_idx + + def reset(self): + self._block_idx = -1 + self._block_tensor_nums = {} + + def get_prefetch_keys(self, block_idx): + prefetch_block_idx = max((idx for idx in self._block_tensor_nums.keys() if idx < block_idx), default=None) + + if prefetch_block_idx is None: + return [] + + block_tensor_nums = self._block_tensor_nums[prefetch_block_idx] + prefetch_idxs = list(range(0, block_tensor_nums)) + return ["{}_{}".format(prefetch_block_idx, prefetch_idx) for prefetch_idx in prefetch_idxs] + + def get_layer_tensor_nums(self, block_idx): + return self._block_tensor_nums[block_idx] + + +class SwapTensor: + def __init__(self, tensor, key, pool): + if not _has_private_dense_storage(tensor): + raise ValueError("Async activation offload requires a tensor with private, dense storage.") + + self.tensor = tensor + self.size = tensor.size() + self.stride = tensor.stride() + self.pool = pool + self.tensor_cpu, self._host_buffer_key = pool.acquire(tensor) + self._host_buffer_released = False + + self.stat = "device" + self.key = key + + self.d2h_event = create_event() + self.h2d_event = create_event() + + def release_host_buffer(self): + if self._host_buffer_released or self.tensor_cpu is None: + return + if self.stat == "d2h": + self.d2h_event.synchronize() + elif self.stat == "host": + self.d2h_event.synchronize() + elif self.stat == "h2d": + self.h2d_event.synchronize() + self.pool.release(self.tensor_cpu, self._host_buffer_key) + self.tensor_cpu = None + self._host_buffer_released = True + + def __del__(self): + try: + self.release_host_buffer() + except Exception: + pass + + def launch_d2h(self, stream): + if self.stat != "device": + return + + forward_event = create_event() + forward_event.record() + with torch.no_grad(): + with switch_to_specified_stream(stream): + stream.wait_event(forward_event) + self.tensor_cpu.copy_(self.tensor, non_blocking=True) + self.d2h_event.record() + self.stat = "d2h" + + def wait_d2h_finished(self): + if self.stat != "d2h": + return + get_current_stream().wait_event(self.d2h_event) + # Stream wait only orders later GPU work. Detaching storage is a host + # operation, so the D2H copy must be host-complete first. + self.d2h_event.synchronize() + # Point this tensor at a fresh empty storage instead of resize_(0). + # resize_ would free a StorageImpl that a Tensor.set_() alias can still + # hold even when _base is None and storage_offset is 0. + placeholder = torch.empty(0, dtype=self.tensor.dtype, device=self.tensor.device) + self.tensor.set_(placeholder.untyped_storage()) + self.stat = "host" + + def launch_h2d(self, h2d_stream): + if self.stat == "h2d": + return + if self.stat != "host": + return + backward_event = create_event() + backward_event.record() + + with torch.no_grad(): + with switch_to_specified_stream(h2d_stream): + h2d_stream.wait_event(backward_event) + restored = torch.empty_strided( + self.size, + self.stride, + dtype=self.tensor_cpu.dtype, + device=self.tensor.device, + ) + restored.copy_(self.tensor_cpu, non_blocking=True) + self.tensor.set_(restored.untyped_storage(), 0, self.size, self.stride) + self.h2d_event.record() + self.stat = "h2d" + + +class OffloadManager: + def __init__(self, host_cache_limit_bytes: int = 4 * 1024**3): + # 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) + + def get_cnt(self, block_idx): + return self.getcnt.get_cnt(block_idx) + + def reset_getcnt(self): + self.getcnt.reset() + + def reset(self): + self.items.clear() + self.getcnt.reset() + + def assert_exist(self, key): + if key not in self.items: + raise RuntimeError(f"Key {key} does not exist in items") + + def exist(self, key): + return key in self.items + + def put(self, key, act): + if key in self.items: + raise RuntimeError(f"Key {key} already exists in items") + self.items[key] = act + + def wait_d2h_finished(self, prefix_key): + for key in list(self.items.keys()): + if key.startswith(prefix_key): + self.items[key].wait_d2h_finished() + + def get_layer_items_keys(self, block_idx): + block_tensor_nums = self.getcnt.get_layer_tensor_nums(block_idx) + layer_items_keys = [] + + for tensor_id in range(block_tensor_nums): + key = f"{block_idx}_{tensor_id}" + if key in self.items.keys(): + layer_items_keys.append(key) + return layer_items_keys + + def get(self, key): + self.assert_exist(key) + return self.items[key] + + def prefetch_get(self, block_idx, h2d_stream, d2h_stream): + prefetch_keys = self.getcnt.get_prefetch_keys(block_idx) + for prefetch_key in prefetch_keys: + if self.exist(prefetch_key): + prefetch_swap_tensor = self.get(prefetch_key) + d2h_stream.wait_stream(h2d_stream) + prefetch_swap_tensor.launch_h2d(h2d_stream) + + def clear(self, key=None): + if key is None: + self.items.clear() + else: + self.assert_exist(key) + self.items.pop(key) + + +def _unpack_swap_tensor(manager, swap_tensor, *, prefetch: bool) -> torch.Tensor: + """Restore an offloaded activation. + + PyTorch may invoke the unpack hook more than once for the same packed + object (``retain_graph=True``, ``create_graph=True``, or repeated saved- + tensor access). The first call moves the tensor back to device and + drops the manager key; later calls must still return the restored tensor + without treating a missing key as an error. + """ + if isinstance(swap_tensor, torch.Tensor): + return swap_tensor + + d2h_stream = manager.swap_stream + h2d_stream = manager.swap_stream + swap_tensor.launch_h2d(h2d_stream) + + get_current_stream().wait_event(swap_tensor.h2d_event) + swap_tensor.release_host_buffer() + swap_tensor.stat = "device" + + if manager.exist(swap_tensor.key): + manager.clear(swap_tensor.key) + if prefetch: + block_idx, _tensor_idx = swap_tensor.key.split("_") + manager.prefetch_get(int(block_idx), h2d_stream, d2h_stream) + return swap_tensor.tensor + + +def reset_async_activation_offload(model): + """Reset all async-offload managers attached to ``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. + """ + managers = {} + for module in model.modules(): + manager = getattr(module, "_veomni_offload_manager", None) + if manager is not None: + managers[id(manager)] = manager + + for manager in managers.values(): + manager.reset() + + +class async_save_on_cpu(saved_tensors_hooks): + def __init__( + self, + block_idx, + depth, + manager, + custom_check_fn=None, + prefetch=True, + ) -> None: + def _pack_to_cpu(tensor): + matched_hidden = ( + custom_check_fn is not None and isinstance(tensor, torch.Tensor) and custom_check_fn(tensor) + ) + if matched_hidden and _is_active_accelerator_tensor(tensor) and not _has_private_dense_storage(tensor): + # Linear/GELU often save a view of hidden_states (e.g. a 2-D + # reshape). Unpack returns this private copy for backward. + tensor = tensor.detach().contiguous().clone() + + if not base_check_fn(tensor): + return tensor + + if custom_check_fn is not None and not matched_hidden: + return tensor + + key, previous_block_idx = manager.get_cnt(block_idx) + d2h_stream = manager.swap_stream + + if previous_block_idx is not None: + manager.wait_d2h_finished("{}_".format(previous_block_idx)) + + if block_idx == depth - 1: + return tensor + + swap_tensor = SwapTensor(tensor, key, manager.host_buffer_pool) + + if block_idx < depth - 1: + swap_tensor.launch_d2h(d2h_stream) + + manager.put(key, swap_tensor) + return swap_tensor + + def _unpack_from_cpu(swap_tensor) -> torch.Tensor: + return _unpack_swap_tensor(manager, swap_tensor, prefetch=prefetch) + + super().__init__(_pack_to_cpu, _unpack_from_cpu) + + +def get_offload_modules(model, plan: List[str]): + matched_submodules = [] + offload_layers = 0 + for plan_name in plan: + if "{*}" in plan_name: + prefix = plan_name.split("{*}")[0].rstrip(".") + parent_module = model + try: + for sub_name in prefix.split("."): + parent_module = getattr(parent_module, sub_name) + if parent_module is None: + continue + depth = len(parent_module) + for layer_idx, module in enumerate(parent_module): + full_module_name = f"{prefix}.{layer_idx}" + matched_submodules.append([full_module_name, module, offload_layers, depth]) + offload_layers += 1 + except AttributeError as e: + logger.info_rank0(f"Skip plan {plan_name}: {e}") + else: + depth = 1 + for name, module in model.named_modules(): + if _module_name_match(plan_name, name): + if not any(item[0] == name for item in matched_submodules): + matched_submodules.append([name, module, offload_layers, depth]) + offload_layers += 1 + + for matched_submodule in matched_submodules: + matched_submodule[-1] = offload_layers + + return matched_submodules + + +def _get_no_split_offload_modules(model): + target_classes = set(getattr(model, "_no_split_modules", None) or []) + if not target_classes: + raise ValueError( + "Async activation offload auto-discovery requires model._no_split_modules when " + "activation_offload_modules is empty." + ) + + matched_submodules = [ + [name, module, layer_idx, 0] + for layer_idx, (name, module) in enumerate( + (item for item in model.named_modules() if type(item[1]).__name__ in target_classes) + ) + ] + if not matched_submodules: + raise ValueError( + "Async activation offload did not match any modules listed in model._no_split_modules: " + f"{sorted(target_classes)!r}" + ) + + depth = len(matched_submodules) + for matched_submodule in matched_submodules: + matched_submodule[-1] = depth + return matched_submodules + + +def async_offload_modules(modules, prefetch=True, hidden_states_idx=0, host_cache_limit_bytes: int = 4 * 1024**3): + """Apply async activation offload via selected-instance ``__call__`` patching. + + For each module, per-instance attributes (``_veomni_offload_layer_idx``, + ``_veomni_offload_depth``, ``_veomni_offload_hidden_states_idx``, + ``_veomni_offload_prefetch``) are set, then the instance is rebound to a + private subclass whose ``__call__`` wraps it with ``async_save_on_cpu``. + + A per-instance subclass is required because Python's dunder-method dispatch + bypasses instance-level assignments. It avoids mutating the original class, + so unselected modules and other models in the process remain unchanged. + + More importantly, class-level ``__call__`` patching places the + ``async_save_on_cpu`` context **outside** the ``GradientCheckpointingLayer`` + checkpoint boundary. This matches MindSpeed-MM's behavior where + ``recompute_wrapper`` is applied first and ``with_async_save_on_cpu`` + wraps it second — so the ``saved_tensors_hooks`` stack order during + forward is: + + [async_save_on_cpu (outer), _checkpoint_hook (inner)] + + With this order, ``_NoopSaveInputs.apply`` (called by checkpoint BEFORE + ``_checkpoint_hook`` is pushed) saves input tensors through + ``async_save_on_cpu.pack``, allowing hidden_states to be offloaded to CPU. + Intermediate activations are then handled by ``_checkpoint_hook`` (GC + recomputation), requiring zero persistent GPU memory. + + **Key uniqueness across multiple forward passes**: VLM models may invoke + the same visual-block sequence twice per training step (image forward + + video forward / FSDP dummy_forward). ``GetCnt`` must produce unique keys + across these passes; the original MindSpeed-MM implementation reset + ``_block_tensor_nums`` when ``block_idx`` decreased, causing key collisions + (e.g. ``"26_0"`` generated twice) that made ``OffloadManager.clear`` fail + 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"``, …. + """ + manager = OffloadManager(host_cache_limit_bytes=host_cache_limit_bytes) + 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}" + ) + module._veomni_offload_layer_idx = layer_idx + module._veomni_offload_depth = depth + module._veomni_offload_hidden_states_idx = hidden_states_idx + module._veomni_offload_prefetch = prefetch + module._veomni_offload_manager = manager + _patch_instance_call(module) + + +def _patch_instance_call(module): + """Rebind one module to a private subclass wrapping ``__call__``. + + This keeps saved-tensor hooks outside gradient checkpointing without + permanently mutating the original module class. + """ + cls = type(module) + if getattr(cls, "_veomni_async_offload_instance_class", False): + return + + original_call = cls.__call__ + + def patched_call(self, *args, **kwargs): + hidden_states_idx = getattr(self, "_veomni_offload_hidden_states_idx", 0) + prefetch = getattr(self, "_veomni_offload_prefetch", True) + layer_idx = self._veomni_offload_layer_idx + depth = self._veomni_offload_depth + manager = self._veomni_offload_manager + + if layer_idx == 0 and not manager.items: + manager.reset_getcnt() + + try: + hidden_states = args[hidden_states_idx] if len(args) > hidden_states_idx else None + if hidden_states is None: + hidden_states = kwargs.get("hidden_states") + if hidden_states is not None and not hasattr(hidden_states, "data_ptr"): + hidden_states = None + except (IndexError, AttributeError): + hidden_states = None + + if hidden_states is not None: + target_ptr = hidden_states.data_ptr() + + def check_fn(x): + return x.data_ptr() == target_ptr + else: + check_fn = None + + ctx = async_save_on_cpu( + block_idx=layer_idx, + depth=depth, + manager=manager, + custom_check_fn=check_fn, + prefetch=prefetch, + ) + + try: + with ctx: + return original_call(self, *args, **kwargs) + except BaseException: + manager.reset() + raise + + async_offload_cls = type( + cls.__name__, + (cls,), + { + "__call__": patched_call, + "_veomni_async_offload_instance_class": True, + }, + ) + module.__class__ = async_offload_cls + + +def apply_async_activation_offload( + model, + activation_offload_modules: List[str], + host_cache_limit_bytes: int = 4 * 1024**3, +): + """Apply async activation offload to matched submodules. + + Uses a selected-instance ``__call__`` subclass so that + ``saved_tensors_hooks`` wraps the ``GradientCheckpointingLayer`` checkpoint + boundary without mutating the shared decoder class. + + Compatible with both GC-enabled and GC-disabled training: + - With GC: async offload intercepts hidden_states (via _NoopSaveInputs), + GC handles intermediate activations via recomputation. + - Without GC: async offload intercepts hidden_states directly, + intermediate activations stay on GPU. + """ + if activation_offload_modules: + matched_modules = get_offload_modules(model, activation_offload_modules) + if not matched_modules: + raise ValueError( + f"activation_offload_modules did not match any model modules: {activation_offload_modules!r}" + ) + else: + matched_modules = _get_no_split_offload_modules(model) + async_offload_modules(matched_modules, host_cache_limit_bytes=host_cache_limit_bytes) diff --git a/veomni/trainer/base.py b/veomni/trainer/base.py index f89ed35fc2..35b435bdb0 100644 --- a/veomni/trainer/base.py +++ b/veomni/trainer/base.py @@ -46,7 +46,7 @@ from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin from transformers.modeling_outputs import ModelOutput -from ..arguments import VeOmniArguments, save_args +from ..arguments import OffloadConfig, VeOmniArguments, save_args from ..checkpoint import CheckpointerBase from ..data import ( DistributedDataloader, @@ -56,6 +56,7 @@ from ..data.chat_template import ChatTemplate from ..data.data_collator import DataCollator, MainCollator from ..data.data_transform import build_data_transform +from ..distributed.async_offload import apply_async_activation_offload, reset_async_activation_offload from ..distributed.clip_grad_norm import veomni_clip_grad_norm from ..distributed.offloading import build_activation_offloading_context from ..distributed.parallel_state import clear_parallel_state, init_parallel_state, use_parallel_state @@ -208,6 +209,13 @@ def state_dict(self): return {} +def _resolve_offload_config(args) -> OffloadConfig: + """Return activation-offload config, or the disabled defaults if a stub omitted it.""" + accelerator = getattr(getattr(args, "train", None), "accelerator", None) + config = getattr(accelerator, "offload_config", None) + return config if config is not None else OffloadConfig() + + def mean_aux_metrics(total_aux_metrics: Dict[str, float], num_micro_steps: int) -> Dict[str, float]: """Reduce accumulated ``aux_metrics`` to the mean over a step's micro batches. @@ -524,6 +532,20 @@ def _build_dataloader(self): def _build_parallelized_model(self): args: VeOmniArguments = self.args + # Apply async activation offload BEFORE FSDP2 sharding. + # Uses per-instance __call__ patching so that async_save_on_cpu is + # OUTER to the checkpoint boundary pushed by GradientCheckpointingLayer, + # matching MindSpeed-MM's GC+async offload behavior: hidden_states + # inputs are offloaded to CPU (via _NoopSaveInputs), while intermediate + # activations are handled by GC recomputation (via _checkpoint_hook). + offload_config = _resolve_offload_config(args) + if offload_config.enable_async_activation: + apply_async_activation_offload( + self.model, + offload_config.activation_offload_modules, + host_cache_limit_bytes=int(offload_config.activation_offload_host_cache_limit_gb * 1024**3), + ) + kwargs = {} cpu_load_param_name = None if hasattr(self.model, "get_parallel_plan"): @@ -610,10 +632,19 @@ def _build_lr_scheduler(self): def _build_training_context(self): """Build training context for distributed training.""" + offload_config = _resolve_offload_config(self.args) + + # Async activation offload uses per-module saved_tensors_hooks (applied + # before FSDP sharding), so the global fwd/bwd contexts are nullcontext. + if offload_config.enable_async_activation: + from contextlib import nullcontext + + self.model_fwd_context, self.model_bwd_context = nullcontext(), nullcontext() + return self.model_fwd_context, self.model_bwd_context = build_activation_offloading_context( - self.args.train.accelerator.offload_config.enable_activation, + offload_config.enable_activation, self.args.train.gradient_checkpointing.enable, - self.args.train.accelerator.offload_config.activation_gpu_limit, + offload_config.activation_gpu_limit, ) def _init_callbacks(self): @@ -808,6 +839,10 @@ def _configure_hsdp_allreduce(self, micro_step: int, num_micro_steps: int): elif micro_step == num_micro_steps - 1: self.model.set_requires_all_reduce(True) + def _reset_async_activation_offload_if_enabled(self): + if _resolve_offload_config(self.args).enable_async_activation: + reset_async_activation_offload(self.model) + def sync_before_train_step(self): if self.args.train.sync_each_train_step: synchronize() @@ -821,6 +856,7 @@ def train_step( micro_batches: List[Dict[str, Any]] = next(data_iterator) + self._reset_async_activation_offload_if_enabled() self.on_step_begin(micro_batches=micro_batches) # Forward and backward for each micro batch diff --git a/veomni/trainer/dit_trainer.py b/veomni/trainer/dit_trainer.py index 5d3ca7eb0a..7f13340322 100644 --- a/veomni/trainer/dit_trainer.py +++ b/veomni/trainer/dit_trainer.py @@ -525,6 +525,7 @@ def train_step(self, data_iterator: Any) -> Dict[str, float]: else: micro_batches = next(data_iterator) + self.base._reset_async_activation_offload_if_enabled() self.on_step_begin(micro_batches=micro_batches) self.base.sync_before_train_step() diff --git a/veomni/trainer/text_dpo_trainer.py b/veomni/trainer/text_dpo_trainer.py index c17a1d0b68..57c91b7a7b 100644 --- a/veomni/trainer/text_dpo_trainer.py +++ b/veomni/trainer/text_dpo_trainer.py @@ -411,6 +411,7 @@ def train_step(self, data_iterator: Any) -> Dict[str, float]: micro_batches: List[Dict[str, Any]] = next(data_iterator) + self.base._reset_async_activation_offload_if_enabled() self.on_step_begin(micro_batches=micro_batches) self.base.sync_before_train_step() diff --git a/veomni/trainer/text_trainer.py b/veomni/trainer/text_trainer.py index ac637b36a8..8a29485ae1 100644 --- a/veomni/trainer/text_trainer.py +++ b/veomni/trainer/text_trainer.py @@ -117,6 +117,7 @@ def train_step( micro_batches: List[Dict[str, Any]] = next(data_iterator) + self.base._reset_async_activation_offload_if_enabled() self.on_step_begin(micro_batches=micro_batches) # Forward and backward for each micro batch diff --git a/veomni/trainer/vlm_trainer.py b/veomni/trainer/vlm_trainer.py index 7fb9928a89..43bb4ef925 100644 --- a/veomni/trainer/vlm_trainer.py +++ b/veomni/trainer/vlm_trainer.py @@ -342,6 +342,7 @@ def train_step( micro_batches: List[Dict[str, Any]] = next(data_iterator) + self.base._reset_async_activation_offload_if_enabled() self.on_step_begin(micro_batches=micro_batches) # Forward and backward for each micro batch diff --git a/veomni/utils/device.py b/veomni/utils/device.py index f2d57646a3..a6d05a55b5 100644 --- a/veomni/utils/device.py +++ b/veomni/utils/device.py @@ -123,6 +123,45 @@ def is_sm90_or_above() -> bool: return get_gpu_compute_capability() >= 90 +def create_stream(device: torch.device | None = None, priority: int = 0) -> Any: + """Create a device stream (CUDA/NPU-agnostic).""" + device_type = get_device_type() + device_api = get_torch_device() + + # ``torch.cpu.Stream`` intentionally has a smaller constructor than the + # CUDA/NPU stream implementations. Passing ``device`` or ``priority`` + # through unconditionally makes CPU-only unit tests fail before they can + # exercise the device-agnostic caller. + if device_type == "cpu": + return device_api.Stream() + + kwargs = {} + if device is not None: + kwargs["device"] = device + if priority: + kwargs["priority"] = priority + return device_api.Stream(**kwargs) + + +def create_event(enable_timing: bool = False, blocking: bool = False) -> Any: + """Create a device event (CUDA/NPU-agnostic).""" + device_api = get_torch_device() + if get_device_type() == "cpu": + # ``torch.cpu.Event`` is a no-argument synchronization marker. + return device_api.Event() + return device_api.Event(enable_timing=enable_timing, blocking=blocking) + + +def get_current_stream() -> Any: + """Get the current device stream.""" + return get_torch_device().current_stream() + + +def switch_to_specified_stream(stream: Any) -> Any: + """Context manager to switch to a specified device stream.""" + return get_torch_device().stream(stream) + + def get_compute_units(): """ Returns the number of streaming multiprocessors (SMs) or equivalent compute units