[Bug]: Layer activation swap frees shared module input storage

Checklist

• 1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues).
• 2. I have read the relevant documentation and activation swap implementation.
• 3. I have a reproducible 8-card case, the full runtime error, and a validated fix.

🐛 Describe the bug

Layer-level activation swap may destructively release storage that is owned outside the wrapped module. This happens when a tensor passed as a module input is also saved for backward inside the wrapper and reused by other layers.

The reproduced case is MindFormers DeepSeek-V3 with 12 decoder layers, 8-card FSDP, and layer swap enabled for layers 0-10. RoPE cos/sin are computed once at model level and the same float32 tensors are passed to every decoder layer through rotary_cos_sin.

Minimal model pattern:

rotary_cos_sin = compute_rotary_cos_sin(...)  # layer-invariant shared tensors
for layer in decoder.layers:
    hidden_states = layer(hidden_states, rotary_cos_sin=rotary_cos_sin)

Swap configuration:

recompute:
  mode: "None"

swap:
  enable: true
  default_prefetch: 1
  layer_swap:
    - layers: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

parallelism:
  data_parallel_shard: -1
  data_parallel_shard_strategy: optim_grads_params
  tensor_parallel: 1
  context_parallel: 1
  expert_parallel: 1
  pipeline_parallel: 1

Launch command:

export PYTHONPATH=<hyper-parallel-root>:<mindformers-root>:$PYTHONPATH
bash scripts/msrun_launcher.sh \
  "run_mindformer.py --config <deepseek-v3-config>.yaml --mode 1" \
  8 8129 <unique-log-dir> True 7200

The first backward fails in an elementwise Mul:

RuntimeError: SyncStream failed for op aclnnMul
The DDR address of the MTE instruction is out of range.
fault kernel_name=Mul_ee98c6628030785f610b924ab1557b31_high_performance
tensor index:0
addr:0
size:1048576 bytes

1048576 = 4096 * 64 * sizeof(float32), matching one layer-invariant RoPE cos/sin tensor in this configuration. The operator dump confirms float32 inputs.

Control results:

• FSDP only: completes 10/10 steps.
• Operator-level swap on self_attention.core_attention: completes 10/10 steps.
• Full-layer swap on all or alternating layers: reproduces the null-address Mul failure.

Root cause

1. The saved-tensor hook registers tensors saved for backward inside SwapWrapper, including module inputs such as shared RoPE cos/sin.
2. SwapTensor.wait_offload() frees the original device storage with storage.resize_(0).
3. set_forward_prefetch_layer() currently calls protect_alias_tensors(group_name, output), so only aliases of the module output are protected.
4. Module inputs and keyword arguments are not protected. A shared input can therefore be offloaded by one layer's swap group and have its storage resized while it is still owned and reused outside that layer.
5. Per-group duplicate registration does not solve external ownership or cross-group lifetime.

Relevant files:

• hyper_parallel/core/activation_checkpoint/swap.py
• hyper_parallel/platform/mindspore/activation_checkpoint/activation_swap.py

Expected behavior

Activation swap may change where saved activations reside, but it must not invalidate storage still owned by module callers or other layers. Enabling layer-level swap should preserve forward/backward correctness without model-specific shape checks or MUST_SAVE rules for RoPE.

Additional context

Proposed solution

Treat tensors crossing the wrapped module boundary as externally owned:

1. Register the swap forward pre-hook with with_kwargs=True.
2. Collect only the device storage pointers from positional and keyword inputs before forward; do not retain Tensor references.
3. At the forward post-hook, union input storage pointers with output storage pointers.
4. Before launch_offload, mark matching SwapTensor entries as keep_on_device through a storage-pointer-level API.
5. Clear the captured input pointer set after the post-hook.

Conceptually:

def forward_pre_hook(module, args, kwargs):
    module._swap_input_storage_ptrs = collect_storage_ptrs((args, kwargs))

def forward_hook(module, args, output):
    alias_ptrs = module._swap_input_storage_ptrs
    alias_ptrs.update(collect_storage_ptrs(output))
    swap_manager.protect_alias_storage_ptrs(group_name, alias_ptrs)
    swap_manager.launch_offload(group_name)
    module._swap_input_storage_ptrs = set()

This is conservative: module inputs and outputs remain on device, while internal Attention/MLP saved activations are still eligible for swap. A future explicit storage-ownership transfer or cross-group storage-lifetime mechanism could safely offload selected boundary tensors, but implicit resize_(0) of externally owned storage is not safe.

This issue is distinct from #185: #185 focuses on copying/restoring multiple tensor views of the same storage efficiently. This bug is about destructive offload of storage owned outside the swap wrapper.

Validation of the proposed change

With the input/output boundary protection applied to HyperParallel master:

8 x Ascend 910B2
FSDP only + layer_swap on decoder layers 0-10
step 1/10:  loss 11.865072, grad_norm 2.578646
step 10/10: loss 11.463392, grad_norm 3.402987
exit code: 0

The loss and gradient norm sequence matches the FSDP-only and operator-level-swap controls. No Mul/MTE/null-address error appears in any worker log.

Recommended regression tests:

• forward hooks protect aliases from both positional inputs and keyword inputs;
• input and output storage-pointer sets are both passed to alias protection;
• matching SwapTensor entries remain in STATE_DEVICE/keep_on_device;
• layer-level swap with a tensor shared by multiple layers completes backward;
• both MindSpore and Torch hook registration preserve the same semantics.

Environment info

HyperParallel: master a7f94f85c962f21964398035d2c652b52ec82c06
MindFormers reproduction commit: a8312a61906c89f717ede457a0328d03dae0263c
MindSpore: 2.10.0
Hardware: 8 x Ascend 910B2 (64 GiB)
npu-smi: 25.5.1
OS: openEuler 22.03 SP4, aarch64, kernel 5.10.0

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 294
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/294

Start with hyper_parallel/core/activation_checkpoint/swap.py and hyper_parallel/platform/mindspore/activation_checkpoint/activation_swap.py, then inspect the forward hook registration and alias protection paths. Run the recommended hook, storage-state, and shared-input regression tests, using the described layer-level swap reproduction if available. Done means positional and keyword boundary tensors remain valid through backward while internal activations can still be swapped, with MindSpore and Torch hooks retaining equivalent semantics.