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